codekingpro/portable-devtools

1"""PII detection and handling middleware for agents."""2 3from __future__ import annotations4 5from typing import TYPE_CHECKING, Any, Literal6 7from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage8from typing_extensions import override9 10from langchain.agents.middleware._redaction import (11    PIIDetectionError,12    PIIMatch,13    RedactionRule,14    ResolvedRedactionRule,15    apply_strategy,16    detect_credit_card,17    detect_email,18    detect_ip,19    detect_mac_address,20    detect_url,21)22from langchain.agents.middleware.types import (23    AgentMiddleware,24    AgentState,25    ContextT,26    ResponseT,27    hook_config,28)29 30if TYPE_CHECKING:31    from collections.abc import Callable32 33    from langgraph.runtime import Runtime34 35 36class PIIMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]):37    """Detect and handle Personally Identifiable Information (PII) in conversations.38 39    This middleware detects common PII types and applies configurable strategies40    to handle them. It can detect emails, credit cards, IP addresses, MAC addresses, and41    URLs in both user input and agent output.42 43    Built-in PII types:44 45    - `email`: Email addresses46    - `credit_card`: Credit card numbers (validated with Luhn algorithm)47    - `ip`: IP addresses (validated with stdlib)48    - `mac_address`: MAC addresses49    - `url`: URLs (both `http`/`https` and bare URLs)50 51    Strategies:52 53    - `block`: Raise an exception when PII is detected54    - `redact`: Replace PII with `[REDACTED_TYPE]` placeholders55    - `mask`: Partially mask PII (e.g., `****-****-****-1234` for credit card)56    - `hash`: Replace PII with deterministic hash (e.g., `<email_hash:a1b2c3d4>`)57 58    Strategy Selection Guide:59 60    | Strategy | Preserves Identity? | Best For                                |61    | -------- | ------------------- | --------------------------------------- |62    | `block`  | N/A                 | Avoid PII completely                    |63    | `redact` | No                  | General compliance, log sanitization    |64    | `mask`   | No                  | Human readability, customer service UIs |65    | `hash`   | Yes (pseudonymous)  | Analytics, debugging                    |66 67    Example:68        ```python69        from langchain.agents.middleware import PIIMiddleware70        from langchain.agents import create_agent71 72        # Redact all emails in user input73        agent = create_agent(74            "openai:gpt-5",75            middleware=[76                PIIMiddleware("email", strategy="redact"),77            ],78        )79 80        # Use different strategies for different PII types81        agent = create_agent(82            "openai:gpt-4o",83            middleware=[84                PIIMiddleware("credit_card", strategy="mask"),85                PIIMiddleware("url", strategy="redact"),86                PIIMiddleware("ip", strategy="hash"),87            ],88        )89 90        # Custom PII type with regex91        agent = create_agent(92            "openai:gpt-5",93            middleware=[94                PIIMiddleware("api_key", detector=r"sk-[a-zA-Z0-9]{32}", strategy="block"),95            ],96        )97        ```98    """99 100    def __init__(101        self,102        # From a typing point of view, the literals are covered by 'str'.103        # Nonetheless, we escape PYI051 to keep hints and autocompletion for the caller.104        pii_type: Literal["email", "credit_card", "ip", "mac_address", "url"] | str,  # noqa: PYI051105        *,106        strategy: Literal["block", "redact", "mask", "hash"] = "redact",107        detector: Callable[[str], list[PIIMatch]] | str | None = None,108        apply_to_input: bool = True,109        apply_to_output: bool = False,110        apply_to_tool_results: bool = False,111    ) -> None:112        """Initialize the PII detection middleware.113 114        Args:115            pii_type: Type of PII to detect.116 117                Can be a built-in type (`email`, `credit_card`, `ip`, `mac_address`,118                `url`) or a custom type name.119            strategy: How to handle detected PII.120 121                Options:122 123                * `block`: Raise `PIIDetectionError` when PII is detected124                * `redact`: Replace with `[REDACTED_TYPE]` placeholders125                * `mask`: Partially mask PII (show last few characters)126                * `hash`: Replace with deterministic hash (format: `<type_hash:digest>`)127 128            detector: Custom detector function or regex pattern.129 130                * If `Callable`: Function that takes content string and returns131                    list of `PIIMatch` objects132                * If `str`: Regex pattern to match PII133                * If `None`: Uses built-in detector for the `pii_type`134            apply_to_input: Whether to check user messages before model call.135            apply_to_output: Whether to check AI messages after model call.136            apply_to_tool_results: Whether to check tool result messages after tool execution.137 138        Raises:139            ValueError: If `pii_type` is not built-in and no detector is provided.140        """141        super().__init__()142 143        self.apply_to_input = apply_to_input144        self.apply_to_output = apply_to_output145        self.apply_to_tool_results = apply_to_tool_results146 147        self._resolved_rule: ResolvedRedactionRule = RedactionRule(148            pii_type=pii_type,149            strategy=strategy,150            detector=detector,151        ).resolve()152        self.pii_type = self._resolved_rule.pii_type153        self.strategy = self._resolved_rule.strategy154        self.detector = self._resolved_rule.detector155 156    @property157    def name(self) -> str:158        """Name of the middleware."""159        return f"{self.__class__.__name__}[{self.pii_type}]"160 161    def _process_content(self, content: str) -> tuple[str, list[PIIMatch]]:162        """Apply the configured redaction rule to the provided content."""163        matches = self.detector(content)164        if not matches:165            return content, []166        sanitized = apply_strategy(content, matches, self.strategy)167        return sanitized, matches168 169    @hook_config(can_jump_to=["end"])170    @override171    def before_model(172        self,173        state: AgentState[Any],174        runtime: Runtime[ContextT],175    ) -> dict[str, Any] | None:176        """Check user messages and tool results for PII before model invocation.177 178        Args:179            state: The current agent state.180            runtime: The langgraph runtime.181 182        Returns:183            Updated state with PII handled according to strategy, or `None` if no PII184                detected.185 186        Raises:187            PIIDetectionError: If PII is detected and strategy is `'block'`.188        """189        if not self.apply_to_input and not self.apply_to_tool_results:190            return None191 192        messages = state["messages"]193        if not messages:194            return None195 196        new_messages = list(messages)197        any_modified = False198 199        # Check user input if enabled200        if self.apply_to_input:201            # Get last user message202            last_user_msg = None203            last_user_idx = None204            for i in range(len(messages) - 1, -1, -1):205                if isinstance(messages[i], HumanMessage):206                    last_user_msg = messages[i]207                    last_user_idx = i208                    break209 210            if last_user_idx is not None and last_user_msg and last_user_msg.content:211                # Detect PII in message content212                content = str(last_user_msg.content)213                new_content, matches = self._process_content(content)214 215                if matches:216                    updated_message: AnyMessage = HumanMessage(217                        content=new_content,218                        id=last_user_msg.id,219                        name=last_user_msg.name,220                    )221 222                    new_messages[last_user_idx] = updated_message223                    any_modified = True224 225        # Check tool results if enabled226        if self.apply_to_tool_results:227            # Find the last AIMessage, then process all `ToolMessage` objects after it228            last_ai_idx = None229            for i in range(len(messages) - 1, -1, -1):230                if isinstance(messages[i], AIMessage):231                    last_ai_idx = i232                    break233 234            if last_ai_idx is not None:235                # Get all tool messages after the last AI message236                for i in range(last_ai_idx + 1, len(messages)):237                    msg = messages[i]238                    if isinstance(msg, ToolMessage):239                        tool_msg = msg240                        if not tool_msg.content:241                            continue242 243                        content = str(tool_msg.content)244                        new_content, matches = self._process_content(content)245 246                        if not matches:247                            continue248 249                        # Create updated tool message250                        updated_message = ToolMessage(251                            content=new_content,252                            id=tool_msg.id,253                            name=tool_msg.name,254                            tool_call_id=tool_msg.tool_call_id,255                        )256 257                        new_messages[i] = updated_message258                        any_modified = True259 260        if any_modified:261            return {"messages": new_messages}262 263        return None264 265    @hook_config(can_jump_to=["end"])266    async def abefore_model(267        self,268        state: AgentState[Any],269        runtime: Runtime[ContextT],270    ) -> dict[str, Any] | None:271        """Async check user messages and tool results for PII before model invocation.272 273        Args:274            state: The current agent state.275            runtime: The langgraph runtime.276 277        Returns:278            Updated state with PII handled according to strategy, or `None` if no PII279                detected.280 281        Raises:282            PIIDetectionError: If PII is detected and strategy is `'block'`.283        """284        return self.before_model(state, runtime)285 286    @override287    def after_model(288        self,289        state: AgentState[Any],290        runtime: Runtime[ContextT],291    ) -> dict[str, Any] | None:292        """Check AI messages for PII after model invocation.293 294        Args:295            state: The current agent state.296            runtime: The langgraph runtime.297 298        Returns:299            Updated state with PII handled according to strategy, or None if no PII300                detected.301 302        Raises:303            PIIDetectionError: If PII is detected and strategy is `'block'`.304        """305        if not self.apply_to_output:306            return None307 308        messages = state["messages"]309        if not messages:310            return None311 312        # Get last AI message313        last_ai_msg = None314        last_ai_idx = None315        for i in range(len(messages) - 1, -1, -1):316            msg = messages[i]317            if isinstance(msg, AIMessage):318                last_ai_msg = msg319                last_ai_idx = i320                break321 322        if last_ai_idx is None or not last_ai_msg or not last_ai_msg.content:323            return None324 325        # Detect PII in message content326        content = str(last_ai_msg.content)327        new_content, matches = self._process_content(content)328 329        if not matches:330            return None331 332        # Create updated message333        updated_message = AIMessage(334            content=new_content,335            id=last_ai_msg.id,336            name=last_ai_msg.name,337            tool_calls=last_ai_msg.tool_calls,338        )339 340        # Return updated messages341        new_messages = list(messages)342        new_messages[last_ai_idx] = updated_message343 344        return {"messages": new_messages}345 346    async def aafter_model(347        self,348        state: AgentState[Any],349        runtime: Runtime[ContextT],350    ) -> dict[str, Any] | None:351        """Async check AI messages for PII after model invocation.352 353        Args:354            state: The current agent state.355            runtime: The langgraph runtime.356 357        Returns:358            Updated state with PII handled according to strategy, or None if no PII359                detected.360 361        Raises:362            PIIDetectionError: If PII is detected and strategy is `'block'`.363        """364        return self.after_model(state, runtime)365 366 367__all__ = [368    "PIIDetectionError",369    "PIIMatch",370    "PIIMiddleware",371    "detect_credit_card",372    "detect_email",373    "detect_ip",374    "detect_mac_address",375    "detect_url",376]377