langchain 1.0.0rc1__py3-none-any.whl → 1.0.1__py3-none-any.whl

langchain/agents/middleware/pii.py

@@ -2,15 +2,22 @@
 
 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._redaction import (
+    PIIDetectionError,
+    PIIMatch,
+    RedactionRule,
+    ResolvedRedactionRule,
+    apply_strategy,
+    detect_credit_card,
+    detect_email,
+    detect_ip,
+    detect_mac_address,
+    detect_url,
+)
 from langchain.agents.middleware.types import AgentMiddleware, AgentState, hook_config
 
 if TYPE_CHECKING:
@@ -19,396 +26,6 @@ if TYPE_CHECKING:
     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.
 
@@ -510,50 +127,34 @@ class PIIMiddleware(AgentMiddleware):
         """
         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
+        self._resolved_rule: ResolvedRedactionRule = RedactionRule(
+            pii_type=pii_type,
+            strategy=strategy,
+            detector=detector,
+        ).resolve()
+        self.pii_type = self._resolved_rule.pii_type
+        self.strategy = self._resolved_rule.strategy
+        self.detector = self._resolved_rule.detector
 
     @property
     def name(self) -> str:
         """Name of the middleware."""
         return f"{self.__class__.__name__}[{self.pii_type}]"
 
+    def _process_content(self, content: str) -> tuple[str, list[PIIMatch]]:
+        """Apply the configured redaction rule to the provided content."""
+        matches = self.detector(content)
+        if not matches:
+            return content, []
+        sanitized = apply_strategy(content, matches, self.strategy)
+        return sanitized, matches
+
     @hook_config(can_jump_to=["end"])
-    def before_model(  # noqa: PLR0915
+    def before_model(
         self,
         state: AgentState,
         runtime: Runtime,  # noqa: ARG002
@@ -594,25 +195,9 @@ class PIIMiddleware(AgentMiddleware):
             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)
+                new_content, matches = self._process_content(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,
@@ -641,26 +226,11 @@ class PIIMiddleware(AgentMiddleware):
                             continue
 
                         content = str(tool_msg.content)
-                        matches = self.detector(content)
+                        new_content, matches = self._process_content(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,
@@ -716,26 +286,11 @@ class PIIMiddleware(AgentMiddleware):
 
         # Detect PII in message content
         content = str(last_ai_msg.content)
-        matches = self.detector(content)
+        new_content, matches = self._process_content(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,
@@ -749,3 +304,14 @@ class PIIMiddleware(AgentMiddleware):
         new_messages[last_ai_idx] = updated_message
 
         return {"messages": new_messages}
+
+
+__all__ = [
+    "PIIDetectionError",
+    "PIIMiddleware",
+    "detect_credit_card",
+    "detect_email",
+    "detect_ip",
+    "detect_mac_address",
+    "detect_url",
+]