coco-framework 0.1.0__py3-none-any.whl

coco/agents/middleware/model_fallback.py

@@ -0,0 +1,417 @@
+"""Model fallback middleware for agents.
+
+When a caching middleware such as `AnthropicPromptCachingMiddleware` wraps this
+middleware from the outside, it applies Anthropic `cache_control` markers to the
+request *before* the fallback loop runs. Those markers are provider-specific and
+cause API errors on non-Anthropic fallback models, so this middleware strips them
+from fallback attempts — but only when the fallback model itself cannot accept
+Anthropic cache markers. When the fallback is another Anthropic model the markers
+are valid and preserve prompt caching, so they are left intact.
+
+The knowledge of the `cache_control` marker is duplicated here (rather than owned
+solely by the Anthropic partner package) because an outer caching middleware
+never re-runs during fallback and therefore cannot clean up after itself.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from coco_core.tools import BaseTool
+from coco_graph.errors import GraphBubbleUp
+
+from coco.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    ContextT,
+    ModelRequest,
+    ModelResponse,
+    ResponseT,
+)
+from coco.chat_models import init_chat_model
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from coco_core.language_models.chat_models import BaseChatModel
+    from coco_core.messages import AIMessage, AnyMessage, SystemMessage
+
+logger = logging.getLogger(__name__)
+
+
+def _sanitize_content_blocks(
+    content: str | list[str | dict[str, Any]],
+) -> str | list[str | dict[str, Any]]:
+    """Remove Anthropic cache markers from message content blocks."""
+    if not isinstance(content, list):
+        return content
+
+    sanitized_content: list[str | dict[str, Any]] = []
+    changed = False
+
+    for block in content:
+        if not isinstance(block, dict):
+            sanitized_content.append(block)
+            continue
+
+        sanitized_block, block_changed = _without_cache_control_from_content_block(block)
+        changed = changed or block_changed
+        sanitized_content.append(sanitized_block)
+
+    return sanitized_content if changed else content
+
+
+def _sanitize_system_message(
+    system_message: SystemMessage | None,
+) -> SystemMessage | None:
+    """Remove Anthropic cache markers from a system message."""
+    if system_message is None:
+        return None
+
+    sanitized_content = _sanitize_content_blocks(system_message.content)
+    if sanitized_content is system_message.content:
+        return system_message
+
+    return system_message.model_copy(update={"content": sanitized_content})
+
+
+def _sanitize_messages(messages: list[AnyMessage]) -> list[AnyMessage]:
+    """Remove Anthropic cache markers from request messages."""
+    sanitized_messages: list[AnyMessage] = []
+    changed = False
+
+    for message in messages:
+        sanitized_message, message_changed = _sanitize_message(message)
+        changed = changed or message_changed
+        sanitized_messages.append(sanitized_message)
+
+    return sanitized_messages if changed else messages
+
+
+def _sanitize_tools(
+    tools: list[BaseTool | dict[str, Any]],
+) -> list[BaseTool | dict[str, Any]]:
+    """Remove Anthropic cache markers from tool payloads."""
+    sanitized_tools: list[BaseTool | dict[str, Any]] = []
+    changed = False
+
+    for tool in tools:
+        sanitized_tool: BaseTool | dict[str, Any]
+        if isinstance(tool, BaseTool):
+            sanitized_tool, tool_changed = _sanitize_base_tool(tool)
+        else:
+            sanitized_tool, tool_changed = _sanitize_dict_tool(tool)
+
+        changed = changed or tool_changed
+        sanitized_tools.append(sanitized_tool)
+
+    return sanitized_tools if changed else tools
+
+
+def _sanitize_request_for_fallback(request: ModelRequest[ContextT]) -> ModelRequest[ContextT]:
+    """Sanitize provider-specific Anthropic cache markers before fallback attempts."""
+    overrides: dict[str, Any] = {}
+
+    model_settings, model_settings_changed = _without_cache_control(request.model_settings)
+    if model_settings_changed:
+        overrides["model_settings"] = model_settings
+
+    system_message = _sanitize_system_message(request.system_message)
+    if system_message is not request.system_message:
+        overrides["system_message"] = system_message
+
+    messages = _sanitize_messages(request.messages)
+    if messages is not request.messages:
+        overrides["messages"] = messages
+
+    tools = _sanitize_tools(request.tools)
+    if tools is not request.tools:
+        overrides["tools"] = tools
+
+    if not overrides:
+        return request
+
+    # Log only the field names that changed, never request content (may contain
+    # prompt data or PII).
+    logger.debug(
+        "Stripped Anthropic cache_control markers from %s before fallback attempt",
+        sorted(overrides),
+    )
+
+    return request.override(**overrides)
+
+
+def _sanitize_message(message: AnyMessage) -> tuple[AnyMessage, bool]:
+    """Remove Anthropic cache markers from a single message.
+
+    Returns:
+        The sanitized message (the original instance when unchanged) and whether
+            any marker was removed.
+    """
+    sanitized_content = _sanitize_content_blocks(message.content)
+    if sanitized_content is message.content:
+        return message, False
+
+    return message.model_copy(update={"content": sanitized_content}), True
+
+
+def _sanitize_base_tool(tool: BaseTool) -> tuple[BaseTool, bool]:
+    """Remove Anthropic cache markers from a `BaseTool` payload.
+
+    Returns:
+        The sanitized tool (the original instance when unchanged) and whether any
+            marker was removed.
+
+            Emptied `extras` collapse back to `None`.
+    """
+    if not tool.extras:
+        return tool, False
+
+    sanitized_extras, changed = _without_cache_control(tool.extras)
+    if not changed:
+        return tool, False
+
+    return tool.model_copy(update={"extras": sanitized_extras or None}), True
+
+
+def _sanitize_dict_tool(tool: dict[str, Any]) -> tuple[dict[str, Any], bool]:
+    """Remove Anthropic cache markers from a dict-style tool payload.
+
+    Returns:
+        The sanitized tool (the original instance when unchanged) and whether any
+            marker was removed.
+
+            Emptied `extras` collapse back to `None`.
+    """
+    sanitized_tool, changed = _without_cache_control(tool)
+
+    extras = sanitized_tool.get("extras")
+    if not isinstance(extras, dict):
+        return sanitized_tool, changed
+
+    sanitized_extras, extras_changed = _without_cache_control(extras)
+    if not extras_changed:
+        return sanitized_tool, changed
+
+    return {**sanitized_tool, "extras": sanitized_extras or None}, True
+
+
+def _without_cache_control(payload: dict[str, Any]) -> tuple[dict[str, Any], bool]:
+    """Return payload without `cache_control`, plus whether anything changed."""
+    if "cache_control" not in payload:
+        return payload, False
+
+    return (
+        {key: value for key, value in payload.items() if key != "cache_control"},
+        True,
+    )
+
+
+def _without_cache_control_from_content_block(
+    block: dict[str, Any],
+) -> tuple[dict[str, Any], bool]:
+    """Return content block without Anthropic cache markers.
+
+    Strips `cache_control` from the block itself and from its nested `extras` and
+    `metadata` payloads.
+
+    Returns:
+        The sanitized block (the original instance when unchanged) and whether any
+            marker was removed.
+    """
+    sanitized_block, changed = _without_cache_control(block)
+
+    for nested_key in ("extras", "metadata"):
+        nested_payload = sanitized_block.get(nested_key)
+        if not isinstance(nested_payload, dict):
+            continue
+
+        sanitized_payload, nested_changed = _without_cache_control(nested_payload)
+        if not nested_changed:
+            continue
+
+        if sanitized_block is block:
+            sanitized_block = dict(block)
+        sanitized_block[nested_key] = sanitized_payload
+        changed = True
+
+    return sanitized_block, changed
+
+
+# `_llm_type` values that indicate a model speaks an Anthropic-compatible API
+# and therefore accepts `cache_control` markers. Direct Anthropic models
+# (`ChatAnthropic`) report `"anthropic-chat"`; Bedrock-hosted Claude
+# (`ChatAnthropicBedrock`, a `ChatAnthropic` subclass in `trustcoco-aws`) reports
+# `"anthropic-bedrock-chat"` and translates the top-level kwarg into block-level
+# breakpoints inside the inherited `ChatAnthropic._get_request_payload`, while
+# content-block and tool `cache_control` markers pass through unchanged.
+# Vertex-hosted Claude (`ChatAnthropicVertex` in `trustcoco-google`) reports
+# `"anthropic-chat-vertexai"` and nests the same marker shape through its own
+# request builder — not the shared `ChatAnthropic` method. All three keep prompt
+# caching intact on fallback.
+#
+# Keep this set in sync with those classes' `_llm_type` values, which live in
+# separate repositories. If a value drifts or a new Anthropic transport ships,
+# the failure mode is silent loss of prompt caching (markers stripped from a
+# model that supports them), not a hard error — so CI here will not catch it.
+_ANTHROPIC_LLM_TYPES: frozenset[str] = frozenset(
+    {
+        "anthropic-chat",
+        "anthropic-bedrock-chat",
+        "anthropic-chat-vertexai",
+    }
+)
+
+
+def _supports_anthropic_cache_control(model: BaseChatModel) -> bool:
+    """Return whether `model` accepts Anthropic `cache_control` markers.
+
+    Checked via `_llm_type` so the decision is provider-based rather than
+    model-name-based: any Anthropic-compatible model (including future model IDs
+    we have not seen) keeps its cache markers on fallback, while OpenAI, Gemini,
+    and other non-Anthropic providers get a sanitized request.
+    """
+    llm_type = getattr(model, "_llm_type", None)
+    return isinstance(llm_type, str) and llm_type in _ANTHROPIC_LLM_TYPES
+
+
+class ModelFallbackMiddleware(AgentMiddleware[AgentState[ResponseT], ContextT, ResponseT]):
+    """Automatic fallback to alternative models on errors.
+
+    Retries failed model calls with alternative models in sequence until
+    success or all models exhausted. Primary model specified in `create_agent`.
+
+    Example:
+        ```python
+        from coco.agents.middleware import ModelFallbackMiddleware
+        from coco.agents import create_agent
+
+        fallback = ModelFallbackMiddleware(
+            "openai:gpt-5.5",  # Try first on error
+            "anthropic:claude-sonnet-4-5-20250929",  # Then this
+        )
+
+        agent = create_agent(
+            model="openai:gpt-5.5",  # Primary model
+            middleware=[fallback],
+        )
+
+        # If primary fails: tries gpt-5.5, then claude-sonnet-4-5-20250929
+        result = await agent.invoke({"messages": [HumanMessage("Hello")]})
+        ```
+    """
+
+    def __init__(
+        self,
+        first_model: str | BaseChatModel,
+        *additional_models: str | BaseChatModel,
+    ) -> None:
+        """Initialize model fallback middleware.
+
+        Args:
+            first_model: First fallback model (string name or instance).
+            *additional_models: Additional fallbacks in order.
+        """
+        super().__init__()
+
+        # Initialize all fallback models
+        all_models = (first_model, *additional_models)
+        self.models: list[BaseChatModel] = []
+        for model in all_models:
+            if isinstance(model, str):
+                self.models.append(init_chat_model(model))
+            else:
+                self.models.append(model)
+
+    def wrap_model_call(
+        self,
+        request: ModelRequest[ContextT],
+        handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],
+    ) -> ModelResponse[ResponseT] | AIMessage:
+        """Try fallback models in sequence on errors.
+
+        Args:
+            request: Initial model request.
+            handler: Callback to execute the model.
+
+        Returns:
+            AIMessage from successful model call.
+
+        Raises:
+            Exception: If all models fail, re-raises last exception.
+        """
+        # Try primary model first
+        last_exception: Exception
+        try:
+            return handler(request)
+        except GraphBubbleUp:
+            raise
+        except Exception as e:
+            last_exception = e
+
+        # Try fallback models — sanitize cache markers only when the fallback
+        # model cannot accept them (i.e. is not an Anthropic-compatible model).
+        # The request is derived outside the try so a sanitizer or `_llm_type`
+        # bug surfaces directly instead of being masked as a model failure.
+        for fallback_model in self.models:
+            fallback_request = (
+                request
+                if _supports_anthropic_cache_control(fallback_model)
+                else _sanitize_request_for_fallback(request)
+            )
+            try:
+                return handler(fallback_request.override(model=fallback_model))
+            except GraphBubbleUp:
+                raise
+            except Exception as e:
+                last_exception = e
+                continue
+
+        raise last_exception
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest[ContextT],
+        handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],
+    ) -> ModelResponse[ResponseT] | AIMessage:
+        """Try fallback models in sequence on errors (async version).
+
+        Args:
+            request: Initial model request.
+            handler: Async callback to execute the model.
+
+        Returns:
+            AIMessage from successful model call.
+
+        Raises:
+            Exception: If all models fail, re-raises last exception.
+        """
+        # Try primary model first
+        last_exception: Exception
+        try:
+            return await handler(request)
+        except GraphBubbleUp:
+            raise
+        except Exception as e:
+            last_exception = e
+
+        # Try fallback models — sanitize cache markers only when the fallback
+        # model cannot accept them (i.e. is not an Anthropic-compatible model).
+        # The request is derived outside the try so a sanitizer or `_llm_type`
+        # bug surfaces directly instead of being masked as a model failure.
+        for fallback_model in self.models:
+            fallback_request = (
+                request
+                if _supports_anthropic_cache_control(fallback_model)
+                else _sanitize_request_for_fallback(request)
+            )
+            try:
+                return await handler(fallback_request.override(model=fallback_model))
+            except GraphBubbleUp:
+                raise
+            except Exception as e:
+                last_exception = e
+                continue
+
+        raise last_exception