test_wrap_tool_call_modify_args() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_tool_call.py lines 115–149

def test_wrap_tool_call_modify_args() -> None:
    """Test modifying tool arguments with wrap_tool_call decorator."""

    @wrap_tool_call
    def modify_args(
        request: ToolCallRequest, handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]]
    ) -> ToolMessage | Command[Any]:
        assert isinstance(request.tool, BaseTool)
        # Modify the query argument before execution
        if request.tool.name == "search":
            request.tool_call["args"]["query"] = "modified query"
        return handler(request)

    model = FakeToolCallingModel(
        tool_calls=[
            [ToolCall(name="search", args={"query": "original"}, id="1")],
            [],
        ]
    )

    agent = create_agent(
        model=model,
        tools=[search],
        middleware=[modify_args],
        checkpointer=InMemorySaver(),
    )

    result = agent.invoke(
        {"messages": [HumanMessage("Search")]},
        {"configurable": {"thread_id": "test"}},
    )

    tool_messages = [m for m in result["messages"] if isinstance(m, ToolMessage)]
    assert len(tool_messages) == 1
    assert "modified query" in tool_messages[0].content