coco-framework 0.1.0__py3-none-any.whl

coco_core/messages/ai.py

@@ -0,0 +1,854 @@
+"""AI message."""
+
+import itertools
+import json
+import logging
+import operator
+from collections.abc import Sequence
+from typing import Any, Literal, cast, overload
+
+from pydantic import Field, model_validator
+from typing_extensions import NotRequired, Self, TypedDict, override
+
+from coco_core.messages import content as types
+from coco_core.messages.base import (
+    BaseMessage,
+    BaseMessageChunk,
+    _extract_reasoning_from_additional_kwargs,
+    merge_content,
+)
+from coco_core.messages.content import InvalidToolCall
+from coco_core.messages.tool import (
+    ToolCall,
+    ToolCallChunk,
+    default_tool_chunk_parser,
+    default_tool_parser,
+)
+from coco_core.messages.tool import invalid_tool_call as create_invalid_tool_call
+from coco_core.messages.tool import tool_call as create_tool_call
+from coco_core.messages.tool import tool_call_chunk as create_tool_call_chunk
+from coco_core.utils._merge import merge_dicts, merge_lists
+from coco_core.utils.json import parse_partial_json
+from coco_core.utils.usage import _dict_int_op
+from coco_core.utils.utils import LC_AUTO_PREFIX, LC_ID_PREFIX
+
+logger = logging.getLogger(__name__)
+
+
+class InputTokenDetails(TypedDict, total=False):
+    """Breakdown of input token counts.
+
+    Does *not* need to sum to full input token count. Does *not* need to have all keys.
+
+    Example:
+        ```python
+        {
+            "audio": 10,
+            "cache_creation": 200,
+            "cache_read": 100,
+        }
+        ```
+
+    May also hold extra provider-specific keys.
+
+    !!! version-added "Added in `trustcoco-core` 0.3.9"
+    """
+
+    audio: int
+    """Audio input tokens."""
+
+    cache_creation: int
+    """Input tokens that were cached and there was a cache miss.
+
+    Since there was a cache miss, the cache was created from these tokens.
+    """
+
+    cache_read: int
+    """Input tokens that were cached and there was a cache hit.
+
+    Since there was a cache hit, the tokens were read from the cache. More precisely,
+    the model state given these tokens was read from the cache.
+    """
+
+
+class OutputTokenDetails(TypedDict, total=False):
+    """Breakdown of output token counts.
+
+    Does *not* need to sum to full output token count. Does *not* need to have all keys.
+
+    Example:
+        ```python
+        {
+            "audio": 10,
+            "reasoning": 200,
+        }
+        ```
+
+    May also hold extra provider-specific keys.
+
+    !!! version-added "Added in `trustcoco-core` 0.3.9"
+
+    """
+
+    audio: int
+    """Audio output tokens."""
+
+    reasoning: int
+    """Reasoning output tokens.
+
+    Tokens generated by the model in a chain of thought process that are not
+    returned as part of model output.
+    """
+
+
+class UsageMetadata(TypedDict):
+    """Usage metadata for a message, such as token counts.
+
+    This is a standard representation of token usage that is consistent across models.
+
+    Example:
+        ```python
+        {
+            "input_tokens": 350,
+            "output_tokens": 240,
+            "total_tokens": 590,
+            "input_token_details": {
+                "audio": 10,
+                "cache_creation": 200,
+                "cache_read": 100,
+            },
+            "output_token_details": {
+                "audio": 10,
+                "reasoning": 200,
+            },
+        }
+        ```
+
+    !!! warning "Behavior changed in `trustcoco-core` 0.3.9"
+
+        Added `input_token_details` and `output_token_details`.
+
+    !!! note "LangSmith SDK"
+
+        The LangSmith SDK also has a `UsageMetadata` class. While the two share fields,
+        LangSmith's `UsageMetadata` has additional fields to capture cost information
+        used by the LangSmith platform.
+    """
+
+    input_tokens: int
+    """Count of input (or prompt) tokens. Sum of all input token types."""
+
+    output_tokens: int
+    """Count of output (or completion) tokens. Sum of all output token types."""
+
+    total_tokens: int
+    """Total token count. Sum of `input_tokens` + `output_tokens`."""
+
+    input_token_details: NotRequired[InputTokenDetails]
+    """Breakdown of input token counts.
+
+    Does *not* need to sum to full input token count. Does *not* need to have all keys.
+    """
+
+    output_token_details: NotRequired[OutputTokenDetails]
+    """Breakdown of output token counts.
+
+    Does *not* need to sum to full output token count. Does *not* need to have all keys.
+    """
+
+
+class AIMessage(BaseMessage):
+    """Message from an AI.
+
+    An `AIMessage` is returned from a chat model as a response to a prompt.
+
+    This message represents the output of the model and consists of both
+    the raw output as returned by the model and standardized fields
+    (e.g., tool calls, usage metadata) added by the Coco framework.
+    """
+
+    tool_calls: list[ToolCall] = Field(default_factory=list)
+    """If present, tool calls associated with the message."""
+
+    invalid_tool_calls: list[InvalidToolCall] = Field(default_factory=list)
+    """If present, tool calls with parsing errors associated with the message."""
+
+    usage_metadata: UsageMetadata | None = None
+    """If present, usage metadata for a message, such as token counts.
+
+    This is a standard representation of token usage that is consistent across models.
+    """
+
+    type: Literal["ai"] = "ai"
+    """The type of the message (used for deserialization)."""
+
+    @overload
+    def __init__(
+        self,
+        content: str | list[str | dict[Any, Any]],
+        **kwargs: Any,
+    ) -> None: ...
+
+    @overload
+    def __init__(
+        self,
+        content: str | list[str | dict[Any, Any]] | None = None,
+        content_blocks: list[types.ContentBlock] | None = None,
+        **kwargs: Any,
+    ) -> None: ...
+
+    def __init__(
+        self,
+        content: str | list[str | dict[Any, Any]] | None = None,
+        content_blocks: list[types.ContentBlock] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        """Initialize an `AIMessage`.
+
+        Specify `content` as positional arg or `content_blocks` for typing.
+
+        Args:
+            content: The content of the message.
+            content_blocks: Typed standard content.
+            **kwargs: Additional arguments to pass to the parent class.
+        """
+        if content_blocks is not None:
+            # If there are tool calls in content_blocks, but not in tool_calls, add them
+            content_tool_calls = [
+                block for block in content_blocks if block.get("type") == "tool_call"
+            ]
+            if content_tool_calls and "tool_calls" not in kwargs:
+                kwargs["tool_calls"] = content_tool_calls
+
+            super().__init__(
+                content=cast("list[str | dict[Any, Any]]", content_blocks),
+                **kwargs,
+            )
+        else:
+            super().__init__(content=content, **kwargs)
+
+    @property
+    def lc_attributes(self) -> dict[str, Any]:
+        """Attributes to be serialized.
+
+        Includes all attributes, even if they are derived from other initialization
+        arguments.
+        """
+        return {
+            "tool_calls": self.tool_calls,
+            "invalid_tool_calls": self.invalid_tool_calls,
+        }
+
+    @property
+    def content_blocks(self) -> list[types.ContentBlock]:
+        """Return standard, typed `ContentBlock` dicts from the message.
+
+        If the message has a known model provider, use the provider-specific translator
+        first before falling back to best-effort parsing. For details, see the property
+        on `BaseMessage`.
+        """
+        if self.response_metadata.get("output_version") == "v1" and isinstance(
+            self.content, list
+        ):
+            # Only short-circuit when content is a list (assumed under v1 to
+            # already hold ContentBlock dicts; the cast is unchecked). See
+            # AIMessageChunk.content_blocks for full rationale.
+            return cast("list[types.ContentBlock]", self.content)
+
+        model_provider = self.response_metadata.get("model_provider")
+        if model_provider:
+            from coco_core.messages.block_translators import (  # noqa: PLC0415
+                get_translator,
+            )
+
+            translator = get_translator(model_provider)
+            if translator:
+                try:
+                    return translator["translate_content"](self)
+                except NotImplementedError:
+                    pass
+
+        # Otherwise, use best-effort parsing
+        blocks = super().content_blocks
+
+        if self.tool_calls:
+            # Add from tool_calls if missing from content
+            content_tool_call_ids = {
+                block.get("id")
+                for block in self.content
+                if isinstance(block, dict) and block.get("type") == "tool_call"
+            }
+            for tool_call in self.tool_calls:
+                if (id_ := tool_call.get("id")) and id_ not in content_tool_call_ids:
+                    tool_call_block: types.ToolCall = {
+                        "type": "tool_call",
+                        "id": id_,
+                        "name": tool_call["name"],
+                        "args": tool_call["args"],
+                    }
+                    if "index" in tool_call:
+                        tool_call_block["index"] = tool_call["index"]  # type: ignore[typeddict-item]
+                    if "extras" in tool_call:
+                        tool_call_block["extras"] = tool_call["extras"]  # type: ignore[typeddict-item]
+                    blocks.append(tool_call_block)
+
+        # Best-effort reasoning extraction from additional_kwargs
+        # Only add reasoning if not already present
+        # Insert before all other blocks to keep reasoning at the start
+        has_reasoning = any(block.get("type") == "reasoning" for block in blocks)
+        if not has_reasoning and (
+            reasoning_block := _extract_reasoning_from_additional_kwargs(self)
+        ):
+            blocks.insert(0, reasoning_block)
+
+        return blocks
+
+    # TODO: remove this logic if possible, reducing breaking nature of changes
+    @model_validator(mode="before")
+    @classmethod
+    def _backwards_compat_tool_calls(cls, values: dict[str, Any]) -> Any:
+        check_additional_kwargs = not any(
+            values.get(k)
+            for k in ("tool_calls", "invalid_tool_calls", "tool_call_chunks")
+        )
+        if check_additional_kwargs and (
+            raw_tool_calls := values.get("additional_kwargs", {}).get("tool_calls")
+        ):
+            try:
+                if issubclass(cls, AIMessageChunk):
+                    values["tool_call_chunks"] = default_tool_chunk_parser(
+                        raw_tool_calls
+                    )
+                else:
+                    parsed_tool_calls, parsed_invalid_tool_calls = default_tool_parser(
+                        raw_tool_calls
+                    )
+                    values["tool_calls"] = parsed_tool_calls
+                    values["invalid_tool_calls"] = parsed_invalid_tool_calls
+            except Exception:
+                logger.debug("Failed to parse tool calls", exc_info=True)
+
+        # Ensure "type" is properly set on all tool call-like dicts.
+        if tool_calls := values.get("tool_calls"):
+            values["tool_calls"] = [
+                create_tool_call(
+                    **{k: v for k, v in tc.items() if k not in {"type", "extras"}}
+                )
+                for tc in tool_calls
+            ]
+        if invalid_tool_calls := values.get("invalid_tool_calls"):
+            values["invalid_tool_calls"] = [
+                create_invalid_tool_call(**{k: v for k, v in tc.items() if k != "type"})
+                for tc in invalid_tool_calls
+            ]
+
+        if tool_call_chunks := values.get("tool_call_chunks"):
+            values["tool_call_chunks"] = [
+                create_tool_call_chunk(**{k: v for k, v in tc.items() if k != "type"})
+                for tc in tool_call_chunks
+            ]
+
+        return values
+
+    @override
+    def pretty_repr(self, html: bool = False) -> str:
+        """Return a pretty representation of the message for display.
+
+        Args:
+            html: Whether to return an HTML-formatted string.
+
+        Returns:
+            A pretty representation of the message.
+
+        Example:
+            ```python
+            from coco_core.messages import AIMessage
+
+            msg = AIMessage(
+                content="Let me check the weather.",
+                tool_calls=[
+                    {"name": "get_weather", "args": {"city": "Paris"}, "id": "1"}
+                ],
+            )
+            ```
+
+            Results in:
+            ```python
+            >>> print(msg.pretty_repr())
+            ================================== Ai Message ==================================
+
+            Let me check the weather.
+            Tool Calls:
+              get_weather (1)
+             Call ID: 1
+              Args:
+                city: Paris
+            ```
+        """  # noqa: E501
+        base = super().pretty_repr(html=html)
+        lines = []
+
+        def _format_tool_args(tc: ToolCall | InvalidToolCall) -> list[str]:
+            lines = [
+                f"  {tc.get('name', 'Tool')} ({tc.get('id')})",
+                f" Call ID: {tc.get('id')}",
+            ]
+            if tc.get("error"):
+                lines.append(f"  Error: {tc.get('error')}")
+            lines.append("  Args:")
+            args = tc.get("args")
+            if isinstance(args, str):
+                lines.append(f"    {args}")
+            elif isinstance(args, dict):
+                for arg, value in args.items():
+                    lines.append(f"    {arg}: {value}")
+            return lines
+
+        if self.tool_calls:
+            lines.append("Tool Calls:")
+            for tc in self.tool_calls:
+                lines.extend(_format_tool_args(tc))
+        if self.invalid_tool_calls:
+            lines.append("Invalid Tool Calls:")
+            for itc in self.invalid_tool_calls:
+                lines.extend(_format_tool_args(itc))
+        return (base.strip() + "\n" + "\n".join(lines)).strip()
+
+
+class AIMessageChunk(AIMessage, BaseMessageChunk):
+    """Message chunk from an AI (yielded when streaming)."""
+
+    # Ignoring mypy re-assignment here since we're overriding the value
+    # to make sure that the chunk variant can be discriminated from the
+    # non-chunk variant.
+    type: Literal["AIMessageChunk"] = "AIMessageChunk"  # type: ignore[assignment]
+    """The type of the message (used for deserialization)."""
+
+    tool_call_chunks: list[ToolCallChunk] = Field(default_factory=list)
+    """If provided, tool call chunks associated with the message."""
+
+    chunk_position: Literal["last"] | None = None
+    """Optional span represented by an aggregated `AIMessageChunk`.
+
+    If a chunk with `chunk_position="last"` is aggregated into a stream,
+    `tool_call_chunks` in message content will be parsed into `tool_calls`.
+    """
+
+    @property
+    @override
+    def lc_attributes(self) -> dict[str, Any]:
+        return {
+            "tool_calls": self.tool_calls,
+            "invalid_tool_calls": self.invalid_tool_calls,
+        }
+
+    @property
+    def content_blocks(self) -> list[types.ContentBlock]:
+        """Return standard, typed `ContentBlock` dicts from the message."""
+        if self.response_metadata.get("output_version") == "v1" and isinstance(
+            self.content, list
+        ):
+            # Only short-circuit when content is already a list of ContentBlock
+            # dicts. Some streaming implementations keep content as a string
+            # even when output_version="v1" is set (e.g., OpenAI Chat
+            # Completions), so it must fall through to the model_provider
+            # translator which builds ContentBlock dicts from tool_calls /
+            # tool_call_chunks. Without this guard, string content would be
+            # returned directly, silently dropping tool calls.
+            return cast("list[types.ContentBlock]", self.content)
+
+        model_provider = self.response_metadata.get("model_provider")
+        if model_provider:
+            from coco_core.messages.block_translators import (  # noqa: PLC0415
+                get_translator,
+            )
+
+            translator = get_translator(model_provider)
+            if translator:
+                try:
+                    return translator["translate_content_chunk"](self)
+                except NotImplementedError:
+                    pass
+
+        # Otherwise, use best-effort parsing
+        blocks = super().content_blocks
+
+        if (
+            self.tool_call_chunks
+            and not self.content
+            and self.chunk_position != "last"  # keep tool_calls if aggregated
+        ):
+            blocks = [
+                block
+                for block in blocks
+                if block["type"] not in {"tool_call", "invalid_tool_call"}
+            ]
+            for tool_call_chunk in self.tool_call_chunks:
+                tc: types.ToolCallChunk = {
+                    "type": "tool_call_chunk",
+                    "id": tool_call_chunk.get("id"),
+                    "name": tool_call_chunk.get("name"),
+                    "args": tool_call_chunk.get("args"),
+                }
+                if (idx := tool_call_chunk.get("index")) is not None:
+                    tc["index"] = idx
+                blocks.append(tc)
+
+        # Best-effort reasoning extraction from additional_kwargs
+        # Only add reasoning if not already present
+        # Insert before all other blocks to keep reasoning at the start
+        has_reasoning = any(block.get("type") == "reasoning" for block in blocks)
+        if not has_reasoning and (
+            reasoning_block := _extract_reasoning_from_additional_kwargs(self)
+        ):
+            blocks.insert(0, reasoning_block)
+
+        return blocks
+
+    @model_validator(mode="after")
+    def init_tool_calls(self) -> Self:
+        """Initialize tool calls from tool call chunks.
+
+        Returns:
+            The values with tool calls initialized.
+
+        Raises:
+            ValueError: If the tool call chunks are malformed.
+        """
+        if not self.tool_call_chunks:
+            if self.tool_calls:
+                self.tool_call_chunks = [
+                    create_tool_call_chunk(
+                        name=tc["name"],
+                        args=json.dumps(tc["args"]),
+                        id=tc["id"],
+                        index=None,
+                    )
+                    for tc in self.tool_calls
+                ]
+            if self.invalid_tool_calls:
+                tool_call_chunks = self.tool_call_chunks
+                tool_call_chunks.extend(
+                    [
+                        create_tool_call_chunk(
+                            name=tc["name"], args=tc["args"], id=tc["id"], index=None
+                        )
+                        for tc in self.invalid_tool_calls
+                    ]
+                )
+                self.tool_call_chunks = tool_call_chunks
+
+            return self
+        tool_calls = []
+        invalid_tool_calls = []
+
+        def add_chunk_to_invalid_tool_calls(chunk: ToolCallChunk) -> None:
+            invalid_tool_calls.append(
+                create_invalid_tool_call(
+                    name=chunk["name"],
+                    args=chunk["args"],
+                    id=chunk["id"],
+                    error=None,
+                )
+            )
+
+        for chunk in self.tool_call_chunks:
+            try:
+                args_ = parse_partial_json(chunk["args"]) if chunk["args"] else {}
+                if isinstance(args_, dict):
+                    tool_calls.append(
+                        create_tool_call(
+                            name=chunk["name"] or "",
+                            args=args_,
+                            id=chunk["id"],
+                        )
+                    )
+                else:
+                    add_chunk_to_invalid_tool_calls(chunk)
+            except Exception:
+                add_chunk_to_invalid_tool_calls(chunk)
+        self.tool_calls = tool_calls
+        self.invalid_tool_calls = invalid_tool_calls
+
+        if (
+            self.chunk_position == "last"
+            and self.tool_call_chunks
+            and self.response_metadata.get("output_version") == "v1"
+            and isinstance(self.content, list)
+        ):
+            id_to_tc: dict[str, types.ToolCall] = {
+                cast("str", tc.get("id")): {
+                    "type": "tool_call",
+                    "name": tc["name"],
+                    "args": tc["args"],
+                    "id": tc.get("id"),
+                }
+                for tc in self.tool_calls
+                if "id" in tc
+            }
+            for idx, block in enumerate(self.content):
+                if (
+                    isinstance(block, dict)
+                    and block.get("type") == "tool_call_chunk"
+                    and (call_id := block.get("id"))
+                    and call_id in id_to_tc
+                ):
+                    self.content[idx] = cast("dict[str, Any]", id_to_tc[call_id])
+                    if "extras" in block:
+                        # mypy does not account for instance check for dict above
+                        self.content[idx]["extras"] = block["extras"]  # type: ignore[index]
+
+        return self
+
+    @model_validator(mode="after")
+    def init_server_tool_calls(self) -> Self:
+        """Initialize server tool calls.
+
+        Parse `server_tool_call_chunks` from
+        [`ServerToolCallChunk`][coco.messages.ServerToolCallChunk] objects.
+        """
+        if (
+            self.chunk_position == "last"
+            and self.response_metadata.get("output_version") == "v1"
+            and isinstance(self.content, list)
+        ):
+            for idx, block in enumerate(self.content):
+                if (
+                    isinstance(block, dict)
+                    and block.get("type")
+                    in {"server_tool_call", "server_tool_call_chunk"}
+                    and (args_str := block.get("args"))
+                    and isinstance(args_str, str)
+                ):
+                    try:
+                        args = json.loads(args_str)
+                        if isinstance(args, dict):
+                            self.content[idx]["type"] = "server_tool_call"  # type: ignore[index]
+                            self.content[idx]["args"] = args  # type: ignore[index]
+                    except json.JSONDecodeError:
+                        pass
+        return self
+
+    @overload  # type: ignore[override]  # summing BaseMessages gives ChatPromptTemplate
+    def __add__(self, other: "AIMessageChunk") -> "AIMessageChunk": ...
+
+    @overload
+    def __add__(self, other: Sequence["AIMessageChunk"]) -> "AIMessageChunk": ...
+
+    @overload
+    def __add__(self, other: Any) -> BaseMessageChunk: ...
+
+    @override
+    def __add__(self, other: Any) -> BaseMessageChunk:
+        if isinstance(other, AIMessageChunk):
+            return add_ai_message_chunks(self, other)
+        if isinstance(other, (list, tuple)) and all(
+            isinstance(o, AIMessageChunk) for o in other
+        ):
+            return add_ai_message_chunks(self, *other)
+        return super().__add__(other)
+
+
+def add_ai_message_chunks(
+    left: AIMessageChunk, *others: AIMessageChunk
+) -> AIMessageChunk:
+    """Add multiple `AIMessageChunk`s together.
+
+    Args:
+        left: The first `AIMessageChunk`.
+        *others: Other `AIMessageChunk`s to add.
+
+    Returns:
+        The resulting `AIMessageChunk`.
+
+    """
+    content = merge_content(left.content, *(o.content for o in others))
+    additional_kwargs = merge_dicts(
+        left.additional_kwargs, *(o.additional_kwargs for o in others)
+    )
+    response_metadata = merge_dicts(
+        left.response_metadata, *(o.response_metadata for o in others)
+    )
+
+    # Merge tool call chunks
+    if raw_tool_calls := merge_lists(
+        left.tool_call_chunks, *(o.tool_call_chunks for o in others)
+    ):
+        tool_call_chunks = [
+            create_tool_call_chunk(
+                name=rtc.get("name"),
+                args=rtc.get("args"),
+                index=rtc.get("index"),
+                id=rtc.get("id"),
+            )
+            for rtc in raw_tool_calls
+        ]
+    else:
+        tool_call_chunks = []
+
+    # Token usage
+    if left.usage_metadata or any(o.usage_metadata is not None for o in others):
+        usage_metadata: UsageMetadata | None = left.usage_metadata
+        for other in others:
+            usage_metadata = add_usage(usage_metadata, other.usage_metadata)
+    else:
+        usage_metadata = None
+
+    # Ranks are defined by the order of preference. Higher is better:
+    # 2. Provider-assigned IDs (non lc_* and non lc_run-*)
+    # 1. lc_run-* IDs
+    # 0. lc_* and other remaining IDs
+    best_rank = -1
+    chunk_id = None
+    candidates = itertools.chain([left.id], (o.id for o in others))
+
+    for id_ in candidates:
+        if not id_:
+            continue
+
+        if not id_.startswith(LC_ID_PREFIX) and not id_.startswith(LC_AUTO_PREFIX):
+            chunk_id = id_
+            # Highest rank, return instantly
+            break
+
+        rank = 1 if id_.startswith(LC_ID_PREFIX) else 0
+
+        if rank > best_rank:
+            best_rank = rank
+            chunk_id = id_
+
+    chunk_position: Literal["last"] | None = (
+        "last" if any(x.chunk_position == "last" for x in [left, *others]) else None
+    )
+
+    return left.__class__(
+        content=content,
+        additional_kwargs=additional_kwargs,
+        tool_call_chunks=tool_call_chunks,
+        response_metadata=response_metadata,
+        usage_metadata=usage_metadata,
+        id=chunk_id,
+        chunk_position=chunk_position,
+    )
+
+
+def add_usage(left: UsageMetadata | None, right: UsageMetadata | None) -> UsageMetadata:
+    """Recursively add two UsageMetadata objects.
+
+    Example:
+        ```python
+        from coco_core.messages.ai import add_usage
+
+        left = UsageMetadata(
+            input_tokens=5,
+            output_tokens=0,
+            total_tokens=5,
+            input_token_details=InputTokenDetails(cache_read=3),
+        )
+        right = UsageMetadata(
+            input_tokens=0,
+            output_tokens=10,
+            total_tokens=10,
+            output_token_details=OutputTokenDetails(reasoning=4),
+        )
+
+        add_usage(left, right)
+        ```
+
+        results in
+
+        ```python
+        UsageMetadata(
+            input_tokens=5,
+            output_tokens=10,
+            total_tokens=15,
+            input_token_details=InputTokenDetails(cache_read=3),
+            output_token_details=OutputTokenDetails(reasoning=4),
+        )
+        ```
+    Args:
+        left: The first `UsageMetadata` object.
+        right: The second `UsageMetadata` object.
+
+    Returns:
+        The sum of the two `UsageMetadata` objects.
+
+    """
+    if not (left or right):
+        return UsageMetadata(input_tokens=0, output_tokens=0, total_tokens=0)
+    if not (left and right):
+        return cast("UsageMetadata", left or right)
+
+    return UsageMetadata(
+        **cast(
+            "UsageMetadata",
+            _dict_int_op(
+                cast("dict[str, Any]", left),
+                cast("dict[str, Any]", right),
+                operator.add,
+            ),
+        )
+    )
+
+
+def subtract_usage(
+    left: UsageMetadata | None, right: UsageMetadata | None
+) -> UsageMetadata:
+    """Recursively subtract two `UsageMetadata` objects.
+
+    Token counts cannot be negative so the actual operation is `max(left - right, 0)`.
+
+    Example:
+        ```python
+        from coco_core.messages.ai import subtract_usage
+
+        left = UsageMetadata(
+            input_tokens=5,
+            output_tokens=10,
+            total_tokens=15,
+            input_token_details=InputTokenDetails(cache_read=4),
+        )
+        right = UsageMetadata(
+            input_tokens=3,
+            output_tokens=8,
+            total_tokens=11,
+            output_token_details=OutputTokenDetails(reasoning=4),
+        )
+
+        subtract_usage(left, right)
+        ```
+
+        results in
+
+        ```python
+        UsageMetadata(
+            input_tokens=2,
+            output_tokens=2,
+            total_tokens=4,
+            input_token_details=InputTokenDetails(cache_read=4),
+            output_token_details=OutputTokenDetails(reasoning=0),
+        )
+        ```
+    Args:
+        left: The first `UsageMetadata` object.
+        right: The second `UsageMetadata` object.
+
+    Returns:
+        The resulting `UsageMetadata` after subtraction.
+
+    """
+    if not (left or right):
+        return UsageMetadata(input_tokens=0, output_tokens=0, total_tokens=0)
+    if not (left and right):
+        return cast("UsageMetadata", left or right)
+
+    return UsageMetadata(
+        **cast(
+            "UsageMetadata",
+            _dict_int_op(
+                cast("dict[str, Any]", left),
+                cast("dict[str, Any]", right),
+                (lambda le, ri: max(le - ri, 0)),
+            ),
+        )
+    )