coco-framework 0.1.0__py3-none-any.whl

coco_anthropic/middleware/prompt_caching.py

@@ -0,0 +1,276 @@
+"""Anthropic prompt caching middleware.
+
+Requires:
+    - `coco`: For agent middleware framework
+    - `trustcoco-anthropic`: For `ChatAnthropic` model (already a dependency)
+"""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable
+from typing import Any, Literal
+from warnings import warn
+
+from coco_core.messages import SystemMessage
+from coco_core.tools import BaseTool
+
+from coco_anthropic.chat_models import ChatAnthropic
+
+try:
+    from coco.agents.middleware.types import (
+        AgentMiddleware,
+        ModelCallResult,
+        ModelRequest,
+        ModelResponse,
+    )
+except ImportError as e:
+    msg = (
+        "AnthropicPromptCachingMiddleware requires 'coco' to be installed. "
+        "This middleware is designed for use with Coco agents. "
+        "Install it with: pip install coco"
+    )
+    raise ImportError(msg) from e
+
+
+def _get_trace_policy() -> Any:
+    """Use tracing optimizations when supported by the installed Coco."""
+    try:
+        from coco.agents.middleware.types import TracePolicy, omit_payload
+    except ImportError:
+        return None
+    return TracePolicy(process_inputs=omit_payload)
+
+
+_TRACE_POLICY = _get_trace_policy()
+
+
+class AnthropicPromptCachingMiddleware(AgentMiddleware):
+    """Prompt Caching Middleware.
+
+    Optimizes API usage by caching conversation prefixes for Anthropic models.
+
+    Requires both `coco` and `trustcoco-anthropic` packages to be installed.
+
+    The middleware tags stable agent content and passes `cache_control` through
+    `model_settings`:
+
+    - **System message**: Tags the last content block of the system message
+        with `cache_control` so static system prompt content is cached.
+    - **Tools**: Tags the last tool definition with `cache_control`. Because
+        tool definitions are sent as one contiguous block, a single trailing
+        breakpoint caches the entire tool set across turns.
+    - **`model_settings`**: Passes `cache_control` to the chat model. The chat
+        model/provider then applies the correct message-tail and
+        provider-specific behavior at request time.
+
+    Learn more about Anthropic prompt caching
+    [here](https://platform.claude.com/docs/en/build-with-claude/prompt-caching).
+    """
+
+    trace_policy = _TRACE_POLICY
+
+    def __init__(
+        self,
+        type: Literal["ephemeral"] = "ephemeral",  # noqa: A002
+        ttl: Literal["5m", "1h"] = "5m",
+        min_messages_to_cache: int = 0,
+        unsupported_model_behavior: Literal["ignore", "warn", "raise"] = "warn",
+    ) -> None:
+        """Initialize the middleware with cache control settings.
+
+        Args:
+            type: The type of cache to use, only `'ephemeral'` is supported.
+            ttl: The time to live for the cache, only `'5m'` and `'1h'` are
+                supported.
+            min_messages_to_cache: The minimum number of messages until the
+                cache is used.
+            unsupported_model_behavior: The behavior to take when an
+                unsupported model is used.
+
+                `'ignore'` will ignore the unsupported model and continue without
+                caching.
+
+                `'warn'` will warn the user and continue without caching.
+
+                `'raise'` will raise an error and stop the agent.
+        """
+        self.type = type
+        self.ttl = ttl
+        self.min_messages_to_cache = min_messages_to_cache
+        self.unsupported_model_behavior = unsupported_model_behavior
+
+    @property
+    def _cache_control(self) -> dict[str, str]:
+        return {"type": self.type, "ttl": self.ttl}
+
+    def _should_apply_caching(self, request: ModelRequest) -> bool:
+        """Check if caching should be applied to the request.
+
+        Args:
+            request: The model request to check.
+
+        Returns:
+            `True` if caching should be applied, `False` otherwise.
+
+        Raises:
+            ValueError: If model is unsupported and behavior is set to `'raise'`.
+        """
+        if not isinstance(request.model, ChatAnthropic):
+            msg = (
+                "AnthropicPromptCachingMiddleware caching middleware only supports "
+                f"Anthropic models, not instances of {type(request.model)}"
+            )
+            if self.unsupported_model_behavior == "raise":
+                raise ValueError(msg)
+            if self.unsupported_model_behavior == "warn":
+                warn(msg, stacklevel=3)
+            return False
+
+        messages_count = (
+            len(request.messages) + 1
+            if request.system_message
+            else len(request.messages)
+        )
+        return messages_count >= self.min_messages_to_cache
+
+    def _apply_caching(self, request: ModelRequest) -> ModelRequest:
+        """Apply cache control to system message, tools, and model settings.
+
+        Args:
+            request: The model request to modify.
+
+        Returns:
+            New request with cache control applied.
+        """
+        overrides: dict[str, Any] = {}
+        cache_control = self._cache_control
+
+        # Always set top-level `cache_control` on model settings. The Anthropic
+        # chat model translates the kwarg to the correct wire format for the
+        # active transport: direct API receives it as-is, while Bedrock has it
+        # expanded into a block-level breakpoint by `_get_request_payload`.
+        overrides["model_settings"] = {
+            **request.model_settings,
+            "cache_control": cache_control,
+        }
+
+        system_message = _tag_system_message(request.system_message, cache_control)
+        if system_message is not request.system_message:
+            overrides["system_message"] = system_message
+
+        tools = _tag_tools(request.tools, cache_control)
+        if tools is not request.tools:
+            overrides["tools"] = tools
+
+        return request.override(**overrides)
+
+    def wrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelCallResult:
+        """Modify the model request to add cache control blocks.
+
+        Args:
+            request: The model request to potentially modify.
+            handler: The handler to execute the model request.
+
+        Returns:
+            The model response from the handler.
+        """
+        if not self._should_apply_caching(request):
+            return handler(request)
+
+        return handler(self._apply_caching(request))
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        """Modify the model request to add cache control blocks (async version).
+
+        Args:
+            request: The model request to potentially modify.
+            handler: The async handler to execute the model request.
+
+        Returns:
+            The model response from the handler.
+        """
+        if not self._should_apply_caching(request):
+            return await handler(request)
+
+        return await handler(self._apply_caching(request))
+
+
+def _tag_system_message(
+    system_message: Any,
+    cache_control: dict[str, str],
+) -> Any:
+    """Tag the last content block of a system message with cache_control.
+
+    Returns the original system_message unchanged if there are no blocks
+    to tag.
+
+    Args:
+        system_message: The system message to tag.
+        cache_control: The cache control dict to apply.
+
+    Returns:
+        A new SystemMessage with cache_control on the last block, or the
+        original if no modification was needed.
+    """
+    if system_message is None:
+        return system_message
+
+    content = system_message.content
+    if isinstance(content, str):
+        if not content:
+            return system_message
+        new_content: list[str | dict[str, Any]] = [
+            {"type": "text", "text": content, "cache_control": cache_control}
+        ]
+    elif isinstance(content, list):
+        if not content:
+            return system_message
+        new_content = list(content)
+        last = new_content[-1]
+        base = last if isinstance(last, dict) else {}
+        new_content[-1] = {**base, "cache_control": cache_control}
+    else:
+        return system_message
+
+    return SystemMessage(content=new_content)
+
+
+def _tag_tools(
+    tools: list[Any] | None,
+    cache_control: dict[str, str],
+) -> list[Any] | None:
+    """Tag the last tool with cache_control via its extras dict.
+
+    Only the last tool is tagged to minimize the number of explicit cache
+    breakpoints (Anthropic limits these to 4 per request). Since tool
+    definitions are sent as a contiguous block, a single breakpoint on the
+    last tool caches the entire set.
+
+    Creates a copy of the last tool with cache_control added to extras,
+    without mutating the original.
+
+    Args:
+        tools: The list of tools to tag.
+        cache_control: The cache control dict to apply.
+
+    Returns:
+        A new list with cache_control on the last tool's extras, or the
+        original if no tools are present.
+    """
+    if not tools:
+        return tools
+
+    last = tools[-1]
+    if not isinstance(last, BaseTool):
+        return tools
+
+    new_extras = {**(last.extras or {}), "cache_control": cache_control}
+    return [*tools[:-1], last.model_copy(update={"extras": new_extras})]

