github-featured/google--adk-python
2
0
Fork 0
Code Issues Pull Requests Actions 7 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

fix: cache read write token counts in LiteLLM and Anthropic models

Browse Source 
Extract cache creation (write) tokens from LiteLLM and Anthropic model
usage metadata and map them to the GenerateContentResponseUsageMetadata
object. This ensures they are recorded in telemetry, allowing correct cost
calculations for prompt caching with providers like Bedrock.

Close #5835

PiperOrigin-RevId: 967259812
This commit is contained in:
Google Team Member
2026-08-19 09:44:11 -07:00
committed by Copybara-Service
parent 1cd6f464e5
commit d0b33a0569
6 changed files with 314 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/models/anthropic_llm.py
+38 -16
View File
@@ -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,
)
src/google/adk/models/lite_llm.py
+53 -1
View File
@@ -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,
src/google/adk/telemetry/_token_usage.py
+12
View File
@@ -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
tests/unittests/models/test_anthropic_llm.py
+118
View File
@@ -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(
tests/unittests/models/test_litellm.py
+80
View File
@@ -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
tests/unittests/telemetry/test_token_usage.py
+13
View File
@@ -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