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

langchain/agents/middleware/tool_retry.py

@@ -3,12 +3,19 @@
 from __future__ import annotations
 
 import asyncio
-import random
 import time
-from typing import TYPE_CHECKING, Literal
+import warnings
+from typing import TYPE_CHECKING
 
 from langchain_core.messages import ToolMessage
 
+from langchain.agents.middleware._retry import (
+    OnFailure,
+    RetryOn,
+    calculate_delay,
+    should_retry_exception,
+    validate_retry_params,
+)
 from langchain.agents.middleware.types import AgentMiddleware
 
 if TYPE_CHECKING:
@@ -26,89 +33,96 @@ class ToolRetryMiddleware(AgentMiddleware):
     Supports retrying on specific exceptions and exponential backoff.
 
     Examples:
-        Basic usage with default settings (2 retries, exponential backoff):
-        ```python
-        from langchain.agents import create_agent
-        from langchain.agents.middleware import ToolRetryMiddleware
-
-        agent = create_agent(model, tools=[search_tool], middleware=[ToolRetryMiddleware()])
-        ```
-
-        Retry specific exceptions only:
-        ```python
-        from requests.exceptions import RequestException, Timeout
-
-        retry = ToolRetryMiddleware(
-            max_retries=4,
-            retry_on=(RequestException, Timeout),
-            backoff_factor=1.5,
-        )
-        ```
+        !!! example "Basic usage with default settings (2 retries, exponential backoff)"
 
-        Custom exception filtering:
-        ```python
-        from requests.exceptions import HTTPError
+            ```python
+            from langchain.agents import create_agent
+            from langchain.agents.middleware import ToolRetryMiddleware
 
+            agent = create_agent(model, tools=[search_tool], middleware=[ToolRetryMiddleware()])
+            ```
 
-        def should_retry(exc: Exception) -> bool:
-            # Only retry on 5xx errors
-            if isinstance(exc, HTTPError):
-                return 500 <= exc.status_code < 600
-            return False
+        !!! example "Retry specific exceptions only"
 
+            ```python
+            from requests.exceptions import RequestException, Timeout
 
-        retry = ToolRetryMiddleware(
-            max_retries=3,
-            retry_on=should_retry,
-        )
-        ```
+            retry = ToolRetryMiddleware(
+                max_retries=4,
+                retry_on=(RequestException, Timeout),
+                backoff_factor=1.5,
+            )
+            ```
 
-        Apply to specific tools with custom error handling:
-        ```python
-        def format_error(exc: Exception) -> str:
-            return "Database temporarily unavailable. Please try again later."
+        !!! example "Custom exception filtering"
 
+            ```python
+            from requests.exceptions import HTTPError
 
-        retry = ToolRetryMiddleware(
-            max_retries=4,
-            tools=["search_database"],
-            on_failure=format_error,
-        )
-        ```
 
-        Apply to specific tools using BaseTool instances:
-        ```python
-        from langchain_core.tools import tool
+            def should_retry(exc: Exception) -> bool:
+                # Only retry on 5xx errors
+                if isinstance(exc, HTTPError):
+                    return 500 <= exc.status_code < 600
+                return False
 
 
-        @tool
-        def search_database(query: str) -> str:
-            '''Search the database.'''
-            return results
+            retry = ToolRetryMiddleware(
+                max_retries=3,
+                retry_on=should_retry,
+            )
+            ```
 
+        !!! example "Apply to specific tools with custom error handling"
 
-        retry = ToolRetryMiddleware(
-            max_retries=4,
-            tools=[search_database],  # Pass BaseTool instance
-        )
-        ```
-
-        Constant backoff (no exponential growth):
-        ```python
-        retry = ToolRetryMiddleware(
-            max_retries=5,
-            backoff_factor=0.0,  # No exponential growth
-            initial_delay=2.0,  # Always wait 2 seconds
-        )
-        ```
+            ```python
+            def format_error(exc: Exception) -> str:
+                return "Database temporarily unavailable. Please try again later."
 
-        Raise exception on failure:
-        ```python
-        retry = ToolRetryMiddleware(
-            max_retries=2,
-            on_failure="raise",  # Re-raise exception instead of returning message
-        )
-        ```
+
+            retry = ToolRetryMiddleware(
+                max_retries=4,
+                tools=["search_database"],
+                on_failure=format_error,
+            )
+            ```
+
+        !!! example "Apply to specific tools using `BaseTool` instances"
+
+            ```python
+            from langchain_core.tools import tool
+
+
+            @tool
+            def search_database(query: str) -> str:
+                '''Search the database.'''
+                return results
+
+
+            retry = ToolRetryMiddleware(
+                max_retries=4,
+                tools=[search_database],  # Pass BaseTool instance
+            )
+            ```
+
+        !!! example "Constant backoff (no exponential growth)"
+
+            ```python
+            retry = ToolRetryMiddleware(
+                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 = ToolRetryMiddleware(
+                max_retries=2,
+                on_failure="error",  # Re-raise exception instead of returning message
+            )
+            ```
     """
 
     def __init__(
@@ -116,59 +130,78 @@ class ToolRetryMiddleware(AgentMiddleware):
         *,
         max_retries: int = 2,
         tools: list[BaseTool | str] | None = None,
-        retry_on: tuple[type[Exception], ...] | Callable[[Exception], bool] = (Exception,),
-        on_failure: (
-            Literal["raise", "return_message"] | Callable[[Exception], str]
-        ) = "return_message",
+        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 ToolRetryMiddleware.
+        """Initialize `ToolRetryMiddleware`.
 
         Args:
             max_retries: Maximum number of retry attempts after the initial call.