coco_anthropic/output_parsers.py

@@ -0,0 +1,93 @@
+"""Output parsers for Anthropic tool calls."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+from coco_core.messages import AIMessage, ToolCall
+from coco_core.messages.tool import tool_call
+from coco_core.output_parsers import BaseGenerationOutputParser
+from coco_core.outputs import ChatGeneration, Generation
+from pydantic import BaseModel, ConfigDict
+
+
+class ToolsOutputParser(BaseGenerationOutputParser):
+    """Output parser for tool calls."""
+
+    first_tool_only: bool = False
+    """Whether to return only the first tool call."""
+    args_only: bool = False
+    """Whether to return only the arguments of the tool calls."""
+    pydantic_schemas: list[type[BaseModel]] | None = None
+    """Pydantic schemas to parse tool calls into."""
+
+    model_config = ConfigDict(
+        extra="forbid",
+    )
+
+    def parse_result(self, result: list[Generation], *, partial: bool = False) -> Any:
+        """Parse a list of candidate model Generations into a specific format.
+
+        Args:
+            result: A list of `Generation` to be parsed. The Generations are assumed
+                to be different candidate outputs for a single model input.
+            partial: (Not used) Whether the result is a partial result. If `True`, the
+                parser may return a partial result, which may not be complete or valid.
+
+        Returns:
+            Structured output.
+
+        """
+        if not result or not isinstance(result[0], ChatGeneration):
+            return None if self.first_tool_only else []
+        message = cast("AIMessage", result[0].message)
+        tool_calls: list = [
+            dict(tc) for tc in _extract_tool_calls_from_message(message)
+        ]
+        if isinstance(message.content, list):
+            # Map tool call id to index
+            id_to_index = {
+                block["id"]: i
+                for i, block in enumerate(message.content)
+                if isinstance(block, dict) and block["type"] == "tool_use"
+            }
+            tool_calls = [{**tc, "index": id_to_index[tc["id"]]} for tc in tool_calls]
+        if self.pydantic_schemas:
+            tool_calls = [self._pydantic_parse(tc) for tc in tool_calls]
+        elif self.args_only:
+            tool_calls = [tc["args"] for tc in tool_calls]
+        else:
+            pass
+
+        if self.first_tool_only:
+            return tool_calls[0] if tool_calls else None
+        return list(tool_calls)
+
+    def _pydantic_parse(self, tool_call: dict) -> BaseModel:
+        cls_ = {schema.__name__: schema for schema in self.pydantic_schemas or []}[
+            tool_call["name"]
+        ]
+        return cls_(**tool_call["args"])
+
+
+def _extract_tool_calls_from_message(message: AIMessage) -> list[ToolCall]:
+    """Extract tool calls from a list of content blocks."""
+    if message.tool_calls:
+        return message.tool_calls
+    return extract_tool_calls(message.content)
+
+
+def extract_tool_calls(content: str | list[str | dict[str, Any]]) -> list[ToolCall]:
+    """Extract tool calls from a list of content blocks."""
+    if isinstance(content, list):
+        tool_calls = []
+        for block in content:
+            if isinstance(block, str):
+                continue
+            if block["type"] != "tool_use":
+                continue
+            tool_calls.append(
+                tool_call(name=block["name"], args=block["input"], id=block["id"]),
+            )
+        return tool_calls
+    return []

coco_core/__init__.py

@@ -0,0 +1,20 @@
+"""`trustcoco-core` defines the base abstractions for the Coco ecosystem.
+
+The interfaces for core components like chat models, LLMs, vector stores, retrievers,
+and more are defined here. The universal invocation protocol (Runnables) along with
+a syntax for combining components are also defined here.
+
+**No third-party integrations are defined here.** The dependencies are kept purposefully
+very lightweight.
+"""
+
+from coco_core._api import (
+    surface_coco_beta_warnings,
+    surface_coco_deprecation_warnings,
+)
+from coco_core.version import VERSION
+
+__version__ = VERSION
+
+surface_coco_deprecation_warnings()
+surface_coco_beta_warnings()

coco_core/_api/__init__.py

@@ -0,0 +1,87 @@
+"""Helper functions for managing the Coco API.
+
+This module is only relevant for Coco developers, not for users.
+
+!!! warning
+
+    This module and its submodules are for internal use only. Do not use them in your
+    own code. We may change the API at any time with no warning.
+"""
+
+from typing import TYPE_CHECKING
+
+from coco_core._import_utils import import_attr
+
+if TYPE_CHECKING:
+    from coco_core._api.beta_decorator import (
+        CocoBetaWarning,
+        beta,
+        suppress_coco_beta_warning,
+        surface_coco_beta_warnings,
+    )
+    from coco_core._api.deprecation import (
+        CocoDeprecationWarning,
+        deprecated,
+        suppress_coco_deprecation_warning,
+        surface_coco_deprecation_warnings,
+        warn_deprecated,
+    )
+    from coco_core._api.path import as_import_path, get_relative_path
+
+__all__ = (
+    "CocoBetaWarning",
+    "CocoDeprecationWarning",
+    "as_import_path",
+    "beta",
+    "deprecated",
+    "get_relative_path",
+    "suppress_coco_beta_warning",
+    "suppress_coco_deprecation_warning",
+    "surface_coco_beta_warnings",
+    "surface_coco_deprecation_warnings",
+    "warn_deprecated",
+)
+
+_dynamic_imports = {
+    "CocoBetaWarning": "beta_decorator",
+    "beta": "beta_decorator",
+    "suppress_coco_beta_warning": "beta_decorator",
+    "surface_coco_beta_warnings": "beta_decorator",
+    "as_import_path": "path",
+    "get_relative_path": "path",
+    "CocoDeprecationWarning": "deprecation",
+    "deprecated": "deprecation",
+    "surface_coco_deprecation_warnings": "deprecation",
+    "suppress_coco_deprecation_warning": "deprecation",
+    "warn_deprecated": "deprecation",
+}
+
+
+def __getattr__(attr_name: str) -> object:
+    """Dynamically import and return an attribute from a submodule.
+
+    This function enables lazy loading of API functions from submodules, reducing
+    initial import time and circular dependency issues.
+
+    Args:
+        attr_name: Name of the attribute to import.
+
+    Returns:
+        The imported attribute object.
+
+    Raises:
+        AttributeError: If the attribute is not a valid dynamic import.
+    """
+    module_name = _dynamic_imports.get(attr_name)
+    result = import_attr(attr_name, module_name, __spec__.parent)
+    globals()[attr_name] = result
+    return result
+
+
+def __dir__() -> list[str]:
+    """Return a list of available attributes for this module.
+
+    Returns:
+        List of attribute names that can be imported from this module.
+    """
+    return list(__all__)