From e7e932e8c03bbc236ba7dfbf3a76708dc21dbe1b 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 a3f06648fce..a6049701290 100644 --- a/src/google/adk/models/anthropic_llm.py +++ b/src/google/adk/models/anthropic_llm.py @@ -710,6 +710,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), ) @@ -1145,6 +1146,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), partial=False, ) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 74a1ecd68f6..9c939508d6b 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -2409,6 +2409,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: @@ -3133,6 +3134,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 fallback_index = 0 @@ -3294,6 +3296,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, @@ -3344,6 +3347,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 @@ -3352,6 +3358,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 8ecdbe290e9..0c575e9d158 100644 --- a/src/google/adk/models/llm_response.py +++ b/src/google/adk/models/llm_response.py @@ -112,6 +112,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 1f25a217006..10ccbdf8cee 100644 --- a/tests/unittests/models/test_anthropic_llm.py +++ b/tests/unittests/models/test_anthropic_llm.py @@ -2010,6 +2010,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 @@ -2041,6 +2051,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 @@ -2097,6 +2108,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 0b89d04f32f..d36b9195dbf 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -1882,6 +1882,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() @@ -4597,6 +4598,7 @@ async def test_generate_content_async_stream_with_bedrock_cache_tokens( assert responses[3].usage_metadata.total_token_count == 15 assert responses[3].usage_metadata.cached_content_token_count == 8 assert responses[3].usage_metadata.cache_creation_input_tokens == 4 + assert responses[3].cache_creation_token_count == 4 @pytest.mark.asyncio