Track Gemini cache token metrics
This commit is contained in:
parent
e86cbac75b
commit
e5841882ad
|
|
@ -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=[],
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in New Issue