diff --git a/instructor/v2/core/usage.py b/instructor/v2/core/usage.py index d220cf462..de8231fd3 100644 --- a/instructor/v2/core/usage.py +++ b/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: diff --git a/instructor/v2/providers/anthropic/usage.py b/instructor/v2/providers/anthropic/usage.py index 1bb27edd8..33f015cff 100644 --- a/instructor/v2/providers/anthropic/usage.py +++ b/instructor/v2/providers/anthropic/usage.py @@ -18,7 +18,15 @@ 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( @@ -26,18 +34,8 @@ def update_total_usage(response_usage: Any, total_usage: Any) -> bool: ): 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 diff --git a/tests/v2/test_usage_accumulation.py b/tests/v2/test_usage_accumulation.py new file mode 100644 index 000000000..6d63c8d90 --- /dev/null +++ b/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