Track Gemini cache token metrics

This commit is contained in:
adavyas 2026-03-14 12:59:55 -07:00
parent e86cbac75b
commit e5841882ad
2 changed files with 75 additions and 0 deletions

View File

@ -506,6 +506,23 @@ def extract_openai_cache_tokens(usage: Any) -> tuple[int, int]:
return cache_creation, cache_read
def extract_gemini_cache_tokens(usage_metadata: Any) -> tuple[int, int]:
"""
Extract cache token counts from Gemini usage metadata when available.
Gemini exposes `cached_content_token_count` on usage metadata when cached
content was reused for a request. Generate responses do not currently expose
a corresponding cache-creation count, so creation is reported as 0.
"""
if not usage_metadata:
return 0, 0
cache_read = getattr(usage_metadata, "cached_content_token_count", 0)
if not isinstance(cache_read, int):
cache_read = 0
return 0, cache_read
class HonchoLLMCallResponse(BaseModel, Generic[T]):
"""
Response object for LLM calls.
@ -2230,6 +2247,9 @@ async def honcho_llm_call_inner(
if gemini_response.usage_metadata
else 0
)
cache_creation_tokens, cache_read_tokens = extract_gemini_cache_tokens(
gemini_response.usage_metadata
)
output_token_count = (
gemini_response.usage_metadata.candidates_token_count or 0
if gemini_response.usage_metadata
@ -2259,6 +2279,8 @@ async def honcho_llm_call_inner(
content=text_content,
input_tokens=input_token_count,
output_tokens=output_token_count,
cache_creation_input_tokens=cache_creation_tokens,
cache_read_input_tokens=cache_read_tokens,
finish_reasons=[finish_reason],
tool_calls_made=gemini_tool_calls,
)
@ -2278,6 +2300,9 @@ async def honcho_llm_call_inner(
if gemini_response.usage_metadata
else 0
)
cache_creation_tokens, cache_read_tokens = extract_gemini_cache_tokens(
gemini_response.usage_metadata
)
output_token_count = (
gemini_response.usage_metadata.candidates_token_count or 0
if gemini_response.usage_metadata
@ -2312,6 +2337,8 @@ async def honcho_llm_call_inner(
content=gemini_response.parsed,
input_tokens=input_token_count,
output_tokens=output_token_count,
cache_creation_input_tokens=cache_creation_tokens,
cache_read_input_tokens=cache_read_tokens,
finish_reasons=[finish_reason],
tool_calls_made=[],
)

View File

@ -31,6 +31,7 @@ from src.utils.clients import (
CLIENTS,
HonchoLLMCallResponse,
HonchoLLMCallStreamChunk,
extract_gemini_cache_tokens,
handle_streaming_response,
honcho_llm_call,
honcho_llm_call_inner,
@ -89,6 +90,16 @@ class TestLLMCallResponse:
assert isinstance(chunk.finish_reasons, list)
assert chunk.finish_reasons == []
def test_extract_gemini_cache_tokens(self):
"""Gemini cached content tokens should map to cache-read metrics."""
usage_metadata = Mock()
usage_metadata.cached_content_token_count = 321
cache_creation, cache_read = extract_gemini_cache_tokens(usage_metadata)
assert cache_creation == 0
assert cache_read == 321
@pytest.mark.asyncio
class TestAnthropicClient:
@ -648,6 +659,43 @@ class TestGoogleClient:
call_args = mock_aio.models.generate_content.call_args
assert call_args.kwargs["config"]["max_output_tokens"] == 100
async def test_google_basic_call_surfaces_cached_content_tokens(self):
"""Gemini usage metadata should populate cache-read input tokens."""
from google import genai
mock_client = Mock(spec=genai.Client)
mock_response = Mock()
mock_part = Mock()
mock_part.text = "Hello from Gemini"
mock_part.function_call = None
mock_content = Mock()
mock_content.parts = [mock_part]
mock_finish_reason = Mock()
mock_finish_reason.name = "STOP"
mock_candidate = Mock()
mock_candidate.content = mock_content
mock_candidate.finish_reason = mock_finish_reason
mock_response.candidates = [mock_candidate]
mock_usage_metadata = Mock()
mock_usage_metadata.prompt_token_count = 100
mock_usage_metadata.candidates_token_count = 5
mock_usage_metadata.cached_content_token_count = 90
mock_response.usage_metadata = mock_usage_metadata
mock_aio = Mock()
mock_aio.models.generate_content = AsyncMock(return_value=mock_response)
mock_client.aio = mock_aio
with patch.dict(CLIENTS, {"google": mock_client}):
response = await honcho_llm_call_inner(
provider="google",
model="gemini-2.5-flash",
prompt="Hello",
max_tokens=100,
)
assert response.cache_creation_input_tokens == 0
assert response.cache_read_input_tokens == 90
async def test_google_json_mode(self):
"""Test Google/Gemini with JSON mode"""
from google import genai