coco-framework 0.1.0__py3-none-any.whl

coco/agents/middleware/tool_emulator.py

@@ -0,0 +1,251 @@
+"""Tool emulator middleware for testing."""
+
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING, Any, Generic
+
+from coco_core.language_models.chat_models import BaseChatModel
+from coco_core.messages import HumanMessage, ToolMessage
+
+from coco.agents.middleware.internal_call_transformer import (
+    InternalCallTransformer,
+    internal_call_metadata,
+)
+from coco.agents.middleware.types import AgentMiddleware, AgentState, ContextT
+from coco.chat_models.base import init_chat_model
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from coco_graph.types import Command
+
+    from coco.agents.middleware.types import ToolCallRequest
+    from coco.tools import BaseTool
+
+_DEFAULT_EMULATOR_MODEL = "anthropic:claude-sonnet-4-5-20250929"
+
+
+class LLMToolEmulator(AgentMiddleware[AgentState[Any], ContextT], Generic[ContextT]):
+    """Emulates specified tools using an LLM instead of executing them.
+
+    This middleware allows selective emulation of tools for testing purposes.
+
+    By default (when `tools=None`), all tools are emulated. You can specify which
+    tools to emulate by passing a list of tool names or `BaseTool` instances.
+
+    Examples:
+        !!! example "Emulate all tools (default behavior)"
+
+            ```python
+            from coco.agents.middleware import LLMToolEmulator
+
+            middleware = LLMToolEmulator()
+
+            agent = create_agent(
+                model="openai:gpt-5.5",
+                tools=[get_weather, get_user_location, calculator],
+                middleware=[middleware],
+            )
+            ```
+
+        !!! example "Emulate specific tools by name"
+
+            ```python
+            middleware = LLMToolEmulator(tools=["get_weather", "get_user_location"])
+            ```
+
+        !!! example "Use a custom model for emulation"
+
+            ```python
+            middleware = LLMToolEmulator(
+                tools=["get_weather"], model="anthropic:claude-sonnet-4-5-20250929"
+            )
+            ```
+
+        !!! example "Emulate specific tools by passing tool instances"
+
+            ```python
+            middleware = LLMToolEmulator(tools=[get_weather, get_user_location])
+            ```
+    """
+
+    transformers = (InternalCallTransformer,)
+    """Keeps the emulation model call's tokens out of `run.messages`.
+
+    Registered only when this middleware is used — see
+    `InternalCallTransformer` for why the call needs tagging and filtering.
+    """
+
+    def __init__(
+        self,
+        *,
+        tools: list[str | BaseTool] | None = None,
+        model: str | BaseChatModel | None = None,
+    ) -> None:
+        """Initialize the tool emulator.
+
+        Args:
+            tools: List of tool names (`str`) or `BaseTool` instances to emulate.
+
+                If `None`, ALL tools will be emulated.
+
+                If empty list, no tools will be emulated.
+            model: Model to use for emulation.
+
+                Defaults to `'anthropic:claude-sonnet-4-5-20250929'`, which requires
+                `trustcoco-anthropic` to be installed.
+
+                !!! warning "Deprecated"
+                    Relying on the implicit default is deprecated and will be
+                    removed in a future release, since it makes this middleware
+                    depend on `trustcoco-anthropic` even when unspecified. Pass
+                    `model` explicitly instead.
+
+                Can be a model identifier string or `BaseChatModel` instance.
+        """
+        super().__init__()
+
+        # Extract tool names from tools
+        # None means emulate all tools
+        self.emulate_all = tools is None
+        self.tools_to_emulate: set[str] = set()
+
+        if tools is not None:
+            for tool in tools:
+                if isinstance(tool, str):
+                    self.tools_to_emulate.add(tool)
+                else:
+                    # Assume BaseTool with .name attribute
+                    self.tools_to_emulate.add(tool.name)
+
+        # Initialize emulator model
+        if model is None:
+            warnings.warn(
+                "LLMToolEmulator's default model "
+                f"({_DEFAULT_EMULATOR_MODEL!r}) is deprecated and will be removed "
+                "in a future release. Pass `model` explicitly instead.",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+            try:
+                self.model = init_chat_model(_DEFAULT_EMULATOR_MODEL, temperature=1)
+            except ImportError as e:
+                msg = (
+                    "LLMToolEmulator's default model requires `trustcoco-anthropic` "
+                    "to be installed. Install it with `pip install trustcoco-anthropic`, "
+                    "or pass `model=...` explicitly to use a different provider."
+                )
+                raise ImportError(msg) from e
+        elif isinstance(model, BaseChatModel):
+            self.model = model
+        else:
+            self.model = init_chat_model(model, temperature=1)
+
+    def wrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
+    ) -> ToolMessage | Command[Any]:
+        """Emulate tool execution using LLM if tool should be emulated.
+
+        Args:
+            request: Tool call request to potentially emulate.
+            handler: Callback to execute the tool (can be called multiple times).
+
+        Returns:
+            ToolMessage with emulated response if tool should be emulated,
+                otherwise calls handler for normal execution.
+        """
+        tool_name = request.tool_call["name"]
+
+        # Check if this tool should be emulated
+        should_emulate = self.emulate_all or tool_name in self.tools_to_emulate
+
+        if not should_emulate:
+            # Let it execute normally by calling the handler
+            return handler(request)
+
+        # Extract tool information for emulation
+        tool_args = request.tool_call["args"]
+        tool_description = request.tool.description if request.tool else "No description available"
+
+        # Build prompt for emulator LLM
+        prompt = (
+            f"You are emulating a tool call for testing purposes.\n\n"
+            f"Tool: {tool_name}\n"
+            f"Description: {tool_description}\n"
+            f"Arguments: {tool_args}\n\n"
+            f"Generate a realistic response that this tool would return "
+            f"given these arguments.\n"
+            f"Return ONLY the tool's output, no explanation or preamble. "
+            f"Introduce variation into your responses."
+        )
+
+        # Get emulated response from LLM
+        response = self.model.invoke(
+            [HumanMessage(prompt)],
+            config={"metadata": {"lc_source": "tool_emulation", **internal_call_metadata()}},
+        )
+
+        # Short-circuit: return emulated result without executing real tool
+        return ToolMessage(
+            content=response.content,
+            tool_call_id=request.tool_call["id"],
+            name=tool_name,
+        )
+
+    async def awrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],
+    ) -> ToolMessage | Command[Any]:
+        """Async version of `wrap_tool_call`.
+
+        Emulate tool execution using LLM if tool should be emulated.
+
+        Args:
+            request: Tool call request to potentially emulate.
+            handler: Async callback to execute the tool (can be called multiple times).
+
+        Returns:
+            ToolMessage with emulated response if tool should be emulated,
+                otherwise calls handler for normal execution.
+        """
+        tool_name = request.tool_call["name"]
+
+        # Check if this tool should be emulated
+        should_emulate = self.emulate_all or tool_name in self.tools_to_emulate
+
+        if not should_emulate:
+            # Let it execute normally by calling the handler
+            return await handler(request)
+
+        # Extract tool information for emulation
+        tool_args = request.tool_call["args"]
+        tool_description = request.tool.description if request.tool else "No description available"
+
+        # Build prompt for emulator LLM
+        prompt = (
+            f"You are emulating a tool call for testing purposes.\n\n"
+            f"Tool: {tool_name}\n"
+            f"Description: {tool_description}\n"
+            f"Arguments: {tool_args}\n\n"
+            f"Generate a realistic response that this tool would return "
+            f"given these arguments.\n"
+            f"Return ONLY the tool's output, no explanation or preamble. "
+            f"Introduce variation into your responses."
+        )
+
+        # Get emulated response from LLM (using async invoke)
+        response = await self.model.ainvoke(
+            [HumanMessage(prompt)],
+            config={"metadata": {"lc_source": "tool_emulation", **internal_call_metadata()}},
+        )
+
+        # Short-circuit: return emulated result without executing real tool
+        return ToolMessage(
+            content=response.content,
+            tool_call_id=request.tool_call["id"],
+            name=tool_name,
+        )

