test_middleware_can_modify_tools() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_tools.py lines 73–123

def test_middleware_can_modify_tools() -> None:
    """Test that middleware can modify the list of tools in ModelRequest."""

    @tool
    def tool_a(value: str) -> str:
        """Tool A."""
        return "A"

    @tool
    def tool_b(value: str) -> str:
        """Tool B."""
        return "B"

    @tool
    def tool_c(value: str) -> str:
        """Tool C."""
        return "C"

    class ToolFilteringMiddleware(AgentMiddleware):
        def wrap_model_call(
            self,
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelCallResult:
            # Only allow tool_a and tool_b
            filtered_tools: list[BaseTool | dict[str, Any]] = []
            for t in request.tools:
                assert isinstance(t, BaseTool)
                if t.name in {"tool_a", "tool_b"}:
                    filtered_tools.append(t)
            return handler(request.override(tools=filtered_tools))

    # Model will try to call tool_a
    model = FakeToolCallingModel(
        tool_calls=[[{"args": {"input": "test"}, "id": "1", "name": "tool_a"}], []]
    )

    agent = create_agent(
        model=model,
        tools=[tool_a, tool_b, tool_c],
        system_prompt="You are a helpful assistant.",
        middleware=[ToolFilteringMiddleware()],
    )

    result = agent.invoke({"messages": [HumanMessage("Use tool_a")]})

    # Verify that the tool was executed successfully
    messages = result["messages"]
    tool_messages = [m for m in messages if isinstance(m, ToolMessage)]
    assert len(tool_messages) == 1
    assert tool_messages[0].name == "tool_a"