feat(langchain_v1): add llm selection middleware (#33272)

* Add llm based tool selection middleware.
* Note that we might want some form of caching for when the agent is
inside an active tool calling loop as the tool selection isn't expected
to change during that time.

API:

```python
class LLMToolSelectorMiddleware(AgentMiddleware):
    """Uses an LLM to select relevant tools before calling the main model.

    When an agent has many tools available, this middleware filters them down
    to only the most relevant ones for the user's query. This reduces token usage
    and helps the main model focus on the right tools.

    Examples:
        Limit to 3 tools:
        ```python
        from langchain.agents.middleware import LLMToolSelectorMiddleware

        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)
        ```
    """

    def __init__(
        self,
        *,
        model: str | BaseChatModel | None = None,
        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
        max_tools: int | None = None,
        always_include: list[str] | None = None,
    ) -> None:
        """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.
            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.
            always_include: Tool names to always include regardless of selection.
                These do not count against the max_tools limit.
        """
```



```python
"""Test script for LLM tool selection middleware."""

from langchain.agents import create_agent
from langchain.agents.middleware import LLMToolSelectorMiddleware
from langchain_core.tools import tool


@tool
def get_weather(location: str) -> str:
    """Get current weather for a location."""
    return f"Weather in {location}: 72°F, sunny"


@tool
def search_web(query: str) -> str:
    """Search the web for information."""
    return f"Search results for: {query}"


@tool
def calculate(expression: str) -> str:
    """Perform mathematical calculations."""
    return f"Result of {expression}: 42"


@tool
def send_email(to: str, subject: str) -> str:
    """Send an email to someone."""
    return f"Email sent to {to} with subject: {subject}"


@tool
def get_stock_price(symbol: str) -> str:
    """Get current stock price for a symbol."""
    return f"Stock price for {symbol}: $150.25"


@tool
def translate_text(text: str, target_language: str) -> str:
    """Translate text to another language."""
    return f"Translated '{text}' to {target_language}"


@tool
def set_reminder(task: str, time: str) -> str:
    """Set a reminder for a task."""
    return f"Reminder set: {task} at {time}"


@tool
def get_news(topic: str) -> str:
    """Get latest news about a topic."""
    return f"Latest news about {topic}"


@tool
def book_flight(destination: str, date: str) -> str:
    """Book a flight to a destination."""
    return f"Flight booked to {destination} on {date}"


@tool
def get_restaurant_recommendations(city: str, cuisine: str) -> str:
    """Get restaurant recommendations."""
    return f"Top {cuisine} restaurants in {city}"


# Create agent with tool selection middleware
middleware = LLMToolSelectorMiddleware(
    model="openai:gpt-4o-mini",
    max_tools=3,
)

agent = create_agent(
    model="openai:gpt-4o",
    tools=[
        get_weather,
        search_web,
        calculate,
        send_email,
        get_stock_price,
        translate_text,
        set_reminder,
        get_news,
        book_flight,
        get_restaurant_recommendations,
    ],
    middleware=[middleware],
)

# Test with a query that should select specific tools
response = agent.invoke(
    {"messages": [{"role": "user", "content": "I need to find restaurants"}]}
)

print(response)
```

Author: eyurtsev

libs/langchain_v1/langchain/agents/middleware/__init__.py

@@ -6,6 +6,7 @@
 from .prompt_caching import AnthropicPromptCachingMiddleware
 from .summarization import SummarizationMiddleware
 from .tool_call_limit import ToolCallLimitMiddleware
