From d8e060ba4a004eb12cd21ae495e90d6aaa339efd Mon Sep 17 00:00:00 2001 From: Gurjot Singh Date: Thu, 20 Aug 2026 01:20:15 +0530 Subject: [PATCH] feat: serialize prompt-cache creation tokens on LlmResponse Commit d0b33a05 extracts cache-creation (write) tokens and attaches them to the GenerateContentResponseUsageMetadata object via object.__setattr__ so they reach OpenTelemetry. But google.genai's usage type forbids extra fields, so that attribute is dropped by model_dump() and never appears in serialized events -- the dev UI, persisted sessions, and SSE consumers cannot see cache-write tokens. Surface the same count on a serializable LlmResponse.cache_creation_token_count field (inherited by Event, emitted as cacheCreationTokenCount, omitted when None). It is the cache-write counterpart to usage_metadata.cached_content_token_count. Populated from the already-extracted value in the Anthropic and LiteLLM paths (streaming and non-streaming); the OpenTelemetry path is left unchanged. --- src/google/adk/models/anthropic_llm.py | 2 ++ src/google/adk/models/lite_llm.py | 9 +++++++++ src/google/adk/models/llm_response.py | 16 ++++++++++++++++ tests/unittests/models/test_anthropic_llm.py | 12 ++++++++++++ tests/unittests/models/test_litellm.py | 2 ++ 5 files changed, 41 insertions(+) diff --git a/src/google/adk/models/anthropic_llm.py b/src/google/adk/models/anthropic_llm.py index 7f0331e33fd..15725d435d5 100644 --- a/src/google/adk/models/anthropic_llm.py +++ b/src/google/adk/models/anthropic_llm.py @@ -804,6 +804,7 @@ def message_to_generate_content_response( parts=parts, ), usage_metadata=usage_metadata, + cache_creation_token_count=cache_creation, finish_reason=to_google_genai_finish_reason(message.stop_reason), model_version=message.model, ) @@ -1359,6 +1360,7 @@ async def _generate_content_streaming( yield LlmResponse( content=types.Content(role="model", parts=all_parts), usage_metadata=usage_metadata, + cache_creation_token_count=cache_creation_tokens, finish_reason=to_google_genai_finish_reason(stop_reason), model_version=model_version, partial=False, diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 1630c3c983e..a98d4406d19 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -2652,6 +2652,7 @@ def _model_response_to_generate_content_response( "cache_creation_input_tokens", cache_creation, ) + llm_response.cache_creation_token_count = cache_creation grounding_metadata = _extract_grounding_metadata(response) if grounding_metadata: @@ -3571,6 +3572,7 @@ async def generate_content_async( aggregated_llm_response = None aggregated_llm_response_with_tool_call = None usage_metadata = None + cache_creation_token_count = None grounding_metadata = None last_finish_reason: str | None = None last_model_version: str | None = None @@ -3795,6 +3797,7 @@ def _reset_stream_buffers() -> None: "cache_creation_input_tokens", chunk.cache_creation_tokens, ) + cache_creation_token_count = chunk.cache_creation_tokens # LiteLLM 1.81+ can set finish_reason="stop" on partial chunks. Only # finalize tool calls on an explicit tool_calls/length finish_reason, @@ -3865,6 +3868,9 @@ def _reset_stream_buffers() -> None: if aggregated_llm_response: if usage_metadata: aggregated_llm_response.usage_metadata = usage_metadata + aggregated_llm_response.cache_creation_token_count = ( + cache_creation_token_count + ) usage_metadata = None if grounding_metadata: aggregated_llm_response.grounding_metadata = grounding_metadata @@ -3873,6 +3879,9 @@ def _reset_stream_buffers() -> None: if aggregated_llm_response_with_tool_call: if usage_metadata: aggregated_llm_response_with_tool_call.usage_metadata = usage_metadata + aggregated_llm_response_with_tool_call.cache_creation_token_count = ( + cache_creation_token_count + ) if grounding_metadata: aggregated_llm_response_with_tool_call.grounding_metadata = ( grounding_metadata diff --git a/src/google/adk/models/llm_response.py b/src/google/adk/models/llm_response.py index 7bf0c2f6596..0ed65801c13 100644 --- a/src/google/adk/models/llm_response.py +++ b/src/google/adk/models/llm_response.py @@ -139,6 +139,22 @@ class LlmResponse(BaseModel): usage_metadata: Optional[types.GenerateContentResponseUsageMetadata] = None """The usage metadata of the LlmResponse""" + cache_creation_token_count: Optional[int] = None + """Number of prompt tokens written to the prompt cache for this response. + + The cache-*write* counterpart to + ``usage_metadata.cached_content_token_count`` (cache-*read*); both are a + breakdown of ``usage_metadata.prompt_token_count``. + + The model backends also attach this value to ``usage_metadata`` as + ``cache_creation_input_tokens`` for OpenTelemetry, but ``google.genai``'s + usage type forbids extra fields, so that attribute is dropped on + serialization. This field surfaces the same count on the (serializable) + ``LlmResponse`` / ``Event`` so it is visible to the dev UI, persisted + sessions, and other event consumers -- not only telemetry. Only populated + for providers/turns that report cache writes; ``None`` otherwise. + """ + live_session_resumption_update: Optional[ types.LiveServerSessionResumptionUpdate ] = None diff --git a/tests/unittests/models/test_anthropic_llm.py b/tests/unittests/models/test_anthropic_llm.py index 0a9dccd546e..2784871f385 100644 --- a/tests/unittests/models/test_anthropic_llm.py +++ b/tests/unittests/models/test_anthropic_llm.py @@ -2303,6 +2303,16 @@ def test_message_to_generate_content_response_reports_cache_creation_tokens(): response = message_to_generate_content_response(message) assert response.usage_metadata.cache_creation_input_tokens == 50 + # Serializable mirror: the usage_metadata attribute above is dropped by + # model_dump (google.genai forbids extra fields), so the count is also + # exposed on LlmResponse.cache_creation_token_count for event consumers. + assert response.cache_creation_token_count == 50 + assert ( + response.model_dump(by_alias=True, exclude_none=True)[ + "cacheCreationTokenCount" + ] + == 50 + ) dumped = response.model_dump() assert "usage_metadata" in dumped @@ -2334,6 +2344,7 @@ def test_message_to_generate_content_response_no_cache_creation_tokens(): response = message_to_generate_content_response(message) assert not hasattr(response.usage_metadata, "cache_creation_input_tokens") + assert response.cache_creation_token_count is None @pytest.mark.asyncio @@ -2391,6 +2402,7 @@ async def test_streaming_reports_cache_creation_tokens(): assert len(responses) == 2 final_response = responses[-1] assert final_response.usage_metadata.cache_creation_input_tokens == 50 + assert final_response.cache_creation_token_count == 50 dumped = final_response.model_dump() assert "usage_metadata" in dumped diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index caf2ceeb97e..be2018dd3e6 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -2606,6 +2606,7 @@ async def test_generate_content_async_with_bedrock_cache_tokens( assert response.usage_metadata.total_token_count == 15 assert response.usage_metadata.cached_content_token_count == 8 assert response.usage_metadata.cache_creation_input_tokens == 4 + assert response.cache_creation_token_count == 4 mock_acompletion.assert_called_once() @@ -5844,6 +5845,7 @@ async def test_generate_content_async_stream_with_bedrock_cache_tokens( assert responses[5].usage_metadata.total_token_count == 15 assert responses[5].usage_metadata.cached_content_token_count == 8 assert responses[5].usage_metadata.cache_creation_input_tokens == 4 + assert responses[5].cache_creation_token_count == 4 @pytest.mark.asyncio