TestLLMToolSelectorBasic Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/implementations/test_tool_selection.py lines 97–207

class TestLLMToolSelectorBasic:
    """Test basic tool selection functionality."""

    def test_sync_basic_selection(self) -> None:
        """Test synchronous tool selection."""
        # First call: selector picks tools
        # Second call: agent uses selected tools

        model_requests = []

        @wrap_model_call
        def trace_model_requests(
            request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            """Middleware to select relevant tools based on state/context."""
            # Select a small, relevant subset of tools based on state/context
            model_requests.append(request)
            return handler(request)

        tool_selection_model = FakeModel(
            messages=cycle(
                [
                    AIMessage(
                        content="",
                        tool_calls=[
                            {
                                "name": "ToolSelectionResponse",
                                "id": "1",
                                "args": {"tools": ["get_weather", "calculate"]},
                            }
                        ],
                    ),
                ]
            )
        )

        model = FakeModel(
            messages=iter(
                [
                    AIMessage(
                        content="",
                        tool_calls=[
                            {"name": "get_weather", "id": "2", "args": {"location": "Paris"}}
                        ],
                    ),
                    AIMessage(content="The weather in Paris is 72°F and sunny."),
                ]
            )
        )

        tool_selector = LLMToolSelectorMiddleware(max_tools=2, model=tool_selection_model)

        agent = create_agent(
            model=model,
            tools=[get_weather, search_web, calculate, send_email, get_stock_price],
            middleware=[tool_selector, trace_model_requests],
        )

        response = agent.invoke({"messages": [HumanMessage("What's the weather in Paris?")]})

        assert isinstance(response["messages"][-1], AIMessage)

        for request in model_requests:
            selected_tool_names = []
            for tool_ in request.tools:
                assert isinstance(tool_, BaseTool)
                selected_tool_names.append(tool_.name)
            assert selected_tool_names == ["get_weather", "calculate"]

    async def test_async_basic_selection(self) -> None:
        """Test asynchronous tool selection."""
        tool_selection_model = FakeModel(
            messages=cycle(
                [
                    AIMessage(
                        content="",
                        tool_calls=[
                            {
                                "name": "ToolSelectionResponse",
                                "id": "1",
                                "args": {"tools": ["search_web"]},