-                Default is 2 retries (3 total attempts). Must be >= 0.
+
+                Must be `>= 0`.
             tools: Optional list of tools or tool names to apply retry logic to.
+
                 Can be a list of `BaseTool` instances or tool name strings.
-                If `None`, applies to all tools. Default is `None`.
+
+                If `None`, applies to all tools.
             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:
-                - `"return_message"` (default): Return a ToolMessage with error details,
-                  allowing the LLM to handle the failure and potentially recover.
-                - `"raise"`: Re-raise the exception, stopping agent execution.
-                - Custom callable: Function that takes the exception and returns a string
-                  for the ToolMessage 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. Default is 2.0.
-            initial_delay: Initial delay in seconds before first retry. Default is 1.0.
-            max_delay: Maximum delay in seconds between retries. Caps exponential
-                backoff growth. Default is 60.0.
-            jitter: Whether to add random jitter (±25%) to delay to avoid thundering herd.
-                Default is `True`.
+            on_failure: Behavior when all retries are exhausted.
+
+                Options:
+
+                - `'continue'`: Return a `ToolMessage` with error details,
+                    allowing the LLM to handle the failure and potentially recover.
+                - `'error'`: Re-raise the exception, stopping agent execution.
+                - **Custom callable:** Function that takes the exception and returns a
+                    string for the `ToolMessage` content, allowing custom error
+                    formatting.
+
+                **Deprecated values** (for backwards compatibility):
+
+                - `'return_message'`: Use `'continue'` instead.
+                - `'raise'`: Use `'error'` instead.
+            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.
+            ValueError: If `max_retries < 0` or delays are negative.
         """
         super().__init__()
 
         # Validate parameters
-        if max_retries < 0:
-            msg = "max_retries must be >= 0"
-            raise ValueError(msg)
-        if initial_delay < 0:
-            msg = "initial_delay must be >= 0"
-            raise ValueError(msg)
-        if max_delay < 0:
-            msg = "max_delay must be >= 0"
-            raise ValueError(msg)
-        if backoff_factor < 0:
-            msg = "backoff_factor must be >= 0"
-            raise ValueError(msg)
+        validate_retry_params(max_retries, initial_delay, max_delay, backoff_factor)
+
+        # Handle backwards compatibility for deprecated on_failure values
+        if on_failure == "raise":  # type: ignore[comparison-overlap]
+            msg = (
+                "on_failure='raise' is deprecated and will be removed in a future version. "
+                "Use on_failure='error' instead."
+            )
+            warnings.warn(msg, DeprecationWarning, stacklevel=2)
+            on_failure = "error"
+        elif on_failure == "return_message":  # type: ignore[comparison-overlap]
+            msg = (
+                "on_failure='return_message' is deprecated and will be removed "
+                "in a future version. Use on_failure='continue' instead."
+            )
+            warnings.warn(msg, DeprecationWarning, stacklevel=2)
+            on_failure = "continue"
 
         self.max_retries = max_retries
 
@@ -200,44 +233,6 @@ class ToolRetryMiddleware(AgentMiddleware):
             return True
         return tool_name in self._tool_filter
 
-    def _should_retry_exception(self, exc: Exception) -> bool:
-        """Check if the exception should trigger a retry.
-
-        Args:
-            exc: The exception that occurred.
-
-        Returns:
-            `True` if the exception should be retried, `False` otherwise.
-        """
-        if callable(self.retry_on):
-            return self.retry_on(exc)
-        return isinstance(exc, self.retry_on)
-
-    def _calculate_delay(self, retry_number: int) -> float:
-        """Calculate delay for the given retry attempt.
-
-        Args:
-            retry_number: The retry attempt number (0-indexed).
-
-        Returns:
-            Delay in seconds before next retry.
-        """
-        if self.backoff_factor == 0.0:
-            delay = self.initial_delay
-        else:
-            delay = self.initial_delay * (self.backoff_factor**retry_number)
-
-        # Cap at max_delay
-        delay = min(delay, self.max_delay)
-
-        if self.jitter and delay > 0:
-            jitter_amount = delay * 0.25
-            delay = delay + random.uniform(-jitter_amount, jitter_amount)  # noqa: S311
-            # Ensure delay is not negative after jitter
-            delay = max(0, delay)
-
-        return delay
-
     def _format_failure_message(self, tool_name: str, exc: Exception, attempts_made: int) -> str:
         """Format the failure message when retries are exhausted.
 
