From c69f4aad2c69943200281c21e4f775c6bbaaed06 Mon Sep 17 00:00:00 2001 From: Kazuhiro Sera Date: Fri, 7 Aug 2026 18:26:21 +0900 Subject: [PATCH] feat(models): preserve raw usage payloads --- src/agents/extensions/models/any_llm_model.py | 85 +++++++++++++++++-- src/agents/items.py | 9 ++ src/agents/model_settings.py | 11 +++ src/agents/models/chatcmpl_stream_handler.py | 15 +++- src/agents/models/openai_chatcompletions.py | 11 ++- src/agents/models/openai_responses.py | 15 +++- src/agents/run_internal/run_loop.py | 7 +- src/agents/usage.py | 60 ++++++++++++- tests/model_settings/test_serialization.py | 11 ++- tests/models/test_any_llm_model.py | 11 ++- .../test_litellm_chatcompletions_stream.py | 5 +- tests/models/test_litellm_logprobs.py | 5 +- tests/models/test_openai_chatcompletions.py | 81 ++++++++++++++++++ .../test_openai_chatcompletions_stream.py | 17 +++- tests/models/test_openai_responses.py | 4 +- tests/test_agent_runner_streamed.py | 25 +++++- tests/test_run_state.py | 4 + tests/test_usage.py | 48 ++++++++++- 18 files changed, 397 insertions(+), 27 deletions(-) diff --git a/src/agents/extensions/models/any_llm_model.py b/src/agents/extensions/models/any_llm_model.py index 74538bcd7d..9169c49508 100644 --- a/src/agents/extensions/models/any_llm_model.py +++ b/src/agents/extensions/models/any_llm_model.py @@ -6,7 +6,7 @@ import inspect import json import time -from collections.abc import AsyncGenerator, AsyncIterator, Iterable +from collections.abc import AsyncGenerator, AsyncIterator, Iterable, Mapping from copy import copy from typing import TYPE_CHECKING, Any, Literal, cast, overload @@ -52,7 +52,12 @@ from ...tracing import generation_span, response_span from ...tracing.span_data import GenerationSpanData from ...tracing.spans import Span -from ...usage import Usage +from ...usage import ( + Usage, + _attach_raw_usage_snapshot, + _extract_raw_usage_snapshot, + _raw_usage_snapshot, +) from ...util._error_tracing import model_span_errors, record_model_error_on_span from ...util._json import _to_dump_compatible @@ -75,6 +80,12 @@ class InternalChatCompletionMessage(ChatCompletionMessage): reasoning_content: str = "" +def _usage_payload(response: Any) -> Any | None: + if isinstance(response, Mapping): + return response.get("usage") + return getattr(response, "usage", None) + + class _AnyLLMResponsesParamsShim: """Fallback shim for tests and older any-llm layouts.""" @@ -417,6 +428,11 @@ async def _get_response_via_responses( usage=usage, response_id=response.id, request_id=getattr(response, "_request_id", None), + raw_usage=( + _extract_raw_usage_snapshot(response, fallback=response.usage) + if model_settings.preserve_raw_usage is True + else None + ), ) async def _stream_response_via_responses( @@ -463,6 +479,8 @@ async def _stream_response_via_responses( chunk_type = getattr(chunk, "type", None) if isinstance(chunk, ResponseCompletedEvent): final_response = chunk.response + if model_settings.preserve_raw_usage is True: + _attach_raw_usage_snapshot(chunk.response, chunk.response.usage) elif chunk_type in {"response.failed", "response.incomplete"}: terminal_response = getattr(chunk, "response", None) terminal_failure_error = response_terminal_failure_error( @@ -640,7 +658,16 @@ async def _get_response_via_chat( if logprob_models: self._attach_logprobs_to_output(items, logprob_models) - return ModelResponse(output=items, usage=usage, response_id=None) + return ModelResponse( + output=items, + usage=usage, + response_id=None, + raw_usage=( + _extract_raw_usage_snapshot(response, fallback=response.usage) + if model_settings.preserve_raw_usage is True + else None + ), + ) async def _stream_response_via_chat( self, @@ -686,11 +713,21 @@ async def _stream_response_via_chat( final_response: Response | None = None yielded_terminal_event = False close_stream_in_background = False + raw_usage_options: dict[str, Any] = ( + {"preserve_raw_usage": True} if model_settings.preserve_raw_usage is True else {} + ) try: async for chunk in ChatCmplStreamHandler.handle_stream( response, - cast(Any, self._normalize_chat_stream(stream)), + cast( + Any, + self._normalize_chat_stream( + stream, + preserve_raw_usage=model_settings.preserve_raw_usage is True, + ), + ), model=self.model, + **raw_usage_options, ): # Record terminal state and populate the span before yielding so a consumer # that stops at the completed event still leaves a fully recorded span. @@ -899,7 +936,18 @@ async def _fetch_chat_response( ) if not stream: - return self._normalize_chat_completion_response(ret) + raw_usage = ( + _raw_usage_snapshot(_usage_payload(ret)) + if model_settings.preserve_raw_usage is True + else None + ) + normalized_response = self._normalize_chat_completion_response(ret) + if model_settings.preserve_raw_usage is True: + _attach_raw_usage_snapshot( + normalized_response, + raw_usage, + ) + return normalized_response responses_tool_choice = OpenAIResponsesConverter.convert_tool_choice( model_settings.tool_choice @@ -1051,7 +1099,18 @@ async def _fetch_responses_response( if stream: return cast(AsyncIterator[ResponseStreamEvent], response) - return self._normalize_response(response) + raw_usage = ( + _raw_usage_snapshot(_usage_payload(response)) + if model_settings.preserve_raw_usage is True + else None + ) + normalized_response = self._normalize_response(response) + if model_settings.preserve_raw_usage is True: + _attach_raw_usage_snapshot( + normalized_response, + raw_usage, + ) + return normalized_response @staticmethod def _split_model_name(model: str) -> tuple[str, str]: @@ -1183,10 +1242,20 @@ def _normalize_chat_completion_response(self, response: Any) -> ChatCompletion: return ChatCompletion.model_validate(response) async def _normalize_chat_stream( - self, stream: AsyncIterator[ChatCompletionChunk] + self, + stream: AsyncIterator[ChatCompletionChunk], + *, + preserve_raw_usage: bool = False, ) -> AsyncIterator[ChatCompletionChunk]: async for chunk in stream: - yield self._normalize_chat_chunk(chunk) + raw_usage = _raw_usage_snapshot(_usage_payload(chunk)) if preserve_raw_usage else None + normalized_chunk = self._normalize_chat_chunk(chunk) + if preserve_raw_usage: + _attach_raw_usage_snapshot( + normalized_chunk, + raw_usage, + ) + yield normalized_chunk def _normalize_chat_chunk(self, chunk: Any) -> ChatCompletionChunk: normalized_chunk = chunk diff --git a/src/agents/items.py b/src/agents/items.py index 012d81b1dd..e44f483adf 100644 --- a/src/agents/items.py +++ b/src/agents/items.py @@ -678,6 +678,15 @@ class ModelResponse: request_id: str | None = None """The transport request ID for this model call, if provided by the model SDK.""" + raw_usage: dict[str, Any] | None = None + """A JSON-compatible snapshot of the provider usage payload, when preservation is enabled. + + The snapshot is captured only while the unnormalized provider payload is available, before the + Agents SDK normalizes missing usage fields. It is ``None`` when preservation is disabled, no + usage payload reaches the model adapter, or upstream normalization has already discarded + field-presence information. + """ + def to_input_items(self) -> list[TResponseInputItem]: """Convert the output into a list of input items suitable for passing to the model.""" # Most output items can be replayed via a direct model_dump. Tool-search items carry diff --git a/src/agents/model_settings.py b/src/agents/model_settings.py index 93fa4112d5..d9db8daefa 100644 --- a/src/agents/model_settings.py +++ b/src/agents/model_settings.py @@ -201,6 +201,16 @@ class ModelSettings: control which prompt prefixes are eligible for caching. """ + preserve_raw_usage: bool | None = None + """Whether to preserve the provider usage payload on completed model responses. + + When enabled and the model adapter still has the unnormalized provider payload, + ``ModelResponse.raw_usage`` contains a JSON-compatible snapshot captured before the Agents + SDK normalizes missing usage fields. It remains ``None`` when usage is absent or upstream + normalization has already discarded field-presence information. This setting does not request + usage from the provider; use ``include_usage`` separately when a streaming provider requires it. + """ + if TYPE_CHECKING: def __init__( @@ -228,6 +238,7 @@ def __init__( retry: ModelRetrySettings | dict[str, Any] | None = None, context_management: list[ContextManagement] | None = None, prompt_cache_options: PromptCacheOptions | None = None, + preserve_raw_usage: bool | None = None, ) -> None: ... def resolve(self, override: ModelSettings | dict[str, Any] | None) -> ModelSettings: diff --git a/src/agents/models/chatcmpl_stream_handler.py b/src/agents/models/chatcmpl_stream_handler.py index 7187b1219d..0f2cd1f152 100644 --- a/src/agents/models/chatcmpl_stream_handler.py +++ b/src/agents/models/chatcmpl_stream_handler.py @@ -51,7 +51,12 @@ from ..exceptions import ModelBehaviorError, UserError from ..items import TResponseStreamEvent from ..logger import logger -from ..usage import _cache_write_tokens, _make_input_tokens_details +from ..usage import ( + _attach_raw_usage_snapshot, + _cache_write_tokens, + _extract_raw_usage_snapshot, + _make_input_tokens_details, +) from .chatcmpl_helpers import ChatCmplHelpers from .fake_id import FAKE_RESPONSES_ID @@ -584,6 +589,7 @@ async def handle_stream( stream: AsyncStream[ChatCompletionChunk], model: str | None = None, strict_feature_validation: bool = False, + preserve_raw_usage: bool = False, ) -> AsyncIterator[TResponseStreamEvent]: """ Handle a streaming chat completion response and yield response events. @@ -593,8 +599,11 @@ async def handle_stream( stream: The async stream of chat completion chunks from the model model: The source model that is generating this stream. Used to handle provider-specific stream processing. + preserve_raw_usage: Whether to retain the last provider usage payload before + converting it to the Responses usage shape. """ usage: CompletionUsage | None = None + raw_usage: dict[str, Any] | None = None state = StreamingState() output_layout = _StreamOutputLayout() sequence_number = SequenceNumber() @@ -616,6 +625,8 @@ async def handle_stream( # Only update when chunk has usage data (not always in the last chunk) if hasattr(chunk, "usage") and chunk.usage is not None: usage = chunk.usage + if preserve_raw_usage: + raw_usage = _extract_raw_usage_snapshot(chunk, fallback=chunk.usage) if not chunk.choices: continue @@ -1272,6 +1283,8 @@ async def handle_stream( if usage else None ) + if preserve_raw_usage: + _attach_raw_usage_snapshot(final_response, raw_usage) yield ResponseCompletedEvent( response=final_response, diff --git a/src/agents/models/openai_chatcompletions.py b/src/agents/models/openai_chatcompletions.py index 8b6c77a557..22d0427d41 100644 --- a/src/agents/models/openai_chatcompletions.py +++ b/src/agents/models/openai_chatcompletions.py @@ -31,7 +31,7 @@ from ..tracing import generation_span from ..tracing.span_data import GenerationSpanData from ..tracing.spans import Span -from ..usage import Usage +from ..usage import Usage, _raw_usage_snapshot from ..util._error_tracing import model_span_errors from ..util._json import _to_dump_compatible from ._openai_retry import get_openai_retry_advice @@ -351,6 +351,11 @@ async def get_response( # The OpenAI SDK records the `x-request-id` header on every parsed response, # so callers can inspect the same debugging handle as on the Responses path. request_id=getattr(response, "_request_id", None), + raw_usage=( + _raw_usage_snapshot(response.usage) + if model_settings.preserve_raw_usage is True + else None + ), ) @staticmethod @@ -444,6 +449,9 @@ async def stream_response( else: stream_for_handler = stream + raw_usage_options: dict[str, Any] = ( + {"preserve_raw_usage": True} if model_settings.preserve_raw_usage is True else {} + ) close_stream_in_background = False yielded_terminal_event = False try: @@ -452,6 +460,7 @@ async def stream_response( cast(AsyncStream[ChatCompletionChunk], stream_for_handler), model=self.model, strict_feature_validation=self._strict_feature_validation, + **raw_usage_options, ): if chunk.type == "response.completed": final_response = chunk.response diff --git a/src/agents/models/openai_responses.py b/src/agents/models/openai_responses.py index 8ff70ef3bf..95b3e2f426 100644 --- a/src/agents/models/openai_responses.py +++ b/src/agents/models/openai_responses.py @@ -74,7 +74,13 @@ validate_responses_tool_search_configuration, ) from ..tracing import SpanError, response_span -from ..usage import Usage, _response_usage_to_usage, model_usage_to_span_usage +from ..usage import ( + Usage, + _attach_raw_usage_snapshot, + _raw_usage_snapshot, + _response_usage_to_usage, + model_usage_to_span_usage, +) from ..util._error_tracing import record_model_error_on_span from ..util._json import _to_dump_compatible from ..version import __version__ @@ -550,6 +556,11 @@ async def get_response( usage=usage, response_id=response.id, request_id=getattr(response, "_request_id", None), + raw_usage=( + _raw_usage_snapshot(response.usage) + if model_settings.preserve_raw_usage is True + else None + ), ) async def stream_response( @@ -592,6 +603,8 @@ async def stream_response( chunk_type = getattr(chunk, "type", None) if isinstance(chunk, ResponseCompletedEvent): final_response = chunk.response + if model_settings.preserve_raw_usage is True: + _attach_raw_usage_snapshot(chunk.response, chunk.response.usage) elif chunk_type in { "response.failed", "response.incomplete", diff --git a/src/agents/run_internal/run_loop.py b/src/agents/run_internal/run_loop.py index 1429270c2e..73ff934700 100644 --- a/src/agents/run_internal/run_loop.py +++ b/src/agents/run_internal/run_loop.py @@ -98,7 +98,7 @@ from ..tracing.config import include_task_and_turn_spans from ..tracing.model_tracing import get_model_tracing_impl from ..tracing.span_data import AgentSpanData, TaskSpanData -from ..usage import Usage, _response_usage_to_usage +from ..usage import Usage, _extract_raw_usage_snapshot, _response_usage_to_usage from ..util import _coro, _error_tracing from ..util._asyncio_tasks import gather_with_cancel from .agent_bindings import AgentBindings, bind_public_agent @@ -1774,6 +1774,11 @@ async def rewind_model_request() -> None: usage=usage, response_id=terminal_response.id, request_id=getattr(terminal_response, "_request_id", None), + raw_usage=( + _extract_raw_usage_snapshot(terminal_response) + if model_settings.preserve_raw_usage is True + else None + ), ) if isinstance(event, ResponseOutputItemDoneEvent): diff --git a/src/agents/usage.py b/src/agents/usage.py index 4880ecdc62..5d5a0c479e 100644 --- a/src/agents/usage.py +++ b/src/agents/usage.py @@ -1,14 +1,70 @@ from __future__ import annotations +import json from collections.abc import Mapping from dataclasses import field -from typing import Annotated, Any +from typing import Annotated, Any, cast from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails -from pydantic import BeforeValidator, TypeAdapter, ValidationError +from pydantic import BeforeValidator, JsonValue, TypeAdapter, ValidationError from pydantic.dataclasses import dataclass +_RAW_USAGE_ATTRIBUTE = "_agents_sdk_raw_usage" +_RAW_USAGE_ADAPTER = TypeAdapter(dict[str, JsonValue]) +_RAW_USAGE_MISSING = object() + + +def _raw_usage_snapshot(raw_usage: Any | None) -> dict[str, Any] | None: + """Return a detached JSON-compatible usage object without adding omitted fields.""" + if raw_usage is None: + return None + + try: + if isinstance(raw_usage, Mapping): + candidate = dict(raw_usage) + else: + model_dump = getattr(raw_usage, "model_dump", None) + if not callable(model_dump): + return None + candidate = model_dump(mode="json", by_alias=True, exclude_unset=True) + + if not isinstance(candidate, dict) or not all(isinstance(key, str) for key in candidate): + return None + + validated = _RAW_USAGE_ADAPTER.validate_python(candidate) + return cast( + dict[str, Any], + json.loads(json.dumps(validated, allow_nan=False)), + ) + except Exception: + # Usage preservation is diagnostic metadata. An adapter-specific value that cannot be + # represented as JSON must not turn an otherwise successful model call into a failure. + return None + + +def _attach_raw_usage_snapshot(target: Any, raw_usage: Any | None) -> None: + """Attach a pre-normalization usage snapshot to an internal response object.""" + snapshot = _raw_usage_snapshot(raw_usage) + try: + object.__setattr__(target, _RAW_USAGE_ATTRIBUTE, snapshot) + except Exception: + # Some custom response objects reject private attributes. Their completed response can + # still be processed normally, but no raw usage snapshot is available downstream. + return + + +def _extract_raw_usage_snapshot( + target: Any, + *, + fallback: Any | None = None, +) -> dict[str, Any] | None: + """Read an attached snapshot, or capture the provided unnormalized fallback.""" + snapshot = getattr(target, _RAW_USAGE_ATTRIBUTE, _RAW_USAGE_MISSING) + if snapshot is not _RAW_USAGE_MISSING: + return snapshot if isinstance(snapshot, dict) else None + return _raw_usage_snapshot(fallback) + def _make_input_tokens_details( *, diff --git a/tests/model_settings/test_serialization.py b/tests/model_settings/test_serialization.py index 073801bd11..d458a8a7b3 100644 --- a/tests/model_settings/test_serialization.py +++ b/tests/model_settings/test_serialization.py @@ -126,6 +126,7 @@ def test_all_fields_serialization() -> None: ), context_management=[{"type": "compaction", "compact_threshold": 200000}], prompt_cache_options={"mode": "explicit", "ttl": "30m"}, + preserve_raw_usage=True, ) # Verify that every single field is set to a non-None value @@ -154,10 +155,14 @@ def test_gpt_5_6_reasoning_and_prompt_cache_serialization() -> None: } -def test_prompt_cache_options_is_appended_to_public_field_order() -> None: +def test_usage_preservation_is_appended_to_public_field_order() -> None: field_names = [field.name for field in fields(ModelSettings)] - assert field_names[-2:] == ["context_management", "prompt_cache_options"] + assert field_names[-3:] == [ + "context_management", + "prompt_cache_options", + "preserve_raw_usage", + ] def test_extra_args_serialization() -> None: @@ -185,6 +190,7 @@ def test_traceable_serialization_omits_request_extras() -> None: extra_query={"api-key": "query-token"}, extra_body={"secret": "body-token"}, extra_args={"api_key": "arg-token"}, + preserve_raw_usage=True, ) json_dict = model_settings.to_json_dict() @@ -199,6 +205,7 @@ def test_traceable_serialization_omits_request_extras() -> None: assert "extra_query" not in traceable assert "extra_body" not in traceable assert "extra_args" not in traceable + assert "preserve_raw_usage" not in traceable def test_extra_args_resolve() -> None: diff --git a/tests/models/test_any_llm_model.py b/tests/models/test_any_llm_model.py index b7788da2a2..9808845898 100644 --- a/tests/models/test_any_llm_model.py +++ b/tests/models/test_any_llm_model.py @@ -421,7 +421,7 @@ async def test_any_llm_chat_path_is_used_when_responses_are_unsupported(monkeypa response = await model.get_response( system_instructions="You are terse.", input="hi", - model_settings=ModelSettings(), + model_settings=ModelSettings(preserve_raw_usage=True), tools=[], output_schema=None, handoffs=[], @@ -445,6 +445,11 @@ async def test_any_llm_chat_path_is_used_when_responses_are_unsupported(monkeypa assert response.output[0].content[0].text == "Hello" assert response.usage.input_tokens_details.cached_tokens == 2 assert getattr(response.usage.input_tokens_details, "cache_write_tokens", None) == 4 + assert response.raw_usage is not None + assert response.raw_usage["prompt_tokens_details"] == { + "cached_tokens": 2, + "cache_write_tokens": 4, + } def _content_filtered_chat_completion(content: str) -> ChatCompletion: @@ -655,7 +660,7 @@ async def test_any_llm_responses_path_defaults_missing_cache_write_tokens( normalized = await model.get_response( system_instructions=None, input="hi", - model_settings=ModelSettings(), + model_settings=ModelSettings(preserve_raw_usage=True), tools=[], output_schema=None, handoffs=[], @@ -668,6 +673,8 @@ async def test_any_llm_responses_path_defaults_missing_cache_write_tokens( assert normalized.output[0].content[0].text == "Hello" assert normalized.usage.input_tokens_details.cache_write_tokens == 0 assert "cache_write_tokens" not in response_payload["usage"]["input_tokens_details"] + assert normalized.raw_usage is not None + assert "cache_write_tokens" not in normalized.raw_usage["input_tokens_details"] @pytest.mark.allow_call_model_methods diff --git a/tests/models/test_litellm_chatcompletions_stream.py b/tests/models/test_litellm_chatcompletions_stream.py index 8bc69eb1e1..ece71165de 100644 --- a/tests/models/test_litellm_chatcompletions_stream.py +++ b/tests/models/test_litellm_chatcompletions_stream.py @@ -91,7 +91,7 @@ async def patched_fetch_response(self, *args, **kwargs): async for event in model.stream_response( system_instructions=None, input="", - model_settings=ModelSettings(), + model_settings=ModelSettings(preserve_raw_usage=True), tools=[], output_schema=None, handoffs=[], @@ -133,6 +133,9 @@ async def patched_fetch_response(self, *args, **kwargs): assert completed_resp.usage.total_tokens == 12 assert completed_resp.usage.input_tokens_details.cached_tokens == 6 assert completed_resp.usage.output_tokens_details.reasoning_tokens == 2 + # LiteLLM has already normalized usage before the Agents adapter receives this chunk, so + # omitted-versus-null provenance is unavailable and no raw snapshot should be attached. + assert not hasattr(completed_resp, "_agents_sdk_raw_usage") @pytest.mark.allow_call_model_methods diff --git a/tests/models/test_litellm_logprobs.py b/tests/models/test_litellm_logprobs.py index 00354ab57e..1cb247a994 100644 --- a/tests/models/test_litellm_logprobs.py +++ b/tests/models/test_litellm_logprobs.py @@ -94,7 +94,7 @@ async def fake_acompletion(model, messages=None, **kwargs): response = await LitellmModel(model="test-model").get_response( system_instructions=None, input=[], - model_settings=ModelSettings(top_logprobs=2), + model_settings=ModelSettings(top_logprobs=2, preserve_raw_usage=True), tools=[], output_schema=None, handoffs=[], @@ -116,3 +116,6 @@ async def fake_acompletion(model, messages=None, **kwargs): assert output_logprobs[0].token == "Hello" assert output_logprobs[0].logprob == -0.25 assert [tlp.token for tlp in output_logprobs[0].top_logprobs] == ["Hello", "Hi"] + # LiteLLM has already normalized usage before the Agents adapter receives this response, so + # omitted-versus-null provenance is unavailable. + assert response.raw_usage is None diff --git a/tests/models/test_openai_chatcompletions.py b/tests/models/test_openai_chatcompletions.py index 9364025275..e756ec020d 100644 --- a/tests/models/test_openai_chatcompletions.py +++ b/tests/models/test_openai_chatcompletions.py @@ -173,6 +173,87 @@ async def patched_fetch_response(self, *args, **kwargs): assert getattr(resp.usage.input_tokens_details, "cache_write_tokens", None) == 4 assert resp.usage.output_tokens_details.reasoning_tokens == 0 assert resp.response_id is None + assert resp.raw_usage is None + + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("prompt_tokens_details", "expected_details"), + [({}, {}), ({"cached_tokens": 0}, {"cached_tokens": 0})], + ids=["omitted-cached-tokens", "explicit-zero-cached-tokens"], +) +async def test_get_response_preserves_raw_usage_field_presence( + monkeypatch: pytest.MonkeyPatch, + prompt_tokens_details: dict[str, int], + expected_details: dict[str, int], +) -> None: + chat = _minimal_chat_completion() + chat.usage = CompletionUsage.model_validate( + { + "completion_tokens": 5, + "prompt_tokens": 7, + "total_tokens": 12, + "prompt_tokens_details": prompt_tokens_details, + } + ) + + async def patched_fetch_response(self, *args, **kwargs): + return chat + + monkeypatch.setattr(OpenAIChatCompletionsModel, "_fetch_response", patched_fetch_response) + model = OpenAIProvider(use_responses=False).get_model("gpt-4") + + response = await model.get_response( + system_instructions=None, + input="", + model_settings=ModelSettings(preserve_raw_usage=True), + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + previous_response_id=None, + conversation_id=None, + prompt=None, + ) + + assert response.raw_usage == { + "completion_tokens": 5, + "prompt_tokens": 7, + "total_tokens": 12, + "prompt_tokens_details": expected_details, + } + assert response.usage.input_tokens_details.cached_tokens == 0 + + +@pytest.mark.allow_call_model_methods +@pytest.mark.asyncio +async def test_get_response_raw_usage_is_none_when_provider_omits_usage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + chat = _minimal_chat_completion() + + async def patched_fetch_response(self, *args, **kwargs): + return chat + + monkeypatch.setattr(OpenAIChatCompletionsModel, "_fetch_response", patched_fetch_response) + model = OpenAIProvider(use_responses=False).get_model("gpt-4") + + response = await model.get_response( + system_instructions=None, + input="", + model_settings=ModelSettings(preserve_raw_usage=True), + tools=[], + output_schema=None, + handoffs=[], + tracing=ModelTracing.DISABLED, + previous_response_id=None, + conversation_id=None, + prompt=None, + ) + + assert response.raw_usage is None + assert response.usage.total_tokens == 0 async def _get_response_for_choice( diff --git a/tests/models/test_openai_chatcompletions_stream.py b/tests/models/test_openai_chatcompletions_stream.py index fcb8bc74ab..014dbe2ac3 100644 --- a/tests/models/test_openai_chatcompletions_stream.py +++ b/tests/models/test_openai_chatcompletions_stream.py @@ -568,7 +568,14 @@ async def test_stream_handler_keeps_empty_choice_usage_chunks() -> None: model="fake", object="chat.completion.chunk", choices=[], - usage=CompletionUsage(completion_tokens=1, prompt_tokens=2, total_tokens=3), + usage=CompletionUsage.model_validate( + { + "completion_tokens": 1, + "prompt_tokens": 2, + "total_tokens": 3, + "prompt_tokens_details": {"cached_tokens": 0}, + } + ), ) async def fake_stream() -> AsyncIterator[ChatCompletionChunk]: @@ -577,7 +584,7 @@ async def fake_stream() -> AsyncIterator[ChatCompletionChunk]: events = [ event async for event in ChatCmplStreamHandler.handle_stream( - _empty_response(), cast(Any, fake_stream()) + _empty_response(), cast(Any, fake_stream()), preserve_raw_usage=True ) ] @@ -587,6 +594,12 @@ async def fake_stream() -> AsyncIterator[ChatCompletionChunk]: assert completed_event.response.output == [] assert completed_event.response.usage assert completed_event.response.usage.total_tokens == 3 + assert cast(Any, completed_event.response)._agents_sdk_raw_usage == { + "completion_tokens": 1, + "prompt_tokens": 2, + "total_tokens": 3, + "prompt_tokens_details": {"cached_tokens": 0}, + } @pytest.mark.asyncio diff --git a/tests/models/test_openai_responses.py b/tests/models/test_openai_responses.py index c5978b2ba9..34bcc7cb5e 100644 --- a/tests/models/test_openai_responses.py +++ b/tests/models/test_openai_responses.py @@ -252,7 +252,7 @@ def __init__(self): response = await model.get_response( system_instructions=None, input="hi", - model_settings=ModelSettings(), + model_settings=ModelSettings(preserve_raw_usage=True), tools=[], output_schema=None, handoffs=[], @@ -261,6 +261,8 @@ def __init__(self): assert response.response_id == "resp-request-id" assert response.request_id == "req_nonstream_123" + assert response.raw_usage is not None + assert response.raw_usage["input_tokens_details"]["cached_tokens"] == 0 @pytest.mark.allow_call_model_methods diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py index 2a3c605817..02cbf76ee3 100644 --- a/tests/test_agent_runner_streamed.py +++ b/tests/test_agent_runner_streamed.py @@ -59,7 +59,7 @@ from agents.stream_events import AgentUpdatedStreamEvent, RawResponsesStreamEvent, StreamEvent from agents.tool import FunctionTool, Tool from agents.tool_guardrails import tool_input_guardrail, tool_output_guardrail -from agents.usage import Usage +from agents.usage import Usage, _attach_raw_usage_snapshot from .fake_model import FakeModel, get_response_obj from .test_responses import ( @@ -317,7 +317,10 @@ async def stream_response( @pytest.mark.asyncio -async def test_streamed_run_exposes_request_id_on_raw_responses() -> None: +@pytest.mark.parametrize("preserve_raw_usage", [None, False, True]) +async def test_streamed_run_exposes_request_id_on_raw_responses( + preserve_raw_usage: bool | None, +) -> None: class RequestIdTerminalFakeModel(FakeModel): async def stream_response( self, @@ -337,6 +340,10 @@ async def stream_response( [get_text_message("partial final")], response_id="resp-partial" ) response._request_id = "req_streamed_result_123" + _attach_raw_usage_snapshot( + response, + {"input_tokens": 3, "input_tokens_details": {"cached_tokens": 0}}, + ) yield ResponseCompletedEvent( type="response.completed", response=response, @@ -344,7 +351,11 @@ async def stream_response( ) model = RequestIdTerminalFakeModel() - agent = Agent(name="test", model=model) + agent = Agent( + name="test", + model=model, + model_settings=ModelSettings(preserve_raw_usage=preserve_raw_usage), + ) result = Runner.run_streamed(agent, input="test") async for _ in result.stream_events(): @@ -352,6 +363,14 @@ async def stream_response( assert len(result.raw_responses) == 1 assert result.raw_responses[0].request_id == "req_streamed_result_123" + assert result.raw_responses[0].raw_usage == ( + { + "input_tokens": 3, + "input_tokens_details": {"cached_tokens": 0}, + } + if preserve_raw_usage is True + else None + ) @pytest.mark.asyncio diff --git a/tests/test_run_state.py b/tests/test_run_state.py index 83bbdb7c1f..78cc09b11b 100644 --- a/tests/test_run_state.py +++ b/tests/test_run_state.py @@ -2668,16 +2668,20 @@ async def test_model_response_serialization_roundtrip(self): ], response_id="resp123", request_id="req123", + raw_usage={"input_tokens": 10, "provider_metric": 0}, ) state._model_responses.append(response) # Round trip + serialized = state.to_json() + assert "raw_usage" not in serialized["model_responses"][0] json_str = state.to_string() restored = await RunState.from_string(agent, json_str) assert len(restored._model_responses) == 1 assert restored._model_responses[0].response_id == "resp123" assert restored._model_responses[0].request_id == "req123" + assert restored._model_responses[0].raw_usage is None assert restored._model_responses[0].usage.requests == 1 assert restored._model_responses[0].usage.input_tokens == 10 diff --git a/tests/test_usage.py b/tests/test_usage.py index 254a88cbea..c6c8444ef0 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -1,7 +1,13 @@ from __future__ import annotations +from typing import Any + import pytest -from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails +from openai.types.completion_usage import ( + CompletionTokensDetails, + CompletionUsage, + PromptTokensDetails, +) from openai.types.responses.response_usage import InputTokensDetails, OutputTokensDetails from agents import Agent, Runner @@ -9,6 +15,7 @@ from agents.usage import ( RequestUsage, Usage, + _raw_usage_snapshot, deserialize_usage, model_usage_to_span_usage, serialize_usage, @@ -24,6 +31,45 @@ def test_usage_defaults_cache_write_tokens_to_zero() -> None: assert getattr(usage.input_tokens_details, "cache_write_tokens", None) == 0 +def test_raw_usage_snapshot_preserves_presence_and_is_detached() -> None: + raw_usage: dict[str, Any] = { + "input_tokens": 3, + "input_tokens_details": {"cached_tokens": 0}, + "provider_metric": None, + } + + snapshot = _raw_usage_snapshot(raw_usage) + raw_usage["input_tokens_details"]["cached_tokens"] = 9 + + assert snapshot == { + "input_tokens": 3, + "input_tokens_details": {"cached_tokens": 0}, + "provider_metric": None, + } + + +def test_raw_usage_snapshot_does_not_add_unset_pydantic_fields() -> None: + usage = CompletionUsage.model_validate( + { + "completion_tokens": 2, + "prompt_tokens": 3, + "total_tokens": 5, + "prompt_tokens_details": {}, + } + ) + + assert _raw_usage_snapshot(usage) == { + "completion_tokens": 2, + "prompt_tokens": 3, + "total_tokens": 5, + "prompt_tokens_details": {}, + } + + +def test_raw_usage_snapshot_rejects_non_json_values() -> None: + assert _raw_usage_snapshot({"provider_metric": object()}) is None + + @pytest.mark.asyncio async def test_runner_run_carries_request_usage_entries() -> None: """Ensure usage produced by the model propagates to RunResult context."""