feat: Add support for extracting cache-related token counts from LiteLLM usage

Closes #3049

Co-authored-by: Eliza Huang <[email protected]>
PiperOrigin-RevId: 828091671

Author: lizzij

src/google/adk/models/lite_llm.py

@@ -101,6 +101,7 @@ class UsageMetadataChunk(BaseModel):
   prompt_tokens: int
   completion_tokens: int
   total_tokens: int
+  cached_prompt_tokens: int = 0
 
 
 class LiteLLMClient:
@@ -217,6 +218,59 @@ def _append_fallback_user_content_if_missing(
   )
 
 
+def _extract_cached_prompt_tokens(usage: Any) -> int:
+  """Extracts cached prompt tokens from LiteLLM usage.
+
+  Providers expose cached token metrics in different shapes. Common patterns:
+  - usage["prompt_tokens_details"]["cached_tokens"] (OpenAI/Azure style)
+  - usage["prompt_tokens_details"] is a list of dicts with cached_tokens
+  - usage["cached_prompt_tokens"] (LiteLLM-normalized for some providers)
+  - usage["cached_tokens"] (flat)
+
+  Args:
+    usage: Usage dictionary from LiteLLM response.
+
+  Returns:
+    Integer number of cached prompt tokens if present; otherwise 0.
+  """
+  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 0
+
+    if not isinstance(usage_dict, dict):
+      return 0
+
+    details = usage_dict.get("prompt_tokens_details")
+    if isinstance(details, dict):
+      value = details.get("cached_tokens")
+      if isinstance(value, int):
+        return value
+    elif isinstance(details, list):
+      total = sum(
+          item.get("cached_tokens", 0)
+          for item in details
+          if isinstance(item, dict)
+          and isinstance(item.get("cached_tokens"), int)
+      )
+      if total > 0:
+        return total
+
+    for key in ("cached_prompt_tokens", "cached_tokens"):
+      value = usage_dict.get(key)
+      if isinstance(value, int):
+        return value
+  except (TypeError, AttributeError) as e:
+    logger.debug("Error extracting cached prompt tokens: %s", e)
+
+  return 0
+
+
 def _content_to_message_param(
     content: types.Content,
 ) -> Union[Message, list[Message]]:
@@ -533,6 +587,7 @@ def _model_response_to_chunk(
         prompt_tokens=response["usage"].get("prompt_tokens", 0),
         completion_tokens=response["usage"].get("completion_tokens", 0),
         total_tokens=response["usage"].get("total_tokens", 0),
+        cached_prompt_tokens=_extract_cached_prompt_tokens(response["usage"]),
     ), None
 
 
@@ -576,6 +631,9 @@ def _model_response_to_generate_content_response(
         prompt_token_count=response["usage"].get("prompt_tokens", 0),
         candidates_token_count=response["usage"].get("completion_tokens", 0),
         total_token_count=response["usage"].get("total_tokens", 0),
+        cached_content_token_count=_extract_cached_prompt_tokens(
+            response["usage"]
+        ),
     )
   return llm_response
 
@@ -965,6 +1023,7 @@ async def generate_content_async(
                 prompt_token_count=chunk.prompt_tokens,
                 candidates_token_count=chunk.completion_tokens,
                 total_token_count=chunk.total_tokens,
+                cached_content_token_count=chunk.cached_prompt_tokens,
             )
 
           if (

tests/unittests/models/test_litellm.py

@@ -1063,6 +1063,7 @@ async def test_generate_content_async_with_usage_metadata(
           "prompt_tokens": 10,
           "completion_tokens": 5,
           "total_tokens": 15,
+          "cached_tokens": 8,
       },
   )
   mock_acompletion.return_value = mock_response_with_usage_metadata
@@ -1083,6 +1084,7 @@ async def test_generate_content_async_with_usage_metadata(
     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
 
   mock_acompletion.assert_called_once()
 
@@ -1718,37 +1720,42 @@ async def test_generate_content_async_stream_with_usage_metadata(
 
 
 @pytest.mark.asyncio
-async def test_generate_content_async_stream_with_usage_metadata_only(
+async def test_generate_content_async_stream_with_usage_metadata(
     mock_completion, lite_llm_instance
 ):
+  """Tests that cached prompt tokens are propagated in streaming mode."""
   streaming_model_response_with_usage_metadata = [
+      *STREAMING_MODEL_RESPONSE,
       ModelResponse(
           usage={
               "prompt_tokens": 10,
               "completion_tokens": 5,
               "total_tokens": 15,
+              "cached_tokens": 8,
           },
           choices=[
               StreamingChoices(
-                  finish_reason="stop",
-                  delta=Delta(content=""),
+                  finish_reason=None,
               )
           ],
       ),
   ]
+
   mock_completion.return_value = iter(
       streaming_model_response_with_usage_metadata
   )
 
-  unused_responses = [
+  responses = [
       response
       async for response in lite_llm_instance.generate_content_async(
           LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
       )
   ]
-  mock_completion.assert_called_once()
-  _, kwargs = mock_completion.call_args
-  assert kwargs["stream_options"] == {"include_usage": 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
 
 
 @pytest.mark.asyncio
@@ -2038,6 +2045,36 @@ def test_function_declaration_to_tool_param_edge_cases():
   assert "required" not in result["function"]["parameters"]
 
 
+@pytest.mark.parametrize(
+    "usage, expected_tokens",
+    [
+        ({"prompt_tokens_details": {"cached_tokens": 123}}, 123),
+        (
+            {
+                "prompt_tokens_details": [
+                    {"cached_tokens": 50},
+                    {"cached_tokens": 25},
+                ]
+            },
+            75,
+        ),
+        ({"cached_prompt_tokens": 45}, 45),
+        ({"cached_tokens": 67}, 67),
+        ({"prompt_tokens": 100}, 0),
+        ({}, 0),
+        ("not a dict", 0),
+        (None, 0),
+        ({"prompt_tokens_details": {"cached_tokens": "not a number"}}, 0),
+        (json.dumps({"cached_tokens": 89}), 89),
+        (json.dumps({"some_key": "some_value"}), 0),
+    ],
+)
+def test_extract_cached_prompt_tokens(usage, expected_tokens):
+  from google.adk.models.lite_llm import _extract_cached_prompt_tokens
+
+  assert _extract_cached_prompt_tokens(usage) == expected_tokens
+
+
 def test_gemini_via_litellm_warning(monkeypatch):
   """Test that Gemini via LiteLLM shows warning."""
   # Ensure environment variable is not set