langchain 1.0.0a9__py3-none-any.whl → 1.0.0a11__py3-none-any.whl

langchain/agents/middleware/pii.py

@@ -0,0 +1,753 @@
+"""PII detection and handling middleware for agents."""
+
+from __future__ import annotations
+
+import hashlib
+import ipaddress
+import re
+from typing import TYPE_CHECKING, Any, Literal
+from urllib.parse import urlparse
+
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage
+from typing_extensions import TypedDict
+
+from langchain.agents.middleware.types import AgentMiddleware, AgentState, hook_config
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from langgraph.runtime import Runtime
+
+
+class PIIMatch(TypedDict):
+    """Represents a detected PII match in text."""
+
+    type: str
+    """The type of PII detected (e.g., 'email', 'ssn', 'credit_card')."""
+    value: str
+    """The actual matched text."""
+    start: int
+    """Starting position of the match in the text."""
+    end: int
+    """Ending position of the match in the text."""
+
+
+class PIIDetectionError(Exception):
+    """Exception raised when PII is detected and strategy is 'block'."""
+
+    def __init__(self, pii_type: str, matches: list[PIIMatch]) -> None:
+        """Initialize the exception with PII detection information.
+
+        Args:
+            pii_type: The type of PII that was detected.
+            matches: List of PII matches found.
+        """
+        self.pii_type = pii_type
+        self.matches = matches
+        count = len(matches)
+        msg = f"Detected {count} instance(s) of {pii_type} in message content"
+        super().__init__(msg)
+
+
+# ============================================================================
+# PII Detection Functions
+# ============================================================================
+
+
+def _luhn_checksum(card_number: str) -> bool:
+    """Validate credit card number using Luhn algorithm.
+
+    Args:
+        card_number: Credit card number string (digits only).
+
+    Returns:
+        True if the number passes Luhn validation, False otherwise.
+    """
+    digits = [int(d) for d in card_number if d.isdigit()]
+
+    if len(digits) < 13 or len(digits) > 19:
+        return False
+
+    checksum = 0
+    for i, digit in enumerate(reversed(digits)):
+        d = digit
+        if i % 2 == 1:
+            d *= 2
+            if d > 9:
+                d -= 9
+        checksum += d
+
+    return checksum % 10 == 0
+
+
+def detect_email(content: str) -> list[PIIMatch]:
+    """Detect email addresses in content.
+
+    Args:
+        content: Text content to scan.
+
+    Returns:
+        List of detected email matches.
+    """
+    pattern = r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"
+    return [
+        PIIMatch(
+            type="email",
+            value=match.group(),
+            start=match.start(),
+            end=match.end(),
+        )
+        for match in re.finditer(pattern, content)
+    ]
+
+
+def detect_credit_card(content: str) -> list[PIIMatch]:
+    """Detect credit card numbers in content using Luhn validation.
+
+    Detects cards in formats like:
+    - 1234567890123456
+    - 1234 5678 9012 3456
+    - 1234-5678-9012-3456
+
+    Args:
+        content: Text content to scan.
+
+    Returns:
+        List of detected credit card matches.
+    """
+    # Match various credit card formats
+    pattern = r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b"
+    matches = []
+
+    for match in re.finditer(pattern, content):
+        card_number = match.group()
+        # Validate with Luhn algorithm
+        if _luhn_checksum(card_number):
+            matches.append(
+                PIIMatch(
+                    type="credit_card",
+                    value=card_number,
+                    start=match.start(),
+                    end=match.end(),
+                )
+            )
+
+    return matches
+
+
+def detect_ip(content: str) -> list[PIIMatch]:
+    """Detect IP addresses in content using stdlib validation.
+
+    Validates both IPv4 and IPv6 addresses.
+
+    Args:
+        content: Text content to scan.
+
+    Returns:
+        List of detected IP address matches.
+    """
+    matches = []
+
+    # IPv4 pattern
+    ipv4_pattern = r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b"
+
+    for match in re.finditer(ipv4_pattern, content):
+        ip_str = match.group()
+        try:
+            # Validate with stdlib
+            ipaddress.ip_address(ip_str)
+            matches.append(
+                PIIMatch(
+                    type="ip",
+                    value=ip_str,
+                    start=match.start(),
+                    end=match.end(),
+                )
+            )
+        except ValueError:
+            # Not a valid IP address
+            pass
+
+    return matches
+
+
+def detect_mac_address(content: str) -> list[PIIMatch]:
+    """Detect MAC addresses in content.
+
+    Detects formats like:
+    - 00:1A:2B:3C:4D:5E
+    - 00-1A-2B-3C-4D-5E
+
+    Args:
+        content: Text content to scan.
+
+    Returns:
+        List of detected MAC address matches.
+    """
+    pattern = r"\b([0-9A-Fa-f]{2}[:-]){5}[0-9A-Fa-f]{2}\b"
+    return [
+        PIIMatch(
+            type="mac_address",
+            value=match.group(),
+            start=match.start(),
+            end=match.end(),
+        )
+        for match in re.finditer(pattern, content)
+    ]
+
+
+def detect_url(content: str) -> list[PIIMatch]:
+    """Detect URLs in content using regex and stdlib validation.
+
+    Detects:
+    - http://example.com
+    - https://example.com/path
+    - www.example.com
+    - example.com/path
+
+    Args:
+        content: Text content to scan.
+
+    Returns:
+        List of detected URL matches.
+    """
+    matches = []
+
+    # Pattern 1: URLs with scheme (http:// or https://)
+    scheme_pattern = r"https?://[^\s<>\"{}|\\^`\[\]]+"
+
+    for match in re.finditer(scheme_pattern, content):
+        url = match.group()
+        try:
+            result = urlparse(url)
+            if result.scheme in ("http", "https") and result.netloc:
+                matches.append(
+                    PIIMatch(
+                        type="url",
+                        value=url,
+                        start=match.start(),
+                        end=match.end(),
+                    )
+                )
+        except Exception:  # noqa: S110, BLE001
+            # Invalid URL, skip
+            pass
+
+    # Pattern 2: URLs without scheme (www.example.com or example.com/path)
+    # More conservative to avoid false positives
+    bare_pattern = r"\b(?:www\.)?[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)+(?:/[^\s]*)?"  # noqa: E501
+
+    for match in re.finditer(bare_pattern, content):
+        # Skip if already matched with scheme
+        if any(
+            m["start"] <= match.start() < m["end"] or m["start"] < match.end() <= m["end"]
+            for m in matches
+        ):
+            continue
+
+        url = match.group()
+        # Only accept if it has a path or starts with www
+        # This reduces false positives like "example.com" in prose
+        if "/" in url or url.startswith("www."):
+            try:
+                # Add scheme for validation (required for urlparse to work correctly)
+                test_url = f"http://{url}"
+                result = urlparse(test_url)
+                if result.netloc and "." in result.netloc:
+                    matches.append(
+                        PIIMatch(
+                            type="url",
+                            value=url,
+                            start=match.start(),
+                            end=match.end(),
+                        )
+                    )
+            except Exception:  # noqa: S110, BLE001
+                # Invalid URL, skip
+                pass
+
+    return matches
+
+
+# Built-in detector registry
+_BUILTIN_DETECTORS: dict[str, Callable[[str], list[PIIMatch]]] = {
+    "email": detect_email,
+    "credit_card": detect_credit_card,
+    "ip": detect_ip,
+    "mac_address": detect_mac_address,
+    "url": detect_url,
+}
+
+
+# ============================================================================
+# Strategy Implementations
+# ============================================================================
+
+
+def _apply_redact_strategy(content: str, matches: list[PIIMatch]) -> str:
+    """Replace PII with [REDACTED_TYPE] placeholders.
+
+    Args:
+        content: Original content.
+        matches: List of PII matches to redact.
+
+    Returns:
+        Content with PII redacted.
+    """
+    if not matches:
+        return content
+
+    # Sort matches by start position in reverse to avoid offset issues
+    sorted_matches = sorted(matches, key=lambda m: m["start"], reverse=True)
+
+    result = content
+    for match in sorted_matches:
+        replacement = f"[REDACTED_{match['type'].upper()}]"
+        result = result[: match["start"]] + replacement + result[match["end"] :]
+
+    return result
+
+
+def _apply_mask_strategy(content: str, matches: list[PIIMatch]) -> str:
+    """Partially mask PII, showing only last few characters.
+
+    Args:
+        content: Original content.
+        matches: List of PII matches to mask.
+
+    Returns:
+        Content with PII masked.
+    """
+    if not matches:
+        return content
+
+    # Sort matches by start position in reverse
+    sorted_matches = sorted(matches, key=lambda m: m["start"], reverse=True)
+
+    result = content
+    for match in sorted_matches:
+        value = match["value"]
+        pii_type = match["type"]
+
+        # Different masking strategies by type
+        if pii_type == "email":
+            # Show only domain: user@****.com
+            parts = value.split("@")
+            if len(parts) == 2:
+                domain_parts = parts[1].split(".")
+                if len(domain_parts) >= 2:
+                    masked = f"{parts[0]}@****.{domain_parts[-1]}"
+                else:
+                    masked = f"{parts[0]}@****"
+            else:
+                masked = "****"
+
+        elif pii_type == "credit_card":
+            # Show last 4: ****-****-****-1234
+            digits_only = "".join(c for c in value if c.isdigit())
+            separator = "-" if "-" in value else " " if " " in value else ""
+            if separator:
+                masked = f"****{separator}****{separator}****{separator}{digits_only[-4:]}"
+            else:
+                masked = f"************{digits_only[-4:]}"
+
+        elif pii_type == "ip":
+            # Show last octet: *.*.*. 123
+            parts = value.split(".")
+            masked = f"*.*.*.{parts[-1]}" if len(parts) == 4 else "****"
+
+        elif pii_type == "mac_address":
+            # Show last byte: **:**:**:**:**:5E
+            separator = ":" if ":" in value else "-"
+            masked = (
+                f"**{separator}**{separator}**{separator}**{separator}**{separator}{value[-2:]}"
+            )
+
+        elif pii_type == "url":
+            # Mask everything: [MASKED_URL]
+            masked = "[MASKED_URL]"
+
+        else:
+            # Default: show last 4 chars
+            masked = f"****{value[-4:]}" if len(value) > 4 else "****"
+
+        result = result[: match["start"]] + masked + result[match["end"] :]
+
+    return result
+
+
+def _apply_hash_strategy(content: str, matches: list[PIIMatch]) -> str:
+    """Replace PII with deterministic hash including type information.
+
+    Args:
+        content: Original content.
+        matches: List of PII matches to hash.
+
+    Returns:
+        Content with PII replaced by hashes in format <type_hash:digest>.
+    """
+    if not matches:
+        return content
+
+    # Sort matches by start position in reverse
+    sorted_matches = sorted(matches, key=lambda m: m["start"], reverse=True)
+
+    result = content
+    for match in sorted_matches:
+        value = match["value"]
+        pii_type = match["type"]
+        # Create deterministic hash
+        hash_digest = hashlib.sha256(value.encode()).hexdigest()[:8]
+        replacement = f"<{pii_type}_hash:{hash_digest}>"
+        result = result[: match["start"]] + replacement + result[match["end"] :]
+
+    return result
+
+
+# ============================================================================
+# PIIMiddleware
+# ============================================================================
+
+
+class PIIMiddleware(AgentMiddleware):
+    """Detect and handle Personally Identifiable Information (PII) in agent conversations.
+
+    This middleware detects common PII types and applies configurable strategies
+    to handle them. It can detect emails, credit cards, IP addresses,
+    MAC addresses, and URLs in both user input and agent output.
+
+    Built-in PII types:
+        - ``email``: Email addresses
+        - ``credit_card``: Credit card numbers (validated with Luhn algorithm)
+        - ``ip``: IP addresses (validated with stdlib)
+        - ``mac_address``: MAC addresses
+        - ``url``: URLs (both http/https and bare URLs)
+
+    Strategies:
+        - ``block``: Raise an exception when PII is detected
+        - ``redact``: Replace PII with ``[REDACTED_TYPE]`` placeholders
+        - ``mask``: Partially mask PII (e.g., ``****-****-****-1234`` for credit card)
+        - ``hash``: Replace PII with deterministic hash (e.g., ``<email_hash:a1b2c3d4>``)
+
+    Strategy Selection Guide:
+
+        ========  ===================  =======================================
+        Strategy  Preserves Identity?  Best For
+        ========  ===================  =======================================
+        `block`   N/A                  Avoid PII completely
+        `redact`  No                   General compliance, log sanitization
+        `mask`    No                   Human readability, customer service UIs
+        `hash`    Yes (pseudonymous)   Analytics, debugging
+        ========  ===================  =======================================
+
+    Example:
+        ```python
+        from langchain.agents.middleware import PIIMiddleware
+        from langchain.agents import create_agent
+
+        # Redact all emails in user input
+        agent = create_agent(
+            "openai:gpt-5",
+            middleware=[
+                PIIMiddleware("email", strategy="redact"),
+            ],
+        )
+
+        # Use different strategies for different PII types
+        agent = create_agent(
+            "openai:gpt-4o",
+            middleware=[
+                PIIMiddleware("credit_card", strategy="mask"),
+                PIIMiddleware("url", strategy="redact"),
+                PIIMiddleware("ip", strategy="hash"),
+            ],
+        )
+
+        # Custom PII type with regex
+        agent = create_agent(
+            "openai:gpt-5",
+            middleware=[
+                PIIMiddleware("api_key", detector=r"sk-[a-zA-Z0-9]{32}", strategy="block"),
+            ],
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        pii_type: Literal["email", "credit_card", "ip", "mac_address", "url"] | str,  # noqa: PYI051
+        *,
+        strategy: Literal["block", "redact", "mask", "hash"] = "redact",
+        detector: Callable[[str], list[PIIMatch]] | str | None = None,
+        apply_to_input: bool = True,
+        apply_to_output: bool = False,
+        apply_to_tool_results: bool = False,
+    ) -> None:
+        """Initialize the PII detection middleware.
+
+        Args:
+            pii_type: Type of PII to detect. Can be a built-in type
+                (``email``, ``credit_card``, ``ip``, ``mac_address``, ``url``)
+                or a custom type name.
+            strategy: How to handle detected PII:
+
+                * ``block``: Raise PIIDetectionError when PII is detected
+                * ``redact``: Replace with ``[REDACTED_TYPE]`` placeholders
+                * ``mask``: Partially mask PII (show last few characters)
+                * ``hash``: Replace with deterministic hash (format: ``<type_hash:digest>``)
+
+            detector: Custom detector function or regex pattern.
+
+                * If ``Callable``: Function that takes content string and returns
+                  list of PIIMatch objects
+                * If ``str``: Regex pattern to match PII
+                * If ``None``: Uses built-in detector for the pii_type
+
+            apply_to_input: Whether to check user messages before model call.
+            apply_to_output: Whether to check AI messages after model call.
+            apply_to_tool_results: Whether to check tool result messages after tool execution.
+
+        Raises:
+            ValueError: If pii_type is not built-in and no detector is provided.
+        """
+        super().__init__()
+
+        self.pii_type = pii_type
+        self.strategy = strategy
+        self.apply_to_input = apply_to_input
+        self.apply_to_output = apply_to_output
+        self.apply_to_tool_results = apply_to_tool_results
+
+        # Resolve detector
+        if detector is None:
+            # Use built-in detector
+            if pii_type not in _BUILTIN_DETECTORS:
+                msg = (
+                    f"Unknown PII type: {pii_type}. "
+                    f"Must be one of {list(_BUILTIN_DETECTORS.keys())} "
+                    "or provide a custom detector."
+                )
+                raise ValueError(msg)
+            self.detector = _BUILTIN_DETECTORS[pii_type]
+        elif isinstance(detector, str):
+            # Custom regex pattern
+            pattern = detector
+
+            def regex_detector(content: str) -> list[PIIMatch]:
+                return [
+                    PIIMatch(
+                        type=pii_type,
+                        value=match.group(),
+                        start=match.start(),
+                        end=match.end(),
+                    )
+                    for match in re.finditer(pattern, content)
+                ]
+
+            self.detector = regex_detector
+        else:
+            # Custom callable detector
+            self.detector = detector
+
+    @property
+    def name(self) -> str:
+        """Name of the middleware."""
+        return f"{self.__class__.__name__}[{self.pii_type}]"
+
+    @hook_config(can_jump_to=["end"])
+    def before_model(  # noqa: PLR0915
+        self,
+        state: AgentState,
+        runtime: Runtime,  # noqa: ARG002
+    ) -> dict[str, Any] | None:
+        """Check user messages and tool results for PII before model invocation.
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            Updated state with PII handled according to strategy, or None if no PII detected.
+
+        Raises:
+            PIIDetectionError: If PII is detected and strategy is "block".
+        """
+        if not self.apply_to_input and not self.apply_to_tool_results:
+            return None
+
+        messages = state["messages"]
+        if not messages:
+            return None
+
+        new_messages = list(messages)
+        any_modified = False
+
+        # Check user input if enabled
+        if self.apply_to_input:
+            # Get last user message
+            last_user_msg = None
+            last_user_idx = None
+            for i in range(len(messages) - 1, -1, -1):
+                if isinstance(messages[i], HumanMessage):
+                    last_user_msg = messages[i]
+                    last_user_idx = i
+                    break
+
+            if last_user_idx is not None and last_user_msg and last_user_msg.content:
+                # Detect PII in message content
+                content = str(last_user_msg.content)
+                matches = self.detector(content)
+
+                if matches:
+                    # Apply strategy
+                    if self.strategy == "block":
+                        raise PIIDetectionError(self.pii_type, matches)
+
+                    if self.strategy == "redact":
+                        new_content = _apply_redact_strategy(content, matches)
+                    elif self.strategy == "mask":
+                        new_content = _apply_mask_strategy(content, matches)
+                    elif self.strategy == "hash":
+                        new_content = _apply_hash_strategy(content, matches)
+                    else:
+                        # Should not reach here due to type hints
+                        msg = f"Unknown strategy: {self.strategy}"
+                        raise ValueError(msg)
+
+                    # Create updated message
+                    updated_message: AnyMessage = HumanMessage(
+                        content=new_content,
+                        id=last_user_msg.id,
+                        name=last_user_msg.name,
+                    )
+
+                    new_messages[last_user_idx] = updated_message
+                    any_modified = True
+
+        # Check tool results if enabled
+        if self.apply_to_tool_results:
+            # Find the last AIMessage, then process all ToolMessages after it
+            last_ai_idx = None
+            for i in range(len(messages) - 1, -1, -1):
+                if isinstance(messages[i], AIMessage):
+                    last_ai_idx = i
+                    break
+
+            if last_ai_idx is not None:
+                # Get all tool messages after the last AI message
+                for i in range(last_ai_idx + 1, len(messages)):
+                    msg = messages[i]
+                    if isinstance(msg, ToolMessage):
+                        tool_msg = msg
+                        if not tool_msg.content:
+                            continue
+
+                        content = str(tool_msg.content)
+                        matches = self.detector(content)
+
+                        if not matches:
+                            continue
+
+                        # Apply strategy
+                        if self.strategy == "block":
+                            raise PIIDetectionError(self.pii_type, matches)
+
+                        if self.strategy == "redact":
+                            new_content = _apply_redact_strategy(content, matches)
+                        elif self.strategy == "mask":
+                            new_content = _apply_mask_strategy(content, matches)
+                        elif self.strategy == "hash":
+                            new_content = _apply_hash_strategy(content, matches)
+                        else:
+                            # Should not reach here due to type hints
+                            msg = f"Unknown strategy: {self.strategy}"
+                            raise ValueError(msg)
+
+                        # Create updated tool message
+                        updated_message = ToolMessage(
+                            content=new_content,
+                            id=tool_msg.id,
+                            name=tool_msg.name,
+                            tool_call_id=tool_msg.tool_call_id,
+                        )
+
+                        new_messages[i] = updated_message
+                        any_modified = True
+
+        if any_modified:
+            return {"messages": new_messages}
+
+        return None
+
+    def after_model(
+        self,
+        state: AgentState,
+        runtime: Runtime,  # noqa: ARG002
+    ) -> dict[str, Any] | None:
+        """Check AI messages for PII after model invocation.
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            Updated state with PII handled according to strategy, or None if no PII detected.
+
+        Raises:
+            PIIDetectionError: If PII is detected and strategy is "block".
+        """
+        if not self.apply_to_output:
+            return None
+
+        messages = state["messages"]
+        if not messages:
+            return None
+
+        # Get last AI message
+        last_ai_msg = None
+        last_ai_idx = None
+        for i in range(len(messages) - 1, -1, -1):
+            msg = messages[i]
+            if isinstance(msg, AIMessage):
+                last_ai_msg = msg
+                last_ai_idx = i
+                break
+
+        if last_ai_idx is None or not last_ai_msg or not last_ai_msg.content:
+            return None
+
+        # Detect PII in message content
+        content = str(last_ai_msg.content)
+        matches = self.detector(content)
+
+        if not matches:
+            return None
+
+        # Apply strategy
+        if self.strategy == "block":
+            raise PIIDetectionError(self.pii_type, matches)
+
+        if self.strategy == "redact":
+            new_content = _apply_redact_strategy(content, matches)
+        elif self.strategy == "mask":
+            new_content = _apply_mask_strategy(content, matches)
+        elif self.strategy == "hash":
+            new_content = _apply_hash_strategy(content, matches)
+        else:
+            # Should not reach here due to type hints
+            msg = f"Unknown strategy: {self.strategy}"
+            raise ValueError(msg)
+
+        # Create updated message
+        updated_message = AIMessage(
+            content=new_content,
+            id=last_ai_msg.id,
+            name=last_ai_msg.name,
+            tool_calls=last_ai_msg.tool_calls,
+        )
+
+        # Return updated messages
+        new_messages = list(messages)
+        new_messages[last_ai_idx] = updated_message
+
+        return {"messages": new_messages}