langchain 1.0.5__py3-none-any.whl → 1.2.4__py3-none-any.whl

langchain/agents/middleware/tool_selection.py

@@ -4,12 +4,7 @@ from __future__ import annotations
 
 import logging
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Annotated, Literal, Union
-
-if TYPE_CHECKING:
-    from collections.abc import Awaitable, Callable
-
-    from langchain.tools import BaseTool
+from typing import TYPE_CHECKING, Annotated, Any, Literal, Union
 
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage
@@ -24,6 +19,11 @@ from langchain.agents.middleware.types import (
 )
 from langchain.chat_models.base import init_chat_model
 
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from langchain.tools import BaseTool
+
 logger = logging.getLogger(__name__)
 
 DEFAULT_SYSTEM_PROMPT = (
@@ -42,21 +42,25 @@ class _SelectionRequest:
     valid_tool_names: list[str]
 
 
-def _create_tool_selection_response(tools: list[BaseTool]) -> TypeAdapter:
+def _create_tool_selection_response(tools: list[BaseTool]) -> TypeAdapter[Any]:
     """Create a structured output schema for tool selection.
 
     Args:
         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.
+
+    Raises:
+        AssertionError: If `tools` is empty.
     """
     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 +96,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],
+            )
+            ```
+
+        !!! example "Use a smaller model for selection"
 
-        Use a smaller model for selection:
-        ```python
-        middleware = LLMToolSelectorMiddleware(model="openai:gpt-4o-mini", max_tools=2)
-        ```
+            ```python
+            middleware = LLMToolSelectorMiddleware(model="openai:gpt-4o-mini", max_tools=2)
+            ```
     """
 
     def __init__(
@@ -122,13 +128,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
@@ -143,8 +156,16 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
     def _prepare_selection_request(self, request: ModelRequest) -> _SelectionRequest | None:
         """Prepare inputs for tool selection.
 
+        Args:
+            request: the model request.
+
         Returns:
-            SelectionRequest with prepared inputs, or None if no selection is needed.
+            `SelectionRequest` with prepared inputs, or `None` if no selection is
+            needed.
+
+        Raises:
+            ValueError: If tools in `always_include` are not found in the request.
+            AssertionError: If no user message is found in the request messages.
         """
         # If no tools available, return None
         if not request.tools or len(request.tools) == 0:
@@ -206,12 +227,12 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
 
     def _process_selection_response(
         self,
-        response: dict,
+        response: dict[str, Any],
         available_tools: list[BaseTool],
         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,15 +265,26 @@ 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,
         request: ModelRequest,
         handler: Callable[[ModelRequest], ModelResponse],
     ) -> ModelCallResult:
-        """Filter tools based on LLM selection before invoking the model via handler."""
+        """Filter tools based on LLM selection before invoking the model via handler.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The model call result.
+
+        Raises:
+            AssertionError: If the selection model response is not a dict.
+        """
         selection_request = self._prepare_selection_request(request)
         if selection_request is None:
             return handler(request)
@@ -272,7 +304,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
         )
@@ -283,7 +315,19 @@ class LLMToolSelectorMiddleware(AgentMiddleware):
         request: ModelRequest,
         handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
     ) -> ModelCallResult:
-        """Filter tools based on LLM selection before invoking the model via handler."""
+        """Filter tools based on LLM selection before invoking the model via handler.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The model call result.
+
+        Raises:
+            AssertionError: If the selection model response is not a dict.
+        """
         selection_request = self._prepare_selection_request(request)
         if selection_request is None:
             return await handler(request)
@@ -303,7 +347,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
         )