langchain 1.0.4__py3-none-any.whl → 1.2.3__py3-none-any.whl

langchain/agents/middleware/model_fallback.py

@@ -22,7 +22,7 @@ class ModelFallbackMiddleware(AgentMiddleware):
     """Automatic fallback to alternative models on errors.
 
     Retries failed model calls with alternative models in sequence until
-    success or all models exhausted. Primary model specified in create_agent().
+    success or all models exhausted. Primary model specified in `create_agent`.
 
     Example:
         ```python
@@ -87,15 +87,14 @@ class ModelFallbackMiddleware(AgentMiddleware):
         last_exception: Exception
         try:
             return handler(request)
-        except Exception as e:  # noqa: BLE001
+        except Exception as e:
             last_exception = e
 
         # Try fallback models
         for fallback_model in self.models:
-            request.model = fallback_model
             try:
-                return handler(request)
-            except Exception as e:  # noqa: BLE001
+                return handler(request.override(model=fallback_model))
+            except Exception as e:
                 last_exception = e
                 continue
 
@@ -122,15 +121,14 @@ class ModelFallbackMiddleware(AgentMiddleware):
         last_exception: Exception
         try:
             return await handler(request)
-        except Exception as e:  # noqa: BLE001
+        except Exception as e:
             last_exception = e
 
         # Try fallback models
         for fallback_model in self.models:
-            request.model = fallback_model
             try:
-                return await handler(request)
-            except Exception as e:  # noqa: BLE001
+                return await handler(request.override(model=fallback_model))
+            except Exception as e:
                 last_exception = e
                 continue
 

langchain/agents/middleware/model_retry.py

@@ -0,0 +1,300 @@
+"""Model retry middleware for agents."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from typing import TYPE_CHECKING
+
+from langchain_core.messages import AIMessage
+
+from langchain.agents.middleware._retry import (
+    OnFailure,
+    RetryOn,
+    calculate_delay,
+    should_retry_exception,
+    validate_retry_params,
+)
+from langchain.agents.middleware.types import AgentMiddleware, ModelResponse
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from langchain.agents.middleware.types import ModelRequest
+
+
+class ModelRetryMiddleware(AgentMiddleware):
+    """Middleware that automatically retries failed model calls with configurable backoff.
+
+    Supports retrying on specific exceptions and exponential backoff.
+
+    Examples:
+        !!! example "Basic usage with default settings (2 retries, exponential backoff)"
+
+            ```python
+            from langchain.agents import create_agent
+            from langchain.agents.middleware import ModelRetryMiddleware
+
+            agent = create_agent(model, tools=[search_tool], middleware=[ModelRetryMiddleware()])
+            ```
+
+        !!! example "Retry specific exceptions only"
+
+            ```python
+            from anthropic import RateLimitError
+            from openai import APITimeoutError
+
+            retry = ModelRetryMiddleware(
+                max_retries=4,
+                retry_on=(APITimeoutError, RateLimitError),
+                backoff_factor=1.5,
+            )
+            ```
+
+        !!! example "Custom exception filtering"
+
+            ```python
+            from anthropic import APIStatusError
+
+
+            def should_retry(exc: Exception) -> bool:
+                # Only retry on 5xx errors
+                if isinstance(exc, APIStatusError):
+                    return 500 <= exc.status_code < 600
+                return False
+
+
+            retry = ModelRetryMiddleware(
+                max_retries=3,
+                retry_on=should_retry,
+            )
+            ```
+
+        !!! example "Custom error handling"
+
+            ```python
+            def format_error(exc: Exception) -> str:
+                return "Model temporarily unavailable. Please try again later."
+
+
+            retry = ModelRetryMiddleware(
+                max_retries=4,
+                on_failure=format_error,
+            )
+            ```
+
+        !!! example "Constant backoff (no exponential growth)"
+
+            ```python
+            retry = ModelRetryMiddleware(
+                max_retries=5,
+                backoff_factor=0.0,  # No exponential growth
+                initial_delay=2.0,  # Always wait 2 seconds
+            )
+            ```
+
+        !!! example "Raise exception on failure"
+
+            ```python
+            retry = ModelRetryMiddleware(
+                max_retries=2,
+                on_failure="error",  # Re-raise exception instead of returning message
+            )
+            ```
+    """
+
+    def __init__(
+        self,
+        *,
+        max_retries: int = 2,
+        retry_on: RetryOn = (Exception,),
+        on_failure: OnFailure = "continue",
+        backoff_factor: float = 2.0,
+        initial_delay: float = 1.0,
+        max_delay: float = 60.0,
+        jitter: bool = True,
+    ) -> None:
+        """Initialize `ModelRetryMiddleware`.
+
+        Args:
+            max_retries: Maximum number of retry attempts after the initial call.
+
+                Must be `>= 0`.
+            retry_on: Either a tuple of exception types to retry on, or a callable
+                that takes an exception and returns `True` if it should be retried.
+
+                Default is to retry on all exceptions.
+            on_failure: Behavior when all retries are exhausted.
+
+                Options:
+
+                - `'continue'`: Return an `AIMessage` with error details,
+                    allowing the agent to continue with an error response.
+                - `'error'`: Re-raise the exception, stopping agent execution.
+                - **Custom callable:** Function that takes the exception and returns a
+                    string for the `AIMessage` content, allowing custom error
+                    formatting.
+            backoff_factor: Multiplier for exponential backoff.
+
+                Each retry waits `initial_delay * (backoff_factor ** retry_number)`
+                seconds.
+
+                Set to `0.0` for constant delay.
+            initial_delay: Initial delay in seconds before first retry.
+            max_delay: Maximum delay in seconds between retries.
+
+                Caps exponential backoff growth.
+            jitter: Whether to add random jitter (`±25%`) to delay to avoid thundering herd.
+
+        Raises:
+            ValueError: If `max_retries < 0` or delays are negative.
+        """
+        super().__init__()
+
+        # Validate parameters
+        validate_retry_params(max_retries, initial_delay, max_delay, backoff_factor)
+
+        self.max_retries = max_retries
+        self.tools = []  # No additional tools registered by this middleware
+        self.retry_on = retry_on
+        self.on_failure = on_failure
+        self.backoff_factor = backoff_factor
+        self.initial_delay = initial_delay
+        self.max_delay = max_delay
+        self.jitter = jitter
+
+    def _format_failure_message(self, exc: Exception, attempts_made: int) -> AIMessage:
+        """Format the failure message when retries are exhausted.
+
+        Args:
+            exc: The exception that caused the failure.
+            attempts_made: Number of attempts actually made.
+
+        Returns:
+            `AIMessage` with formatted error message.
+        """
+        exc_type = type(exc).__name__
+        exc_msg = str(exc)
+        attempt_word = "attempt" if attempts_made == 1 else "attempts"
+        content = (
+            f"Model call failed after {attempts_made} {attempt_word} with {exc_type}: {exc_msg}"
+        )
+        return AIMessage(content=content)
+
+    def _handle_failure(self, exc: Exception, attempts_made: int) -> ModelResponse:
+        """Handle failure when all retries are exhausted.
+
+        Args:
+            exc: The exception that caused the failure.
+            attempts_made: Number of attempts actually made.
+
+        Returns:
+            `ModelResponse` with error details.
+
+        Raises:
+            Exception: If `on_failure` is `'error'`, re-raises the exception.
+        """
+        if self.on_failure == "error":
+            raise exc
+
+        if callable(self.on_failure):
+            content = self.on_failure(exc)
+            ai_msg = AIMessage(content=content)
+        else:
+            ai_msg = self._format_failure_message(exc, attempts_made)
+
+        return ModelResponse(result=[ai_msg])
+
+    def wrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelResponse | AIMessage:
+        """Intercept model execution and retry on failure.
+
+        Args:
+            request: Model request with model, messages, state, and runtime.
+            handler: Callable to execute the model (can be called multiple times).
+
+        Returns:
+            `ModelResponse` or `AIMessage` (the final result).
+        """
+        # Initial attempt + retries
+        for attempt in range(self.max_retries + 1):
+            try:
+                return handler(request)
+            except Exception as exc:
+                attempts_made = attempt + 1  # attempt is 0-indexed
+
+                # Check if we should retry this exception
+                if not should_retry_exception(exc, self.retry_on):
+                    # Exception is not retryable, handle failure immediately
+                    return self._handle_failure(exc, attempts_made)
+
+                # Check if we have more retries left
+                if attempt < self.max_retries:
+                    # Calculate and apply backoff delay
+                    delay = calculate_delay(
+                        attempt,
+                        backoff_factor=self.backoff_factor,
+                        initial_delay=self.initial_delay,
+                        max_delay=self.max_delay,
+                        jitter=self.jitter,
+                    )
+                    if delay > 0:
+                        time.sleep(delay)
+                    # Continue to next retry
+                else:
+                    # No more retries, handle failure
+                    return self._handle_failure(exc, attempts_made)
+
+        # Unreachable: loop always returns via handler success or _handle_failure
+        msg = "Unexpected: retry loop completed without returning"
+        raise RuntimeError(msg)
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelResponse | AIMessage:
+        """Intercept and control async model execution with retry logic.
+
+        Args:
+            request: Model request with model, messages, state, and runtime.
+            handler: Async callable to execute the model and returns `ModelResponse`.
+
+        Returns:
+            `ModelResponse` or `AIMessage` (the final result).
+        """
+        # Initial attempt + retries
+        for attempt in range(self.max_retries + 1):
+            try:
+                return await handler(request)
+            except Exception as exc:
+                attempts_made = attempt + 1  # attempt is 0-indexed
+
+                # Check if we should retry this exception
+                if not should_retry_exception(exc, self.retry_on):
+                    # Exception is not retryable, handle failure immediately
+                    return self._handle_failure(exc, attempts_made)
+
+                # Check if we have more retries left
+                if attempt < self.max_retries:
+                    # Calculate and apply backoff delay
+                    delay = calculate_delay(
+                        attempt,
+                        backoff_factor=self.backoff_factor,
+                        initial_delay=self.initial_delay,
+                        max_delay=self.max_delay,
+                        jitter=self.jitter,
+                    )
+                    if delay > 0:
+                        await asyncio.sleep(delay)
+                    # Continue to next retry
+                else:
+                    # No more retries, handle failure
+                    return self._handle_failure(exc, attempts_made)
+
+        # Unreachable: loop always returns via handler success or _handle_failure
+        msg = "Unexpected: retry loop completed without returning"
+        raise RuntimeError(msg)

