langchain 0.4.0.dev0__py3-none-any.whl → 1.0.0__py3-none-any.whl

langchain/agents/middleware/_redaction.py

@@ -0,0 +1,350 @@
+"""Shared redaction utilities for middleware components."""
+
+from __future__ import annotations
+
+import hashlib
+import ipaddress
+import re
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Literal
+from urllib.parse import urlparse
+
+from typing_extensions import TypedDict
+
+RedactionStrategy = Literal["block", "redact", "mask", "hash"]
+"""Supported strategies for handling detected sensitive values."""
+
+
+class PIIMatch(TypedDict):
+    """Represents an individual match of sensitive data."""
+
+    type: str
+    value: str
+    start: int
+    end: int
+
+
+class PIIDetectionError(Exception):
+    """Raised when configured to block on detected sensitive values."""
+
+    def __init__(self, pii_type: str, matches: Sequence[PIIMatch]) -> None:
+        """Initialize the exception with match context.
+
+        Args:
+            pii_type: Name of the detected sensitive type.
+            matches: All matches that were detected for that type.
+        """
+        self.pii_type = pii_type
+        self.matches = list(matches)
+        count = len(matches)
+        msg = f"Detected {count} instance(s) of {pii_type} in text content"
+        super().__init__(msg)
+
+
+Detector = Callable[[str], list[PIIMatch]]
+"""Callable signature for detectors that locate sensitive values."""
+
+
+def detect_email(content: str) -> list[PIIMatch]:
+    """Detect email addresses in content."""
+    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."""
+    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()
+        if _passes_luhn(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 IPv4 or IPv6 addresses in content."""
+    matches: list[PIIMatch] = []
+    ipv4_pattern = r"\b(?:[0-9]{1,3}\.){3}[0-9]{1,3}\b"
+
+    for match in re.finditer(ipv4_pattern, content):
+        ip_candidate = match.group()
+        try:
+            ipaddress.ip_address(ip_candidate)
+        except ValueError:
+            continue
+        matches.append(
+            PIIMatch(
+                type="ip",
+                value=ip_candidate,
+                start=match.start(),
+                end=match.end(),
+            )
+        )
+
+    return matches
+
+
+def detect_mac_address(content: str) -> list[PIIMatch]:
+    """Detect MAC addresses in content."""
+    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."""
+    matches: list[PIIMatch] = []
+
+    # Pattern 1: URLs with scheme (http:// or https://)
+    scheme_pattern = r"https?://[^\s<>\"{}|\\^`\[\]]+"
+
+    for match in re.finditer(scheme_pattern, content):
+        url = match.group()
+        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(),
+                )
+            )
+
+    # 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])?"
+        r"(?:\.[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?)+(?:/[^\s]*)?"
+    )
+
+    for match in re.finditer(bare_pattern, content):
+        start, end = match.start(), match.end()
+        # Skip if already matched with scheme
+        if any(m["start"] <= start < m["end"] or m["start"] < 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."):
+            # 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=start,
+                        end=end,
+                    )
+                )
+
+    return matches
+
+
+BUILTIN_DETECTORS: dict[str, Detector] = {
+    "email": detect_email,
+    "credit_card": detect_credit_card,
+    "ip": detect_ip,
+    "mac_address": detect_mac_address,
+    "url": detect_url,
+}
+"""Registry of built-in detectors keyed by type name."""
+
+
+def _passes_luhn(card_number: str) -> bool:
+    """Validate credit card number using the Luhn checksum."""
+    digits = [int(d) for d in card_number if d.isdigit()]
+    if not 13 <= len(digits) <= 19:
+        return False
+
+    checksum = 0
+    for index, digit in enumerate(reversed(digits)):
+        value = digit
+        if index % 2 == 1:
+            value *= 2
+            if value > 9:
+                value -= 9
+        checksum += value
+    return checksum % 10 == 0
+
+
+def _apply_redact_strategy(content: str, matches: list[PIIMatch]) -> str:
+    result = content
+    for match in sorted(matches, key=lambda item: item["start"], reverse=True):
+        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:
+    result = content
+    for match in sorted(matches, key=lambda item: item["start"], reverse=True):
+        value = match["value"]
+        pii_type = match["type"]
+        if pii_type == "email":
+            parts = value.split("@")
+            if len(parts) == 2:
+                domain_parts = parts[1].split(".")
+                masked = (
+                    f"{parts[0]}@****.{domain_parts[-1]}"
+                    if len(domain_parts) >= 2
+                    else f"{parts[0]}@****"
+                )
+            else:
+                masked = "****"
+        elif pii_type == "credit_card":
+            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":
+            octets = value.split(".")
+            masked = f"*.*.*.{octets[-1]}" if len(octets) == 4 else "****"
+        elif pii_type == "mac_address":
+            separator = ":" if ":" in value else "-"
+            masked = (
+                f"**{separator}**{separator}**{separator}**{separator}**{separator}{value[-2:]}"
+            )
+        elif pii_type == "url":
+            masked = "[MASKED_URL]"
+        else:
+            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:
+    result = content
+    for match in sorted(matches, key=lambda item: item["start"], reverse=True):
+        digest = hashlib.sha256(match["value"].encode()).hexdigest()[:8]
+        replacement = f"<{match['type']}_hash:{digest}>"
+        result = result[: match["start"]] + replacement + result[match["end"] :]
+    return result
+
+
+def apply_strategy(
+    content: str,
+    matches: list[PIIMatch],
+    strategy: RedactionStrategy,
+) -> str:
+    """Apply the configured strategy to matches within content."""
+    if not matches:
+        return content
+    if strategy == "redact":
+        return _apply_redact_strategy(content, matches)
+    if strategy == "mask":
+        return _apply_mask_strategy(content, matches)
+    if strategy == "hash":
+        return _apply_hash_strategy(content, matches)
+    if strategy == "block":
+        raise PIIDetectionError(matches[0]["type"], matches)
+    msg = f"Unknown redaction strategy: {strategy}"
+    raise ValueError(msg)
+
+
+def resolve_detector(pii_type: str, detector: Detector | str | None) -> Detector:
+    """Return a callable detector for the given configuration."""
+    if detector is None:
+        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)
+        return BUILTIN_DETECTORS[pii_type]
+    if isinstance(detector, str):
+        pattern = re.compile(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 pattern.finditer(content)
+            ]
+
+        return regex_detector
+    return detector
+
+
+@dataclass(frozen=True)
+class RedactionRule:
+    """Configuration for handling a single PII type."""
+
+    pii_type: str
+    strategy: RedactionStrategy = "redact"
+    detector: Detector | str | None = None
+
+    def resolve(self) -> ResolvedRedactionRule:
+        """Resolve runtime detector and return an immutable rule."""
+        resolved_detector = resolve_detector(self.pii_type, self.detector)
+        return ResolvedRedactionRule(
+            pii_type=self.pii_type,
+            strategy=self.strategy,
+            detector=resolved_detector,
+        )
+
+
+@dataclass(frozen=True)
+class ResolvedRedactionRule:
+    """Resolved redaction rule ready for execution."""
+
+    pii_type: str
+    strategy: RedactionStrategy
+    detector: Detector
+
+    def apply(self, content: str) -> tuple[str, list[PIIMatch]]:
+        """Apply this rule to content, returning new content and matches."""
+        matches = self.detector(content)
+        if not matches:
+            return content, []
+        updated = apply_strategy(content, matches, self.strategy)
+        return updated, matches
+
+
+__all__ = [
+    "PIIDetectionError",
+    "PIIMatch",
+    "RedactionRule",
+    "ResolvedRedactionRule",
+    "apply_strategy",
+    "detect_credit_card",
+    "detect_email",
+    "detect_ip",
+    "detect_mac_address",
+    "detect_url",
+]

langchain/agents/middleware/context_editing.py

@@ -0,0 +1,274 @@
+"""Context editing middleware.
+
+This middleware mirrors Anthropic's context editing capabilities by clearing
+older tool results once the conversation grows beyond a configurable token
+threshold. The implementation is intentionally model-agnostic so it can be used
+with any LangChain chat model.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Iterable, Sequence
+from dataclasses import dataclass
+from typing import Literal
+
+from langchain_core.messages import (
+    AIMessage,
+    AnyMessage,
+    BaseMessage,
+    SystemMessage,
+    ToolMessage,
+)
+from langchain_core.messages.utils import count_tokens_approximately
+from typing_extensions import Protocol
+
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    ModelCallResult,
+    ModelRequest,
+    ModelResponse,
+)
+
+DEFAULT_TOOL_PLACEHOLDER = "[cleared]"
+
+
+TokenCounter = Callable[
+    [Sequence[BaseMessage]],
+    int,
+]
+
+
+class ContextEdit(Protocol):
+    """Protocol describing a context editing strategy."""
+
+    def apply(
+        self,
+        messages: list[AnyMessage],
+        *,
+        count_tokens: TokenCounter,
+    ) -> None:
+        """Apply an edit to the message list in place."""
+        ...
+
+
+@dataclass(slots=True)
+class ClearToolUsesEdit(ContextEdit):
+    """Configuration for clearing tool outputs when token limits are exceeded."""
+
+    trigger: int = 100_000
+    """Token count that triggers the edit."""
+
+    clear_at_least: int = 0
+    """Minimum number of tokens to reclaim when the edit runs."""
+
+    keep: int = 3
+    """Number of most recent tool results that must be preserved."""
+
+    clear_tool_inputs: bool = False
+    """Whether to clear the originating tool call parameters on the AI message."""
+
+    exclude_tools: Sequence[str] = ()
+    """List of tool names to exclude from clearing."""
+
+    placeholder: str = DEFAULT_TOOL_PLACEHOLDER
+    """Placeholder text inserted for cleared tool outputs."""
+
+    def apply(
+        self,
+        messages: list[AnyMessage],
+        *,
+        count_tokens: TokenCounter,
+    ) -> None:
+        """Apply the clear-tool-uses strategy."""
+        tokens = count_tokens(messages)
+
+        if tokens <= self.trigger:
+            return
+
+        candidates = [
+            (idx, msg) for idx, msg in enumerate(messages) if isinstance(msg, ToolMessage)
+        ]
+
+        if self.keep >= len(candidates):
+            candidates = []
+        elif self.keep:
+            candidates = candidates[: -self.keep]
+
+        cleared_tokens = 0
+        excluded_tools = set(self.exclude_tools)
+
+        for idx, tool_message in candidates:
+            if tool_message.response_metadata.get("context_editing", {}).get("cleared"):
+                continue
+
+            ai_message = next(
+                (m for m in reversed(messages[:idx]) if isinstance(m, AIMessage)), None
+            )
+
+            if ai_message is None:
+                continue
+
+            tool_call = next(
+                (
+                    call
+                    for call in ai_message.tool_calls
+                    if call.get("id") == tool_message.tool_call_id
+                ),
+                None,
+            )
+
+            if tool_call is None:
+                continue
+
+            if (tool_message.name or tool_call["name"]) in excluded_tools:
+                continue
+
+            messages[idx] = tool_message.model_copy(
+                update={
+                    "artifact": None,
+                    "content": self.placeholder,
+                    "response_metadata": {
+                        **tool_message.response_metadata,
+                        "context_editing": {
+                            "cleared": True,
+                            "strategy": "clear_tool_uses",
+                        },
+                    },
+                }
+            )
+
+            if self.clear_tool_inputs:
+                messages[messages.index(ai_message)] = self._build_cleared_tool_input_message(
+                    ai_message,
+                    tool_message.tool_call_id,
+                )
+
+            if self.clear_at_least > 0:
+                new_token_count = count_tokens(messages)
+                cleared_tokens = max(0, tokens - new_token_count)
+                if cleared_tokens >= self.clear_at_least:
+                    break
+
+        return
+
+    def _build_cleared_tool_input_message(
+        self,
+        message: AIMessage,
+        tool_call_id: str,
+    ) -> AIMessage:
+        updated_tool_calls = []
+        cleared_any = False
+        for tool_call in message.tool_calls:
+            updated_call = dict(tool_call)
+            if updated_call.get("id") == tool_call_id:
+                updated_call["args"] = {}
+                cleared_any = True
+            updated_tool_calls.append(updated_call)
+
+        metadata = dict(getattr(message, "response_metadata", {}))
+        context_entry = dict(metadata.get("context_editing", {}))
+        if cleared_any:
+            cleared_ids = set(context_entry.get("cleared_tool_inputs", []))
+            cleared_ids.add(tool_call_id)
+            context_entry["cleared_tool_inputs"] = sorted(cleared_ids)
+            metadata["context_editing"] = context_entry
+
+        return message.model_copy(
+            update={
+                "tool_calls": updated_tool_calls,
+                "response_metadata": metadata,
+            }
+        )
+
+
+class ContextEditingMiddleware(AgentMiddleware):
+    """Middleware that automatically prunes tool results to manage context size.
+
+    The middleware applies a sequence of edits when the total input token count
+    exceeds configured thresholds. Currently the `ClearToolUsesEdit` strategy is
+    supported, aligning with Anthropic's `clear_tool_uses_20250919` behaviour.
+    """
+
+    edits: list[ContextEdit]
+    token_count_method: Literal["approximate", "model"]
+
+    def __init__(
+        self,
+        *,
+        edits: Iterable[ContextEdit] | None = None,
+        token_count_method: Literal["approximate", "model"] = "approximate",  # noqa: S107
+    ) -> None:
+        """Initializes a context editing middleware instance.
+
+        Args:
+            edits: Sequence of edit strategies to apply. Defaults to a single
+                `ClearToolUsesEdit` mirroring Anthropic defaults.
+            token_count_method: Whether to use approximate token counting
+                (faster, less accurate) or exact counting implemented by the
+                chat model (potentially slower, more accurate).
+        """
+        super().__init__()
+        self.edits = list(edits or (ClearToolUsesEdit(),))
+        self.token_count_method = token_count_method
+
+    def wrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelCallResult:
+        """Apply context edits before invoking the model via handler."""
+        if not request.messages:
+            return handler(request)
+
+        if self.token_count_method == "approximate":  # noqa: S105
+
+            def count_tokens(messages: Sequence[BaseMessage]) -> int:
+                return count_tokens_approximately(messages)
+        else:
+            system_msg = (
+                [SystemMessage(content=request.system_prompt)] if request.system_prompt else []
+            )
+
+            def count_tokens(messages: Sequence[BaseMessage]) -> int:
+                return request.model.get_num_tokens_from_messages(
+                    system_msg + list(messages), request.tools
+                )
+
+        for edit in self.edits:
+            edit.apply(request.messages, count_tokens=count_tokens)
+
+        return handler(request)
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        """Apply context edits before invoking the model via handler (async version)."""
+        if not request.messages:
+            return await handler(request)
+
+        if self.token_count_method == "approximate":  # noqa: S105
+
+            def count_tokens(messages: Sequence[BaseMessage]) -> int:
+                return count_tokens_approximately(messages)
+        else:
+            system_msg = (
+                [SystemMessage(content=request.system_prompt)] if request.system_prompt else []
+            )
+
+            def count_tokens(messages: Sequence[BaseMessage]) -> int:
+                return request.model.get_num_tokens_from_messages(
+                    system_msg + list(messages), request.tools
+                )
+
+        for edit in self.edits:
+            edit.apply(request.messages, count_tokens=count_tokens)
+
+        return await handler(request)
+
+
+__all__ = [
+    "ClearToolUsesEdit",
+    "ContextEditingMiddleware",
+]