diff --git a/src/google/adk/models/anthropic_llm.py b/src/google/adk/models/anthropic_llm.py index 9e1118f5..a3f06648 100644 --- a/src/google/adk/models/anthropic_llm.py +++ b/src/google/adk/models/anthropic_llm.py @@ -670,6 +670,12 @@ def _extract_thinking_token_count( return min(thinking, output_tokens) +def _extract_cache_creation_token_count(usage: Any) -> int | None: + """Returns Anthropic cache-write tokens, the analog of cache_creation tokens.""" + cached = getattr(usage, "cache_creation_input_tokens", None) + return cached if isinstance(cached, int) else None + + def message_to_generate_content_response( message: anthropic_types.Message, ) -> LlmResponse: @@ -683,21 +689,27 @@ def message_to_generate_content_response( prompt_tokens = _extract_prompt_token_count(message.usage) thinking_tokens = _extract_thinking_token_count(message.usage) + usage_metadata = types.GenerateContentResponseUsageMetadata( + prompt_token_count=prompt_tokens, + candidates_token_count=( + message.usage.output_tokens - (thinking_tokens or 0) + ), + total_token_count=prompt_tokens + message.usage.output_tokens, + cached_content_token_count=_extract_cached_token_count(message.usage), + thoughts_token_count=thinking_tokens, + ) + cache_creation = _extract_cache_creation_token_count(message.usage) + if cache_creation is not None: + object.__setattr__( + usage_metadata, "cache_creation_input_tokens", cache_creation + ) return LlmResponse( content=types.Content( role="model", parts=parts, ), - usage_metadata=types.GenerateContentResponseUsageMetadata( - prompt_token_count=prompt_tokens, - candidates_token_count=( - message.usage.output_tokens - (thinking_tokens or 0) - ), - total_token_count=prompt_tokens + message.usage.output_tokens, - cached_content_token_count=_extract_cached_token_count(message.usage), - thoughts_token_count=thinking_tokens, - ), + usage_metadata=usage_metadata, finish_reason=to_google_genai_finish_reason(message.stop_reason), ) @@ -993,6 +1005,7 @@ class AnthropicLlm(BaseLlm): output_tokens = 0 thinking_tokens: int | None = None cached_input_tokens: int | None = None + cache_creation_tokens: int | None = None stop_reason: Optional[anthropic_types.StopReason] = None async for event in raw_stream: @@ -1001,6 +1014,9 @@ class AnthropicLlm(BaseLlm): output_tokens = event.message.usage.output_tokens thinking_tokens = _extract_thinking_token_count(event.message.usage) cached_input_tokens = _extract_cached_token_count(event.message.usage) + cache_creation_tokens = _extract_cache_creation_token_count( + event.message.usage + ) elif event.type == "content_block_start": block = event.content_block @@ -1114,15 +1130,21 @@ class AnthropicLlm(BaseLlm): function_call.id = tool_acc.id all_parts.append(part) + usage_metadata = types.GenerateContentResponseUsageMetadata( + prompt_token_count=input_tokens, + candidates_token_count=output_tokens - (thinking_tokens or 0), + total_token_count=input_tokens + output_tokens, + cached_content_token_count=cached_input_tokens, + thoughts_token_count=thinking_tokens, + ) + if cache_creation_tokens is not None: + object.__setattr__( + usage_metadata, "cache_creation_input_tokens", cache_creation_tokens + ) + yield LlmResponse( content=types.Content(role="model", parts=all_parts), - usage_metadata=types.GenerateContentResponseUsageMetadata( - prompt_token_count=input_tokens, - candidates_token_count=output_tokens - (thinking_tokens or 0), - total_token_count=input_tokens + output_tokens, - cached_content_token_count=cached_input_tokens, - thoughts_token_count=thinking_tokens, - ), + usage_metadata=usage_metadata, 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 bfe31023..74a1ecd6 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -856,6 +856,7 @@ class UsageMetadataChunk(BaseModel): total_tokens: int cached_prompt_tokens: int = 0 reasoning_tokens: int = 0 + cache_creation_tokens: Optional[int] = None class LiteLLMClient: @@ -1027,7 +1028,11 @@ def _extract_cached_prompt_tokens(usage: Any) -> int: if total > 0: return total - for key in ("cached_prompt_tokens", "cached_tokens"): + for key in ( + "cached_prompt_tokens", + "cached_tokens", + "cache_read_input_tokens", + ): value = usage_dict.get(key) if isinstance(value, int): return value @@ -1037,6 +1042,39 @@ def _extract_cached_prompt_tokens(usage: Any) -> int: return 0 +def _extract_cache_creation_tokens(usage: Any) -> Optional[int]: + """Extracts cache creation (write) tokens from LiteLLM usage. + + Args: + usage: Usage dictionary from LiteLLM response. + + Returns: + Integer number of cache creation tokens if present; otherwise None. + """ + try: + usage_dict = usage + if hasattr(usage, "model_dump"): + usage_dict = usage.model_dump() + elif isinstance(usage, str): + try: + usage_dict = json.loads(usage) + except json.JSONDecodeError: + return None + + if not isinstance(usage_dict, dict): + return None + + for key in ("cache_creation_input_tokens", "cache_write_input_tokens"): + if key in usage_dict: + value = usage_dict.get(key) + if isinstance(value, int): + return value + except (TypeError, AttributeError) as e: + logger.debug("Error extracting cache creation tokens: %s", e) + + return None + + def _decode_thought_signature(value: Any) -> Optional[bytes]: """Safely decodes a thought_signature value to bytes. @@ -2268,6 +2306,7 @@ def _model_response_to_chunk( total_tokens=usage.get("total_tokens", 0) or 0, cached_prompt_tokens=_extract_cached_prompt_tokens(usage), reasoning_tokens=_extract_reasoning_tokens(usage), + cache_creation_tokens=_extract_cache_creation_tokens(usage), ), None except AttributeError as e: raise TypeError( @@ -2363,6 +2402,13 @@ def _model_response_to_generate_content_response( cached_content_token_count=_extract_cached_prompt_tokens(usage_dict), thoughts_token_count=reasoning_tokens if reasoning_tokens else None, ) + cache_creation = _extract_cache_creation_tokens(usage_dict) + if cache_creation is not None: + object.__setattr__( + llm_response.usage_metadata, + "cache_creation_input_tokens", + cache_creation, + ) grounding_metadata = _extract_grounding_metadata(response) if grounding_metadata: @@ -3242,6 +3288,12 @@ class LiteLlm(BaseLlm): if chunk.reasoning_tokens else None, ) + if chunk.cache_creation_tokens is not None: + object.__setattr__( + usage_metadata, + "cache_creation_input_tokens", + 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, diff --git a/src/google/adk/telemetry/_token_usage.py b/src/google/adk/telemetry/_token_usage.py index 0ab1788e..4f437aad 100644 --- a/src/google/adk/telemetry/_token_usage.py +++ b/src/google/adk/telemetry/_token_usage.py @@ -29,6 +29,12 @@ from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_A # from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS = 'gen_ai.usage.cache_read.input_tokens' +# Use the import symbol once the minimum OpenTelemetry SDK version is updated to 1.41.0 +# from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS +GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS = ( + 'gen_ai.usage.cache_creation.input_tokens' +) + # Use the import symbol once the minimum OpenTelemetry SDK version is updated to 1.42.0 # from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_USAGE_REASONING_OUTPUT_TOKENS GEN_AI_USAGE_REASONING_OUTPUT_TOKENS = 'gen_ai.usage.reasoning.output_tokens' @@ -79,6 +85,12 @@ class TokenUsage: if cached_tokens is not None: attrs[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] = cached_tokens + cache_creation_tokens = getattr( + self.usage_metadata, 'cache_creation_input_tokens', None + ) + if cache_creation_tokens is not None: + attrs[GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS] = cache_creation_tokens + thoughts_tokens = self.usage_metadata.thoughts_token_count if thoughts_tokens is not None: attrs[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] = thoughts_tokens diff --git a/tests/unittests/models/test_anthropic_llm.py b/tests/unittests/models/test_anthropic_llm.py index b03abdcd..1f25a217 100644 --- a/tests/unittests/models/test_anthropic_llm.py +++ b/tests/unittests/models/test_anthropic_llm.py @@ -1983,6 +1983,124 @@ def test_message_to_generate_content_response_maps_finish_reason( assert response.finish_reason == expected_finish_reason +def test_message_to_generate_content_response_reports_cache_creation_tokens(): + """cache_creation_input_tokens maps to usage_metadata.cache_creation_input_tokens.""" + from google.adk.models.anthropic_llm import message_to_generate_content_response + + message = anthropic_types.Message( + id="msg_cache_creation", + content=[ + anthropic_types.TextBlock(text="hi", type="text", citations=None) + ], + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage=anthropic_types.Usage( + input_tokens=100, + output_tokens=20, + cache_creation_input_tokens=50, + cache_read_input_tokens=0, + server_tool_use=None, + service_tier=None, + ), + ) + + response = message_to_generate_content_response(message) + + assert response.usage_metadata.cache_creation_input_tokens == 50 + dumped = response.model_dump() + assert "usage_metadata" in dumped + + +def test_message_to_generate_content_response_no_cache_creation_tokens(): + """Absent cache_creation_input_tokens yields cache_creation_input_tokens=None.""" + from google.adk.models.anthropic_llm import message_to_generate_content_response + + message = anthropic_types.Message( + id="msg_no_cache_creation", + content=[ + anthropic_types.TextBlock(text="hi", type="text", citations=None) + ], + model="claude-sonnet-4-20250514", + role="assistant", + stop_reason="end_turn", + stop_sequence=None, + type="message", + usage=anthropic_types.Usage( + input_tokens=100, + output_tokens=20, + cache_creation_input_tokens=None, + cache_read_input_tokens=0, + server_tool_use=None, + service_tier=None, + ), + ) + + response = message_to_generate_content_response(message) + + assert not hasattr(response.usage_metadata, "cache_creation_input_tokens") + + +@pytest.mark.asyncio +async def test_streaming_reports_cache_creation_tokens(): + """Anthropic streaming extracts and attaches cache_creation_input_tokens.""" + llm = AnthropicLlm(model="claude-sonnet-4-20250514") + + events = [ + MagicMock( + type="message_start", + message=MagicMock( + usage=MagicMock( + input_tokens=100, + output_tokens=0, + cache_creation_input_tokens=50, + cache_read_input_tokens=0, + ) + ), + ), + MagicMock( + type="content_block_start", + index=0, + content_block=anthropic_types.TextBlock(text="", type="text"), + ), + MagicMock( + type="content_block_delta", + index=0, + delta=anthropic_types.TextDelta(text="Hi", type="text_delta"), + ), + MagicMock(type="content_block_stop", index=0), + MagicMock( + type="message_delta", + delta=MagicMock(stop_reason="end_turn"), + usage=MagicMock(output_tokens=20), + ), + MagicMock(type="message_stop"), + ] + + mock_client = MagicMock() + mock_client.messages.create = AsyncMock( + return_value=_make_mock_stream_events(events) + ) + + llm_request = LlmRequest( + model="claude-sonnet-4-20250514", + contents=[Content(role="user", parts=[Part.from_text(text="Hi")])], + ) + + with mock.patch.object(llm, "_anthropic_client", mock_client): + responses = [ + r async for r in llm.generate_content_async(llm_request, stream=True) + ] + + assert len(responses) == 2 + final_response = responses[-1] + assert final_response.usage_metadata.cache_creation_input_tokens == 50 + dumped = final_response.model_dump() + assert "usage_metadata" in dumped + + def test_part_to_message_block_thinking_roundtrip(): """Part with thought=True and signature creates ThinkingBlockParam.""" part = Part( diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 7f67a87f..0b89d04f 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -1846,6 +1846,46 @@ async def test_generate_content_async_with_usage_metadata( mock_acompletion.assert_called_once() +@pytest.mark.asyncio +async def test_generate_content_async_with_bedrock_cache_tokens( + lite_llm_instance, mock_acompletion +): + mock_response_with_usage_metadata = ModelResponse( + choices=[ + Choices( + message=ChatCompletionAssistantMessage( + role="assistant", + content="Test response", + ) + ) + ], + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "cache_read_input_tokens": 8, + "cache_creation_input_tokens": 4, + }, + ) + mock_acompletion.return_value = mock_response_with_usage_metadata + + llm_request = LlmRequest( + contents=[ + types.Content( + role="user", parts=[types.Part.from_text(text="Test prompt")] + ), + ], + ) + async for response in lite_llm_instance.generate_content_async(llm_request): + assert response.usage_metadata.prompt_token_count == 10 + assert response.usage_metadata.candidates_token_count == 5 + 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 + + mock_acompletion.assert_called_once() + + @pytest.mark.asyncio async def test_generate_content_async_ollama_chat_preserves_multimodal_content( mock_acompletion, mock_completion @@ -4519,6 +4559,46 @@ async def test_generate_content_async_stream_with_usage_metadata( assert responses[3].usage_metadata.thoughts_token_count == 5 +@pytest.mark.asyncio +async def test_generate_content_async_stream_with_bedrock_cache_tokens( + mock_completion, lite_llm_instance +): + streaming_model_response_with_usage_metadata = [ + *STREAMING_MODEL_RESPONSE, + ModelResponseStream( + usage={ + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + "cache_read_input_tokens": 8, + "cache_creation_input_tokens": 4, + }, + choices=[ + StreamingChoices( + finish_reason=None, + ) + ], + ), + ] + + mock_completion.return_value = iter( + streaming_model_response_with_usage_metadata + ) + + responses = [ + response + async for response in lite_llm_instance.generate_content_async( + LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True + ) + ] + assert len(responses) == 4 + assert responses[3].usage_metadata.prompt_token_count == 10 + assert responses[3].usage_metadata.candidates_token_count == 5 + 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 + + @pytest.mark.asyncio async def test_generate_content_async_multiple_function_calls( mock_completion, lite_llm_instance diff --git a/tests/unittests/telemetry/test_token_usage.py b/tests/unittests/telemetry/test_token_usage.py index fcb37adc..b496b98c 100644 --- a/tests/unittests/telemetry/test_token_usage.py +++ b/tests/unittests/telemetry/test_token_usage.py @@ -219,3 +219,16 @@ def test_to_attributes_missing_optional_attrs(): attrs = token_usage.to_attributes() assert attrs[_token_usage.GEN_AI_USAGE_INPUT_TOKENS] == 10 assert attrs[_token_usage.GEN_AI_USAGE_OUTPUT_TOKENS] == 20 + + +def test_to_attributes_cache_creation( + usage_metadata: types.GenerateContentResponseUsageMetadata, +): + """Tests to_attributes when cache_creation_input_tokens is present.""" + usage_metadata.prompt_token_count = 10 + object.__setattr__(usage_metadata, "cache_creation_input_tokens", 50) + + token_usage = _token_usage.TokenUsage(usage_metadata) + attrs = token_usage.to_attributes() + assert attrs[_token_usage.GEN_AI_USAGE_INPUT_TOKENS] == 10 + assert attrs[_token_usage.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS] == 50