langchain/agents/middleware/pii.py

@@ -5,6 +5,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Literal
 
 from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage
+from typing_extensions import override
 
 from langchain.agents.middleware._redaction import (
     PIIDetectionError,
@@ -27,24 +28,26 @@ if TYPE_CHECKING:
 
 
 class PIIMiddleware(AgentMiddleware):
-    """Detect and handle Personally Identifiable Information (PII) in agent conversations.
+    """Detect and handle Personally Identifiable Information (PII) in 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.
+    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)
+
+    - `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>`)
+
+    - `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:
 
@@ -90,6 +93,8 @@ class PIIMiddleware(AgentMiddleware):
 
     def __init__(
         self,
+        # From a typing point of view, the literals are covered by 'str'.
+        # Nonetheless, we escape PYI051 to keep hints and autocompletion for the caller.
         pii_type: Literal["email", "credit_card", "ip", "mac_address", "url"] | str,  # noqa: PYI051
         *,
         strategy: Literal["block", "redact", "mask", "hash"] = "redact",
@@ -101,12 +106,15 @@ class PIIMiddleware(AgentMiddleware):
         """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:
+            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.
+
+                Options:
 
-                * `block`: Raise PIIDetectionError when PII is detected
+                * `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>`)
@@ -114,16 +122,15 @@ class PIIMiddleware(AgentMiddleware):
             detector: Custom detector function or regex pattern.
 
                 * If `Callable`: Function that takes content string and returns
-                    list of PIIMatch objects
+                    list of `PIIMatch` objects
                 * If `str`: Regex pattern to match PII
-                * If `None`: Uses built-in detector for the pii_type
-
+                * 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.
+            ValueError: If `pii_type` is not built-in and no detector is provided.
         """
         super().__init__()
 
@@ -154,10 +161,11 @@ class PIIMiddleware(AgentMiddleware):
         return sanitized, matches
 
     @hook_config(can_jump_to=["end"])
+    @override
     def before_model(
         self,
         state: AgentState,
-        runtime: Runtime,  # noqa: ARG002
+        runtime: Runtime,
     ) -> dict[str, Any] | None:
         """Check user messages and tool results for PII before model invocation.
 
@@ -166,10 +174,11 @@ class PIIMiddleware(AgentMiddleware):
             runtime: The langgraph runtime.
 
         Returns:
-            Updated state with PII handled according to strategy, or None if no PII detected.
+            Updated state with PII handled according to strategy, or `None` if no PII
+                detected.
 
         Raises:
-            PIIDetectionError: If PII is detected and strategy is "block".
+            PIIDetectionError: If PII is detected and strategy is `'block'`.
         """
         if not self.apply_to_input and not self.apply_to_tool_results:
             return None
@@ -247,10 +256,32 @@ class PIIMiddleware(AgentMiddleware):
 
         return None
 
+    @hook_config(can_jump_to=["end"])
+    async def abefore_model(
+        self,
+        state: AgentState,
+        runtime: Runtime,
+    ) -> dict[str, Any] | None:
+        """Async 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'`.
+        """
+        return self.before_model(state, runtime)
+
+    @override
     def after_model(
         self,
         state: AgentState,
-        runtime: Runtime,  # noqa: ARG002
+        runtime: Runtime,
     ) -> dict[str, Any] | None:
         """Check AI messages for PII after model invocation.
 
@@ -259,10 +290,11 @@ class PIIMiddleware(AgentMiddleware):
             runtime: The langgraph runtime.
 
         Returns:
-            Updated state with PII handled according to strategy, or None if no PII detected.
+            Updated state with PII handled according to strategy, or None if no PII
+                detected.
 
         Raises:
-            PIIDetectionError: If PII is detected and strategy is "block".
+            PIIDetectionError: If PII is detected and strategy is `'block'`.
         """
         if not self.apply_to_output:
             return None
@@ -305,9 +337,30 @@ class PIIMiddleware(AgentMiddleware):
 
         return {"messages": new_messages}
 
+    async def aafter_model(
+        self,
+        state: AgentState,
+        runtime: Runtime,
+    ) -> dict[str, Any] | None:
+        """Async 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'`.
+        """
+        return self.after_model(state, runtime)
+
 
 __all__ = [
     "PIIDetectionError",
+    "PIIMatch",
     "PIIMiddleware",
     "detect_credit_card",
     "detect_email",