codekingpro/portable-devtools

1"""Tool emulator middleware for testing."""2 3from __future__ import annotations4 5from typing import TYPE_CHECKING, Any, Generic6 7from langchain_core.language_models.chat_models import BaseChatModel8from langchain_core.messages import HumanMessage, ToolMessage9 10from langchain.agents.middleware.types import AgentMiddleware, AgentState, ContextT11from langchain.chat_models.base import init_chat_model12 13if TYPE_CHECKING:14    from collections.abc import Awaitable, Callable15 16    from langgraph.types import Command17 18    from langchain.agents.middleware.types import ToolCallRequest19    from langchain.tools import BaseTool20 21 22class LLMToolEmulator(AgentMiddleware[AgentState[Any], ContextT], Generic[ContextT]):23    """Emulates specified tools using an LLM instead of executing them.24 25    This middleware allows selective emulation of tools for testing purposes.26 27    By default (when `tools=None`), all tools are emulated. You can specify which28    tools to emulate by passing a list of tool names or `BaseTool` instances.29 30    Examples:31        !!! example "Emulate all tools (default behavior)"32 33            ```python34            from langchain.agents.middleware import LLMToolEmulator35 36            middleware = LLMToolEmulator()37 38            agent = create_agent(39                model="openai:gpt-4o",40                tools=[get_weather, get_user_location, calculator],41                middleware=[middleware],42            )43            ```44 45        !!! example "Emulate specific tools by name"46 47            ```python48            middleware = LLMToolEmulator(tools=["get_weather", "get_user_location"])49            ```50 51        !!! example "Use a custom model for emulation"52 53            ```python54            middleware = LLMToolEmulator(55                tools=["get_weather"], model="anthropic:claude-sonnet-4-5-20250929"56            )57            ```58 59        !!! example "Emulate specific tools by passing tool instances"60 61            ```python62            middleware = LLMToolEmulator(tools=[get_weather, get_user_location])63            ```64    """65 66    def __init__(67        self,68        *,69        tools: list[str | BaseTool] | None = None,70        model: str | BaseChatModel | None = None,71    ) -> None:72        """Initialize the tool emulator.73 74        Args:75            tools: List of tool names (`str`) or `BaseTool` instances to emulate.76 77                If `None`, ALL tools will be emulated.78 79                If empty list, no tools will be emulated.80            model: Model to use for emulation.81 82                Defaults to `'anthropic:claude-sonnet-4-5-20250929'`.83 84                Can be a model identifier string or `BaseChatModel` instance.85        """86        super().__init__()87 88        # Extract tool names from tools89        # None means emulate all tools90        self.emulate_all = tools is None91        self.tools_to_emulate: set[str] = set()92 93        if not self.emulate_all and tools is not None:94            for tool in tools:95                if isinstance(tool, str):96                    self.tools_to_emulate.add(tool)97                else:98                    # Assume BaseTool with .name attribute99                    self.tools_to_emulate.add(tool.name)100 101        # Initialize emulator model102        if model is None:103            self.model = init_chat_model("anthropic:claude-sonnet-4-5-20250929", temperature=1)104        elif isinstance(model, BaseChatModel):105            self.model = model106        else:107            self.model = init_chat_model(model, temperature=1)108 109    def wrap_tool_call(110        self,111        request: ToolCallRequest,112        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],113    ) -> ToolMessage | Command[Any]:114        """Emulate tool execution using LLM if tool should be emulated.115 116        Args:117            request: Tool call request to potentially emulate.118            handler: Callback to execute the tool (can be called multiple times).119 120        Returns:121            ToolMessage with emulated response if tool should be emulated,122                otherwise calls handler for normal execution.123        """124        tool_name = request.tool_call["name"]125 126        # Check if this tool should be emulated127        should_emulate = self.emulate_all or tool_name in self.tools_to_emulate128 129        if not should_emulate:130            # Let it execute normally by calling the handler131            return handler(request)132 133        # Extract tool information for emulation134        tool_args = request.tool_call["args"]135        tool_description = request.tool.description if request.tool else "No description available"136 137        # Build prompt for emulator LLM138        prompt = (139            f"You are emulating a tool call for testing purposes.\n\n"140            f"Tool: {tool_name}\n"141            f"Description: {tool_description}\n"142            f"Arguments: {tool_args}\n\n"143            f"Generate a realistic response that this tool would return "144            f"given these arguments.\n"145            f"Return ONLY the tool's output, no explanation or preamble. "146            f"Introduce variation into your responses."147        )148 149        # Get emulated response from LLM150        response = self.model.invoke([HumanMessage(prompt)])151 152        # Short-circuit: return emulated result without executing real tool153        return ToolMessage(154            content=response.content,155            tool_call_id=request.tool_call["id"],156            name=tool_name,157        )158 159    async def awrap_tool_call(160        self,161        request: ToolCallRequest,162        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],163    ) -> ToolMessage | Command[Any]:164        """Async version of `wrap_tool_call`.165 166        Emulate tool execution using LLM if tool should be emulated.167 168        Args:169            request: Tool call request to potentially emulate.170            handler: Async callback to execute the tool (can be called multiple times).171 172        Returns:173            ToolMessage with emulated response if tool should be emulated,174                otherwise calls handler for normal execution.175        """176        tool_name = request.tool_call["name"]177 178        # Check if this tool should be emulated179        should_emulate = self.emulate_all or tool_name in self.tools_to_emulate180 181        if not should_emulate:182            # Let it execute normally by calling the handler183            return await handler(request)184 185        # Extract tool information for emulation186        tool_args = request.tool_call["args"]187        tool_description = request.tool.description if request.tool else "No description available"188 189        # Build prompt for emulator LLM190        prompt = (191            f"You are emulating a tool call for testing purposes.\n\n"192            f"Tool: {tool_name}\n"193            f"Description: {tool_description}\n"194            f"Arguments: {tool_args}\n\n"195            f"Generate a realistic response that this tool would return "196            f"given these arguments.\n"197            f"Return ONLY the tool's output, no explanation or preamble. "198            f"Introduce variation into your responses."199        )200 201        # Get emulated response from LLM (using async invoke)202        response = await self.model.ainvoke([HumanMessage(prompt)])203 204        # Short-circuit: return emulated result without executing real tool205        return ToolMessage(206            content=response.content,207            tool_call_id=request.tool_call["id"],208            name=tool_name,209        )210