@@ -250,8 +245,12 @@ class ToolRetryMiddleware(AgentMiddleware):
             Formatted error message string.
         """
         exc_type = type(exc).__name__
+        exc_msg = str(exc)
         attempt_word = "attempt" if attempts_made == 1 else "attempts"
-        return f"Tool '{tool_name}' failed after {attempts_made} {attempt_word} with {exc_type}"
+        return (
+            f"Tool '{tool_name}' failed after {attempts_made} {attempt_word} "
+            f"with {exc_type}: {exc_msg}. Please try again."
+        )
 
     def _handle_failure(
         self, tool_name: str, tool_call_id: str | None, exc: Exception, attempts_made: int
@@ -260,17 +259,17 @@ class ToolRetryMiddleware(AgentMiddleware):
 
         Args:
             tool_name: Name of the tool that failed.
-            tool_call_id: ID of the tool call (may be None).
+            tool_call_id: ID of the tool call (may be `None`).
             exc: The exception that caused the failure.
             attempts_made: Number of attempts actually made.
 
         Returns:
-            ToolMessage with error details.
+            `ToolMessage` with error details.
 
         Raises:
-            Exception: If on_failure is "raise", re-raises the exception.
+            Exception: If `on_failure` is `'error'`, re-raises the exception.
         """
-        if self.on_failure == "raise":
+        if self.on_failure == "error":
             raise exc
 
         if callable(self.on_failure):
@@ -293,11 +292,11 @@ class ToolRetryMiddleware(AgentMiddleware):
         """Intercept tool execution and retry on failure.
 
         Args:
-            request: Tool call request with call dict, BaseTool, state, and runtime.
+            request: Tool call request with call dict, `BaseTool`, state, and runtime.
             handler: Callable to execute the tool (can be called multiple times).
 
         Returns:
-            ToolMessage or Command (the final result).
+            `ToolMessage` or `Command` (the final result).
         """
         tool_name = request.tool.name if request.tool else request.tool_call["name"]
 
@@ -311,18 +310,24 @@ class ToolRetryMiddleware(AgentMiddleware):
         for attempt in range(self.max_retries + 1):
             try:
                 return handler(request)
-            except Exception as exc:  # noqa: BLE001
+            except Exception as exc:
                 attempts_made = attempt + 1  # attempt is 0-indexed
 
                 # Check if we should retry this exception
-                if not self._should_retry_exception(exc):
+                if not should_retry_exception(exc, self.retry_on):
                     # Exception is not retryable, handle failure immediately
                     return self._handle_failure(tool_name, tool_call_id, exc, attempts_made)
 
                 # Check if we have more retries left
                 if attempt < self.max_retries:
                     # Calculate and apply backoff delay
-                    delay = self._calculate_delay(attempt)
+                    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
@@ -342,11 +347,12 @@ class ToolRetryMiddleware(AgentMiddleware):
         """Intercept and control async tool execution with retry logic.
 
         Args:
-            request: Tool call request with call dict, BaseTool, state, and runtime.
-            handler: Async callable to execute the tool and returns ToolMessage or Command.
+            request: Tool call request with call `dict`, `BaseTool`, state, and runtime.
+            handler: Async callable to execute the tool and returns `ToolMessage` or
+                `Command`.
 
         Returns:
-            ToolMessage or Command (the final result).
+            `ToolMessage` or `Command` (the final result).
         """
         tool_name = request.tool.name if request.tool else request.tool_call["name"]
 
@@ -360,18 +366,24 @@ class ToolRetryMiddleware(AgentMiddleware):
         for attempt in range(self.max_retries + 1):
             try:
                 return await handler(request)
-            except Exception as exc:  # noqa: BLE001
+            except Exception as exc:
                 attempts_made = attempt + 1  # attempt is 0-indexed
 
                 # Check if we should retry this exception
-                if not self._should_retry_exception(exc):
+                if not should_retry_exception(exc, self.retry_on):
                     # Exception is not retryable, handle failure immediately
                     return self._handle_failure(tool_name, tool_call_id, exc, attempts_made)
 
                 # Check if we have more retries left
                 if attempt < self.max_retries:
                     # Calculate and apply backoff delay
-                    delay = self._calculate_delay(attempt)
+                    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

langchain/agents/middleware/tool_selection.py

