test_wrap_tool_call.py — langchain Source File

Source Code

"""Tests for wrap_tool_call decorator functionality.

These tests verify the decorator-based approach for wrapping tool calls,
focusing on the handler pattern (not generators).
"""

import time
from collections.abc import Callable
from typing import Any

from langchain_core.messages import HumanMessage, ToolCall, ToolMessage
from langchain_core.tools import BaseTool, tool
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command

from langchain.agents.factory import create_agent
from langchain.agents.middleware.types import ToolCallRequest, wrap_tool_call
from tests.unit_tests.agents.model import FakeToolCallingModel


@tool
def search(query: str) -> str:
    """Search for information."""
    return f"Results for: {query}"


@tool
def calculator(expression: str) -> str:
    """Calculate an expression."""
    return f"Calculated: {expression}"


@tool
def failing_tool(value: str) -> str:
    """Tool that always fails."""
    msg = f"Failed: {value}"
    raise ValueError(msg)


def test_wrap_tool_call_basic_passthrough() -> None:
    """Test basic passthrough with wrap_tool_call decorator."""
    call_log = []

    @wrap_tool_call
    def passthrough(
        request: ToolCallRequest, handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]]
    ) -> ToolMessage | Command[Any]:
        call_log.append("called")
        return handler(request)

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

    agent = create_agent(
        model=model,
        tools=[search],
// ... (808 more lines)