fix(v2): accumulate all billable usage fields across retries (#2493)

instructor/v2/core/usage.py

@@ -1,10 +1,26 @@
-"""Usage accumulation helpers owned by the v2 runtime."""
+"""Usage accumulation helpers owned by the v2 runtime.
+
+Accumulating token usage across retries used to enumerate the fields to sum by
+hand, once per provider. Every counter a provider SDK added afterwards was then
+silently dropped -- left stale (Anthropic) or overwritten with ``None`` when the
+accumulator's details object was copied back onto the response (OpenAI). The
+counters people check first (``input_tokens`` / ``prompt_tokens``) stayed
+correct, so the under-count only surfaced when reconciling against an invoice.
+
+The accumulator below instead discovers fields from the model and sums them
+generically -- including the numeric leaves of nested sub-models such as
+``cache_creation``, ``server_tool_use`` and ``*_tokens_details`` -- so new
+billable fields are picked up automatically rather than needing a matching edit
+in two places every time.
+"""
 
 from __future__ import annotations
 
 import logging
 from typing import TYPE_CHECKING, TypeVar
 
+from pydantic import BaseModel
+
 if TYPE_CHECKING:
     from anthropic.types import Usage as AnthropicUsage
     from openai.types import CompletionUsage as OpenAIUsage
@@ -13,6 +29,82 @@
 T_Response = TypeVar("T_Response")
 
 
+def _zero_numeric(model: BaseModel) -> None:
+    """Set every numeric leaf of ``model`` to 0, recursing into sub-models."""
+    for name in type(model).model_fields:
+        value = getattr(model, name, None)
+        if isinstance(value, BaseModel):
+            _zero_numeric(value)
+        elif isinstance(value, bool):
+            continue  # bool is an int subclass; treat as a flag, not a counter
+        elif isinstance(value, (int, float)):
+            setattr(model, name, 0)
+
+
+def _zeroed_copy(model: BaseModel) -> BaseModel:
+    """A deep copy of ``model`` with all numeric leaves zeroed."""
+    clone = model.model_copy(deep=True)
+    _zero_numeric(clone)
+    return clone
+
+
+def _accumulate_into(total: BaseModel, response: BaseModel) -> None:
+    """Add the numeric fields of ``response`` into ``total``, in place.
+
+    Numeric fields (and the numeric leaves of nested sub-models) are summed;
+    non-numeric fields (``service_tier``, ``inference_geo``, ...) take the
+    latest reported value; a field that ``response`` reports as ``None`` leaves
+    the running total untouched.
+    """
+    for name in type(response).model_fields:
+        value = getattr(response, name, None)
+        if isinstance(value, BaseModel):
+            current = getattr(total, name, None)
+            if not isinstance(current, BaseModel):
+                # First attempt to report this sub-model: adopt a zeroed copy so
+                # the response's own numbers are counted rather than skipped
+                # (skipping them is what left the accumulator's field ``None``).
+                current = _zeroed_copy(value)
+                setattr(total, name, current)
+            _accumulate_into(current, value)
+        elif isinstance(value, bool):
+            setattr(total, name, value)
+        elif isinstance(value, (int, float)):
+            running = getattr(total, name, None)
+            setattr(total, name, (running or 0) + value)
+        elif value is not None:
+            setattr(total, name, value)
+
+
+def _sync_into(response: BaseModel, total: BaseModel) -> None:
+    """Mirror ``total`` back onto ``response``, in place.
+
+    Existing (sub-)model instances on ``response`` are mutated rather than
+    replaced, so a provider-specific ``Usage`` subclass keeps its type.
+    """
+    for name in type(total).model_fields:
+        value = getattr(total, name, None)
+        if isinstance(value, BaseModel):
+            existing = getattr(response, name, None)
+            if isinstance(existing, BaseModel):
+                _sync_into(existing, value)
+            else:
+                setattr(response, name, value.model_copy(deep=True))
+        else:
+            setattr(response, name, value)
+
+
+def accumulate_usage(total: BaseModel, response_usage: BaseModel) -> None:
+    """Accumulate ``response_usage`` into ``total`` and mirror the running total
+    back onto ``response_usage`` (both mutated in place).
+
+    Works for any pydantic usage model -- the OpenAI ``CompletionUsage`` and the
+    Anthropic ``Usage`` are both handled by the same generic walk.
+    """
+    _accumulate_into(total, response_usage)
+    _sync_into(response_usage, total)
+
+
 def update_total_usage(
     response: T_Response | None,
     total_usage: OpenAIUsage | AnthropicUsage,
@@ -26,34 +118,7 @@ def update_total_usage(
     if isinstance(response_usage, _OpenAIUsage) and isinstance(
         total_usage, _OpenAIUsage
     ):
-        total_usage.completion_tokens += response_usage.completion_tokens or 0
-        total_usage.prompt_tokens += response_usage.prompt_tokens or 0
-        total_usage.total_tokens += response_usage.total_tokens or 0
-        if (rtd := response_usage.completion_tokens_details) and (
-            ttd := total_usage.completion_tokens_details
-        ):
-            ttd.audio_tokens = (ttd.audio_tokens or 0) + (rtd.audio_tokens or 0)
-            ttd.reasoning_tokens = (ttd.reasoning_tokens or 0) + (
-                rtd.reasoning_tokens or 0
-            )
-        if (rpd := response_usage.prompt_tokens_details) and (
-            tpd := total_usage.prompt_tokens_details
-        ):
-            tpd.audio_tokens = (tpd.audio_tokens or 0) + (rpd.audio_tokens or 0)
-            tpd.cached_tokens = (tpd.cached_tokens or 0) + (rpd.cached_tokens or 0)
-        response_usage.completion_tokens = total_usage.completion_tokens
-        response_usage.prompt_tokens = total_usage.prompt_tokens
-        response_usage.total_tokens = total_usage.total_tokens
-        response_usage.completion_tokens_details = (
-            total_usage.completion_tokens_details.model_copy(deep=True)
-            if total_usage.completion_tokens_details is not None
-            else None
-        )
-        response_usage.prompt_tokens_details = (
-            total_usage.prompt_tokens_details.model_copy(deep=True)
-            if total_usage.prompt_tokens_details is not None
-            else None
-        )
+        accumulate_usage(total_usage, response_usage)
         return response
 
     try:

instructor/v2/providers/anthropic/usage.py

@@ -18,26 +18,24 @@ def initialize_usage() -> Any:
 
 
 def update_total_usage(response_usage: Any, total_usage: Any) -> bool:
-    """Accumulate Anthropic token usage into a running total when applicable."""
+    """Accumulate Anthropic token usage into a running total when applicable.
+
+    Every numeric field is summed generically, including the leaves of nested
+    sub-models such as ``cache_creation`` (``ephemeral_*_input_tokens``),
+    ``server_tool_use`` (``web_search_requests`` / ``web_fetch_requests``) and
+    ``output_tokens_details`` (``thinking_tokens``). Billable counters added by
+    newer ``anthropic`` SDK releases are therefore picked up automatically
+    instead of being left stale at whatever the last attempt reported.
+    """
     from anthropic.types import Usage as AnthropicUsage
 
     if not isinstance(response_usage, AnthropicUsage) or not isinstance(
         total_usage, AnthropicUsage
     ):
         return False
 
-    if not total_usage.cache_creation_input_tokens:
-        total_usage.cache_creation_input_tokens = 0
-    if not total_usage.cache_read_input_tokens:
-        total_usage.cache_read_input_tokens = 0
-    total_usage.input_tokens += response_usage.input_tokens or 0
-    total_usage.output_tokens += response_usage.output_tokens or 0
-    total_usage.cache_creation_input_tokens += (
-        response_usage.cache_creation_input_tokens or 0
-    )
-    total_usage.cache_read_input_tokens += response_usage.cache_read_input_tokens or 0
-    response_usage.input_tokens = total_usage.input_tokens
-    response_usage.output_tokens = total_usage.output_tokens
-    response_usage.cache_creation_input_tokens = total_usage.cache_creation_input_tokens
-    response_usage.cache_read_input_tokens = total_usage.cache_read_input_tokens
+    # Imported lazily to avoid a circular import at module load time.
+    from instructor.v2.core.usage import accumulate_usage
+
+    accumulate_usage(total_usage, response_usage)
     return True

tests/v2/test_usage_accumulation.py

@@ -0,0 +1,146 @@
+"""Regression tests for retry usage accumulation (issue #2493).
+
+Retries used to sum a hand-written list of token fields, so nested billable
+counters were either left stale (Anthropic) or wiped to ``None`` when the
+accumulator's details object was copied back onto the response (OpenAI). These
+tests reproduce the reported accounting bugs and assert the generic accumulator
+now sums them. All of it is pure accounting logic -- no API keys, no network.
+"""
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from instructor.v2.core.usage import accumulate_usage
+
+
+def test_anthropic_accumulates_nested_billable_fields() -> None:
+    from anthropic.types import Usage
+    from anthropic.types.cache_creation import CacheCreation
+    from anthropic.types.server_tool_usage import ServerToolUsage
+
+    from instructor.v2.providers.anthropic.usage import (
+        initialize_usage,
+        update_total_usage,
+    )
+
+    total = initialize_usage()
+    response: Usage | None = None
+    for _ in range(3):
+        response = Usage(
+            input_tokens=100,
+            output_tokens=50,
+            cache_creation_input_tokens=500,
+            cache_read_input_tokens=0,
+            cache_creation=CacheCreation(
+                ephemeral_5m_input_tokens=500,
+                ephemeral_1h_input_tokens=0,
+            ),
+            server_tool_use=ServerToolUsage(
+                web_search_requests=2,
+                web_fetch_requests=0,
+            ),
+        )
+        assert update_total_usage(response, total) is True
+
+    # Flat counters were already cumulative before the fix.
+    assert total.input_tokens == 300
+    assert total.output_tokens == 150
+    assert total.cache_creation_input_tokens == 1500
+    # Nested sub-models used to be frozen at the last attempt's value.
+    assert total.cache_creation.ephemeral_5m_input_tokens == 1500
+    assert total.server_tool_use.web_search_requests == 6
+    # The response is mirrored to the running total.
+    assert response is not None
+    assert response.input_tokens == 300
+    assert response.cache_creation.ephemeral_5m_input_tokens == 1500
+    assert response.server_tool_use.web_search_requests == 6
+
+
+def test_openai_accumulates_and_does_not_wipe_detail_fields() -> None:
+    from openai.types import CompletionUsage
+    from openai.types.completion_usage import (
+        CompletionTokensDetails,
+        PromptTokensDetails,
+    )
+
+    # Mirrors how the v2 retry runtime seeds the accumulator: the details
+    # objects only carry the two fields it happened to enumerate.
+    total = CompletionUsage(
+        completion_tokens=0,
+        prompt_tokens=0,
+        total_tokens=0,
+        completion_tokens_details=CompletionTokensDetails(
+            audio_tokens=0, reasoning_tokens=0
+        ),
+        prompt_tokens_details=PromptTokensDetails(audio_tokens=0, cached_tokens=0),
+    )
+    response: CompletionUsage | None = None
+    for _ in range(3):
+        response = CompletionUsage(
+            completion_tokens=50,
+            prompt_tokens=100,
+            total_tokens=150,
+            completion_tokens_details=CompletionTokensDetails(
+                reasoning_tokens=40,
+                accepted_prediction_tokens=7,
+                rejected_prediction_tokens=3,
+            ),
+            prompt_tokens_details=PromptTokensDetails(cached_tokens=500),
+        )
+        accumulate_usage(total, response)
+
+    assert total.prompt_tokens == 300
+    assert total.completion_tokens == 150
+    assert total.total_tokens == 450
+    assert total.completion_tokens_details.reasoning_tokens == 120
+    # These were overwritten with None before the fix (accumulator never
+    # populated them, then its details object was copied onto the response).
+    assert total.completion_tokens_details.accepted_prediction_tokens == 21
+    assert total.completion_tokens_details.rejected_prediction_tokens == 9
+    assert total.prompt_tokens_details.cached_tokens == 1500
+    assert response is not None
+    assert response.completion_tokens_details.accepted_prediction_tokens == 21
+
+
+def test_openai_usage_subclass_type_is_preserved() -> None:
+    # A provider may hand back a CompletionUsage subclass; mirroring the total
+    # must not reconstruct it as the base class.
+    from openai.types import CompletionUsage
+
+    class ProviderUsage(CompletionUsage):
+        pass
+
+    total = CompletionUsage(prompt_tokens=5, completion_tokens=3, total_tokens=8)
+    response = ProviderUsage(prompt_tokens=11, completion_tokens=7, total_tokens=18)
+
+    accumulate_usage(total, response)
+
+    assert isinstance(response, ProviderUsage)
+    assert response.prompt_tokens == 16
+    assert response.completion_tokens == 10
+    assert response.total_tokens == 26
+
+
+def test_generic_engine_handles_nesting_none_and_non_numeric() -> None:
+    # Version-independent proof of the accumulation rules on a synthetic model.
+    class Details(BaseModel):
+        a: int | None = None
+        b: int | None = None
+
+    class Usage(BaseModel):
+        x: int = 0
+        tier: str | None = None
+        details: Details | None = None
+
+    total = Usage()
+    # First attempt: accumulator's sub-model is None and must be adopted zeroed.
+    accumulate_usage(total, Usage(x=10, tier="standard", details=Details(a=1, b=2)))
+    # Second attempt: b is None and must not reset the running total.
+    accumulate_usage(total, Usage(x=5, tier="priority", details=Details(a=3)))
+
+    assert total.x == 15
+    assert total.tier == "priority"  # non-numeric: latest value wins
+    assert total.details is not None
+    assert total.details.a == 4
+    assert total.details.b == 2  # a None report left the prior total intact