@@ -49,14 +49,15 @@ def _create_tool_selection_response(tools: list[BaseTool]) -> TypeAdapter:
         tools: Available tools to include in the schema.
 
     Returns:
-        TypeAdapter for a schema where each tool name is a Literal with its description.
+        `TypeAdapter` for a schema where each tool name is a `Literal` with its
+            description.
     """
     if not tools:
         msg = "Invalid usage: tools must be non-empty"
         raise AssertionError(msg)
 
     # Create a Union of Annotated Literal types for each tool name with description
-    # Example: Union[Annotated[Literal["tool1"], Field(description="...")], ...] noqa: ERA001
+    # For instance: Union[Annotated[Literal["tool1"], Field(description="...")], ...]
     literals = [
         Annotated[Literal[tool.name], Field(description=tool.description)] for tool in tools
     ]
@@ -92,23 +93,25 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
     and helps the main model focus on the right tools.
 
     Examples:
-        Limit to 3 tools:
-        ```python
-        from langchain.agents.middleware import LLMToolSelectorMiddleware
+        !!! example "Limit to 3 tools"
 
-        middleware = LLMToolSelectorMiddleware(max_tools=3)
+            ```python
+            from langchain.agents.middleware import LLMToolSelectorMiddleware
 
-        agent = create_agent(
-            model="openai:gpt-4o",
-            tools=[tool1, tool2, tool3, tool4, tool5],
-            middleware=[middleware],
-        )
-        ```
+            middleware = LLMToolSelectorMiddleware(max_tools=3)
+
+            agent = create_agent(
+                model="openai:gpt-4o",
+                tools=[tool1, tool2, tool3, tool4, tool5],
+                middleware=[middleware],
+            )
+            ```
 
-        Use a smaller model for selection:
-        ```python
-        middleware = LLMToolSelectorMiddleware(model="openai:gpt-4o-mini", max_tools=2)
-        ```
+        !!! example "Use a smaller model for selection"
+
+            ```python
+            middleware = LLMToolSelectorMiddleware(model="openai:gpt-4o-mini", max_tools=2)
+            ```
     """
 
     def __init__(
@@ -122,13 +125,20 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         """Initialize the tool selector.
 
         Args:
-            model: Model to use for selection. If not provided, uses the agent's main model.
-                Can be a model identifier string or BaseChatModel instance.
+            model: Model to use for selection.
+
+                If not provided, uses the agent's main model.
+
+                Can be a model identifier string or `BaseChatModel` instance.
             system_prompt: Instructions for the selection model.
-            max_tools: Maximum number of tools to select. If the model selects more,
-                only the first max_tools will be used. No limit if not specified.
+            max_tools: Maximum number of tools to select.
+
+                If the model selects more, only the first `max_tools` will be used.
+
+                If not specified, there is no limit.
             always_include: Tool names to always include regardless of selection.
-                These do not count against the max_tools limit.
+
+                These do not count against the `max_tools` limit.
         """
         super().__init__()
         self.system_prompt = system_prompt
@@ -144,7 +154,8 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         """Prepare inputs for tool selection.
 
         Returns:
-            SelectionRequest with prepared inputs, or None if no selection is needed.
+            `SelectionRequest` with prepared inputs, or `None` if no selection is
+                needed.
         """
         # If no tools available, return None
         if not request.tools or len(request.tools) == 0:
@@ -211,7 +222,7 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         valid_tool_names: list[str],
         request: ModelRequest,
     ) -> ModelRequest:
-        """Process the selection response and return filtered ModelRequest."""
+        """Process the selection response and return filtered `ModelRequest`."""
         selected_tool_names: list[str] = []
         invalid_tool_selections = []
 
@@ -244,8 +255,7 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         # Also preserve any provider-specific tool dicts from the original request
         provider_tools = [tool for tool in request.tools if isinstance(tool, dict)]
 
-        request.tools = [*selected_tools, *provider_tools]
-        return request
+        return request.override(tools=[*selected_tools, *provider_tools])
 
     def wrap_model_call(
         self,
@@ -272,7 +282,7 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         # Response should be a dict since we're passing a schema (not a Pydantic model class)
         if not isinstance(response, dict):
             msg = f"Expected dict response, got {type(response)}"
-            raise AssertionError(msg)
+            raise AssertionError(msg)  # noqa: TRY004
         modified_request = self._process_selection_response(
             response, selection_request.available_tools, selection_request.valid_tool_names, request
         )
@@ -303,7 +313,7 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         # Response should be a dict since we're passing a schema (not a Pydantic model class)
         if not isinstance(response, dict):
             msg = f"Expected dict response, got {type(response)}"
-            raise AssertionError(msg)
+            raise AssertionError(msg)  # noqa: TRY004
         modified_request = self._process_selection_response(
             response, selection_request.available_tools, selection_request.valid_tool_names, request
         )