coco/agents/middleware/tool_error.py

@@ -0,0 +1,209 @@
+"""Tool error middleware for agents."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, cast
+
+from coco_core.messages import ToolMessage
+from coco_graph.errors import GraphBubbleUp
+
+from coco.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    ContextT,
+    ResponseT,
+    TracePolicy,
+    omit_payload,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from coco_core.messages import ContentBlock
+    from coco_graph.types import Command
+
+    from coco.agents.middleware.types import ToolCallRequest
+    from coco.tools import BaseTool
+
+    OnError = Callable[[Exception, ToolCallRequest], str | list[ContentBlock] | None]
+    """Sync handler: return content to surface the error as a `ToolMessage`; return
+    `None` (or nothing) to let the exception propagate."""
+
+    AOnError = Callable[[Exception, ToolCallRequest], Awaitable[str | list[ContentBlock] | None]]
+    """Async handler: return content to surface the error as a `ToolMessage`; return
+    `None` (or nothing) to let the exception propagate."""
+
+
+class ToolErrorMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]):
+    """Return selected tool-execution exceptions to the model as error `ToolMessage`s.
+
+    `on_error` is called for each exception raised by tool execution. Return content
+    (a `str` or a list of content blocks) to convert the exception into a
+    `ToolMessage(status="error")`; return `None` — or simply don't return — to let the
+    exception propagate (halting the run). Handling is therefore opt-in — exceptions you
+    do not return content for propagate unchanged, so arbitrary internal exceptions are
+    never serialized to the model or end user unless you choose to surface them.
+
+    CocoGraph control-flow signals (interrupts, parent commands) always propagate and
+    never reach `on_error`.
+
+    Prefer returning content that names the exception type over the raw exception message,
+    which may carry sensitive or internal detail.
+
+    Provide at least one of `on_error` or `aon_error`. `aon_error` handles errors on the
+    async execution path (falling back to `on_error` when omitted); the sync path only
+    ever calls `on_error`. For async-only usage, pass `aon_error` alone — running such a
+    middleware on the sync path raises, since the async handler cannot be awaited there.
+
+    This middleware does not retry. For retries, compose with `ToolRetryMiddleware`
+    placed *inner* and configured with `on_failure="error"` so exceptions reach this
+    middleware.
+
+    This middleware only sees exceptions raised by tool *execution*. Argument-binding
+    and validation errors are handled upstream by `ToolNode` (converted to an error
+    `ToolMessage` before the tool runs), so they do not reach `on_error`.
+
+    Example:
+        ```python
+        from coco.agents import create_agent
+        from coco.agents.middleware import ToolErrorMiddleware
+
+
+        def on_error(exc: Exception, request: ToolCallRequest) -> str | None:
+            if isinstance(exc, ValueError):
+                return f"`{request.tool_call['name']}` failed; fix the input and retry."
+            return None  # propagate everything else
+
+
+        agent = create_agent(model, tools=[...], middleware=[ToolErrorMiddleware(on_error)])
+        ```
+    """
+
+    trace_policy = TracePolicy(process_inputs=omit_payload)
+
+    def __init__(
+        self,
+        on_error: OnError | None = None,
+        *,
+        aon_error: AOnError | None = None,
+        tools: list[BaseTool | str] | None = None,
+    ) -> None:
+        """Initialize `ToolErrorMiddleware`.
+
+        Args:
+            on_error: Handler called for each exception raised by tool execution. Return
+                content (`str` or list of content blocks) to convert the exception into an
+                error `ToolMessage`. Return `None` — or simply don't return — to let the
+                exception propagate. Falling through without a return therefore re-raises,
+                so handle only the exceptions you mean to. Receives the exception and the
+                tool call request (tool name, args, call id). Used on the sync path and,
+                unless `aon_error` is given, on the async path.
+            aon_error: Optional async handler, used on the async execution path. Falls back
+                to `on_error` when not provided.
+            tools: Optional list of tools or tool names to apply handling to. If `None`,
+                applies to all tools.
+
+        Raises:
+            ValueError: If neither `on_error` nor `aon_error` is provided.
+        """
+        super().__init__()
+
+        if on_error is None and aon_error is None:
+            msg = "ToolErrorMiddleware requires `on_error` and/or `aon_error`."
+            raise ValueError(msg)
+
+        self.on_error = on_error
+        self.aon_error = aon_error
+
+        # Extract tool names from BaseTool instances or strings
+        self._tool_filter: list[str] | None
+        if tools is not None:
+            self._tool_filter = [tool.name if not isinstance(tool, str) else tool for tool in tools]
+        else:
+            self._tool_filter = None
+
+        self.tools = []  # No additional tools registered by this middleware
+
+    def _should_handle_tool(self, tool_name: str) -> bool:
+        """Check if error handling should apply to this tool."""
+        if self._tool_filter is None:
+            return True
+        return tool_name in self._tool_filter
+
+    def wrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], ToolMessage | Command[Any]],
+    ) -> ToolMessage | Command[Any]:
+        """Intercept tool execution and convert handled exceptions to error messages.
+
+        Args:
+            request: Tool call request with call dict, `BaseTool`, state, and runtime.
+            handler: Callable to execute the tool.
+
+        Returns:
+            `ToolMessage` or `Command` (the final result).
+        """
+        tool_name = request.tool.name if request.tool else request.tool_call["name"]
+
+        if not self._should_handle_tool(tool_name):
+            return handler(request)
+
+        try:
+            return handler(request)
+        except GraphBubbleUp:
+            # Control-flow signals (interrupts, parent commands) must propagate.
+            raise
+        except Exception as exc:
+            if self.on_error is None:
+                # Async-only config (aon_error) cannot be awaited on the sync path.
+                msg = (
+                    "ToolErrorMiddleware has no sync `on_error`; run async "
+                    "(ainvoke/astream) or provide `on_error`."
+                )
+                raise RuntimeError(msg) from exc
+            content = self.on_error(exc, request)
+            if content is None:
+                raise
+            return ToolMessage(
+                content=cast("str | list[str | dict[Any, Any]]", content),
+                tool_call_id=request.tool_call["id"],
+                name=tool_name,
+                status="error",
+            )
+
+    async def awrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command[Any]]],
+    ) -> ToolMessage | Command[Any]:
+        """Async version of `wrap_tool_call`.
+
+        Uses `aon_error` if provided, otherwise the sync `on_error`. The sync path never
+        awaits.
+        """
+        tool_name = request.tool.name if request.tool else request.tool_call["name"]
+
+        if not self._should_handle_tool(tool_name):
+            return await handler(request)
+
+        try:
+            return await handler(request)
+        except GraphBubbleUp:
+            # Control-flow signals (interrupts, parent commands) must propagate.
+            raise
+        except Exception as exc:
+            if self.aon_error is not None:
+                content = await self.aon_error(exc, request)
+            elif self.on_error is not None:
+                content = self.on_error(exc, request)
+            else:  # pragma: no cover - __init__ guarantees at least one handler
+                raise
+            if content is None:
+                raise
+            return ToolMessage(
+                content=cast("str | list[str | dict[Any, Any]]", content),
+                tool_call_id=request.tool_call["id"],
+                name=tool_name,
+                status="error",
+            )