From e5841882ad63a6a4a30892e95c21750bae3c3d78 Mon Sep 17 00:00:00 2001
From: adavyas <adavyasharma@gmail.com>
Date: Sat, 14 Mar 2026 12:59:55 -0700
Subject: [PATCH] Track Gemini cache token metrics

---
 src/utils/clients.py        | 27 +++++++++++++++++++++
 tests/utils/test_clients.py | 48 +++++++++++++++++++++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/src/utils/clients.py b/src/utils/clients.py
index 9c25dce1..7bc71bfd 100644
--- a/src/utils/clients.py
+++ b/src/utils/clients.py
@@ -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=[],
                 )
diff --git a/tests/utils/test_clients.py b/tests/utils/test_clients.py
index d9546658..e30cf160 100644
--- a/tests/utils/test_clients.py
+++ b/tests/utils/test_clients.py
@@ -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