+from .tool_selection import LLMToolSelectorMiddleware
 from .types import (
     AgentMiddleware,
     AgentState,
@@ -23,6 +24,7 @@
     # should move to langchain-anthropic if we decide to keep it
     "AnthropicPromptCachingMiddleware",
     "HumanInTheLoopMiddleware",
+    "LLMToolSelectorMiddleware",
     "ModelRequest",
     "PIIDetectionError",
     "PIIMiddleware",

libs/langchain_v1/langchain/agents/middleware/tool_selection.py

@@ -0,0 +1,293 @@
+"""LLM-based tool selector middleware."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Annotated, Literal, Union
+
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import HumanMessage
+from pydantic import Field, TypeAdapter
+from typing_extensions import TypedDict
+
+from langchain.agents.middleware.types import AgentMiddleware, AgentState, ModelRequest, StateT
+from langchain.chat_models.base import init_chat_model
+
+if TYPE_CHECKING:
+    from langgraph.runtime import Runtime
+    from langgraph.typing import ContextT
+
+    from langchain.tools import BaseTool
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_SYSTEM_PROMPT = (
+    "Your goal is to select the most relevant tools for answering the user's query."
+)
+
+
+@dataclass
+class _SelectionRequest:
+    """Prepared inputs for tool selection."""
+
+    available_tools: list[BaseTool]
+    system_message: str
+    last_user_message: HumanMessage
+    model: BaseChatModel
+    valid_tool_names: list[str]
+
+
+def _create_tool_selection_response(tools: list[BaseTool]) -> TypeAdapter:
+    """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.
+    """
+    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
+    literals = [
+        Annotated[Literal[tool.name], Field(description=tool.description)] for tool in tools
+    ]
+    selected_tool_type = Union[tuple(literals)]  # type: ignore[valid-type]  # noqa: UP007
+
+    description = "Tools to use. Place the most relevant tools first."
+
+    class ToolSelectionResponse(TypedDict):
+        """Use to select relevant tools."""
+
+        tools: Annotated[list[selected_tool_type], Field(description=description)]  # type: ignore[valid-type]
+
+    return TypeAdapter(ToolSelectionResponse)
+
+
+def _render_tool_list(tools: list[BaseTool]) -> str:
+    """Format tools as markdown list.
+
+    Args:
+        tools: Tools to format.
+
+    Returns:
+        Markdown string with each tool on a new line.
+    """
+    return "\n".join(f"- {tool.name}: {tool.description}" for tool in tools)
+
+
+class LLMToolSelectorMiddleware(AgentMiddleware):
+    """Uses an LLM to select relevant tools before calling the main model.
+
+    When an agent has many tools available, this middleware filters them down
+    to only the most relevant ones for the user's query. This reduces token usage
+    and helps the main model focus on the right tools.
+
+    Examples:
+        Limit to 3 tools:
+        ```python
+        from langchain.agents.middleware import LLMToolSelectorMiddleware
+
+        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)
+        ```
+    """
+
+    def __init__(
+        self,
+        *,
+        model: str | BaseChatModel | None = None,
+        system_prompt: str = DEFAULT_SYSTEM_PROMPT,
+        max_tools: int | None = None,
+        always_include: list[str] | None = None,
+    ) -> None:
+        """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.
+            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.
+            always_include: Tool names to always include regardless of selection.
+                These do not count against the max_tools limit.
+        """
+        super().__init__()
+        self.system_prompt = system_prompt
+        self.max_tools = max_tools
+        self.always_include = always_include or []
+
+        if isinstance(model, (BaseChatModel, type(None))):
+            self.model: BaseChatModel | None = model
+        else:
+            self.model = init_chat_model(model)
+
+    def _prepare_selection_request(self, request: ModelRequest) -> _SelectionRequest | None:
+        """Prepare inputs for tool selection.
+
+        Returns:
+            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:
+            return None
+
+        # Validate that always_include tools exist
+        if self.always_include:
+            available_tool_names = {tool.name for tool in request.tools}
+            missing_tools = [
+                name for name in self.always_include if name not in available_tool_names
+            ]
+            if missing_tools:
+                msg = (
+                    f"Tools in always_include not found in request: {missing_tools}. "
+                    f"Available tools: {sorted(available_tool_names)}"
+                )
+                raise ValueError(msg)
+
+        # Separate tools that are always included from those available for selection
+        available_tools = [tool for tool in request.tools if tool.name not in self.always_include]
+
+        # If no tools available for selection, return None
+        if not available_tools:
+            return None
+
+        system_message = self.system_prompt
+        # If there's a max_tools limit, append instructions to the system prompt
+        if self.max_tools is not None:
+            system_message += (
+                f"\nIMPORTANT: List the tool names in order of relevance, "
+                f"with the most relevant first. "
+                f"If you exceed the maximum number of tools, "
+                f"only the first {self.max_tools} will be used."
+            )
+
+        # Get the last user message from the conversation history
+        last_user_message: HumanMessage
+        for message in request.messages:
+            if isinstance(message, HumanMessage):
+                last_user_message = message
+                break
+        else:
+            msg = "No user message found in request messages"
+            raise AssertionError(msg)
+
+        model = self.model or request.model
+        valid_tool_names = [tool.name for tool in available_tools]
+
+        return _SelectionRequest(
+            available_tools=available_tools,
+            system_message=system_message,
+            last_user_message=last_user_message,
+            model=model,
+            valid_tool_names=valid_tool_names,
+        )
+
+    def _process_selection_response(
+        self,
+        response: dict,
+        available_tools: list[BaseTool],
+        valid_tool_names: list[str],
+        request: ModelRequest,
+    ) -> ModelRequest:
+        """Process the selection response and return filtered ModelRequest."""
+        selected_tool_names: list[str] = []
+        invalid_tool_selections = []
+
+        for tool_name in response["tools"]:
+            if tool_name not in valid_tool_names:
+                invalid_tool_selections.append(tool_name)
+                continue
+
+            # Only add if not already selected and within max_tools limit
+            if tool_name not in selected_tool_names and (
+                self.max_tools is None or len(selected_tool_names) < self.max_tools
+            ):
+                selected_tool_names.append(tool_name)
+
+        if invalid_tool_selections:
+            msg = f"Model selected invalid tools: {invalid_tool_selections}"
+            raise ValueError(msg)
+
+        # Filter tools based on selection and append always-included tools
+        selected_tools = [tool for tool in available_tools if tool.name in selected_tool_names]
+        always_included_tools = [tool for tool in request.tools if tool.name in self.always_include]
+        selected_tools.extend(always_included_tools)
+        request.tools = selected_tools
+        return request
+
+    def modify_model_request(
+        self,
+        request: ModelRequest,
+        state: StateT,  # noqa: ARG002
+        runtime: Runtime[ContextT],  # noqa: ARG002
+    ) -> ModelRequest:
+        """Modify the model request to filter tools based on LLM selection."""
+        selection_request = self._prepare_selection_request(request)
+        if selection_request is None:
+            return request
+
+        # Create dynamic response model with Literal enum of available tool names
+        type_adapter = _create_tool_selection_response(selection_request.available_tools)
+        schema = type_adapter.json_schema()
+        structured_model = selection_request.model.with_structured_output(schema)
+
+        response = structured_model.invoke(
+            [
+                {"role": "system", "content": selection_request.system_message},
+                selection_request.last_user_message,
+            ]
+        )
+
+        # 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)
+        return self._process_selection_response(
+            response, selection_request.available_tools, selection_request.valid_tool_names, request
+        )
+
+    async def amodify_model_request(
+        self,
+        request: ModelRequest,
+        state: AgentState,  # noqa: ARG002
+        runtime: Runtime,  # noqa: ARG002
+    ) -> ModelRequest:
+        """Modify the model request to filter tools based on LLM selection."""
+        selection_request = self._prepare_selection_request(request)
+        if selection_request is None:
+            return request
+
+        # Create dynamic response model with Literal enum of available tool names
+        type_adapter = _create_tool_selection_response(selection_request.available_tools)
+        schema = type_adapter.json_schema()
+        structured_model = selection_request.model.with_structured_output(schema)
+
+        response = await structured_model.ainvoke(
+            [
+                {"role": "system", "content": selection_request.system_message},
+                selection_request.last_user_message,
+            ]
+        )
+
+        # 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)
+        return self._process_selection_response(
+            response, selection_request.available_tools, selection_request.valid_tool_names, request
+        )