feat(tracing): report cached input token usage (#1148)

Author: evanlowe

tests/test_llm_attributes_extractors.py

@@ -16,12 +16,17 @@
 from types import SimpleNamespace
 from unittest.mock import Mock, call
 
+import pytest
 from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
 from google.adk.tools.function_tool import FunctionTool
+from google.genai import types
 from opentelemetry.sdk.trace import TracerProvider
 
 from veadk.tracing.telemetry.attributes.extractors.llm_attributes_extractors import (
     llm_gen_ai_request_functions,
+    llm_gen_ai_usage_cache_creation_input_tokens,
+    llm_gen_ai_usage_cache_read_input_tokens,
     llm_gen_ai_usage_output_tokens,
 )
 from veadk.tracing.telemetry.attributes.extractors.types import ExtractorResponse
@@ -118,6 +123,45 @@ def test_missing_output_token_count_is_not_written_to_span():
     span.set_attribute.assert_not_called()
 
 
+@pytest.mark.parametrize("cached_tokens", [600, 0, None])
+def test_cached_tokens_are_only_reported_as_cache_reads(cached_tokens):
+    params = SimpleNamespace(
+        llm_response=LlmResponse(
+            usage_metadata=types.GenerateContentResponseUsageMetadata(
+                cached_content_token_count=cached_tokens,
+            )
+        )
+    )
+    provider = TracerProvider()
+    try:
+        with provider.get_tracer(__name__).start_as_current_span("cache-usage") as span:
+            ExtractorResponse.update_span(
+                span,
+                "gen_ai.usage.cache_read_input_tokens",
+                llm_gen_ai_usage_cache_read_input_tokens(params),
+            )
+            ExtractorResponse.update_span(
+                span,
+                "gen_ai.usage.cache_creation_input_tokens",
+                llm_gen_ai_usage_cache_creation_input_tokens(params),
+            )
+            expected = (
+                {"gen_ai.usage.cache_read_input_tokens": cached_tokens}
+                if cached_tokens is not None
+                else {}
+            )
+            assert dict(span.attributes) == expected
+    finally:
+        provider.shutdown()
+
+
+def test_missing_usage_omits_cache_attributes():
+    params = SimpleNamespace(llm_response=LlmResponse())
+
+    assert llm_gen_ai_usage_cache_read_input_tokens(params).content is None
+    assert llm_gen_ai_usage_cache_creation_input_tokens(params).content is None
+
+
 def test_falsy_attribute_values_are_written_to_span():
     span = Mock()
 

tests/test_portal_metrics.py

@@ -16,6 +16,10 @@
 from types import SimpleNamespace
 
 import pytest
+from google.adk.agents.run_config import RunConfig, StreamingMode
+from google.adk.models.llm_request import LlmRequest
+from google.adk.models.llm_response import LlmResponse
+from google.genai import types
 from opentelemetry import metrics as metrics_api
 from opentelemetry.metrics import _internal as metrics_internal
 from opentelemetry.sdk import metrics as metrics_sdk
@@ -138,6 +142,79 @@ def test_proxy_instruments_follow_provider_installed_later(
     ]
 
 
+@pytest.mark.parametrize(
+    "server_address",
+    [
+        "https://ark.cn-beijing.volces.com/api/v3/",
+        "https://ark.ap-southeast.bytepluses.com/api/v3",
+    ],
+    ids=["volcengine", "byteplus"],
+)
+@pytest.mark.parametrize("streaming_mode", [StreamingMode.NONE, StreamingMode.SSE])
+@pytest.mark.parametrize("cached_tokens", [600, 0, None])
+def test_cache_read_token_metrics_preserve_input_usage(
+    fresh_global_meter_provider, server_address, streaming_mode, cached_tokens
+):
+    reader = InMemoryMetricReader()
+    metrics_api.set_meter_provider(metrics_sdk.MeterProvider(metric_readers=[reader]))
+    recorder = PortalMetricRecorder(name="test-cache-read")
+    context = SimpleNamespace(
+        agent=SimpleNamespace(model_api_base=server_address),
+        run_config=RunConfig(streaming_mode=streaming_mode),
+    )
+    request = LlmRequest(model="test-model")
+
+    if streaming_mode == StreamingMode.SSE:
+        recorder.record_call_llm(context, "partial", request, LlmResponse(partial=True))
+
+    recorder.record_call_llm(
+        context,
+        "final",
+        request,
+        LlmResponse(
+            usage_metadata=types.GenerateContentResponseUsageMetadata(
+                prompt_token_count=1000,
+                candidates_token_count=200,
+                total_token_count=1200,
+                cached_content_token_count=cached_tokens,
+            ),
+        ),
+    )
+
+    metrics_data = reader.get_metrics_data()
+    recorded_metrics = {
+        metric.name: list(metric.data.data_points)
+        for resource_metrics in metrics_data.resource_metrics
+        for scope_metrics in resource_metrics.scope_metrics
+        for metric in scope_metrics.metrics
+    }
+    points = recorded_metrics["gen_ai.client.token.usage"]
+    expected = {"input": 1000, "output": 200}
+    if cached_tokens is not None:
+        expected["cache_read"] = cached_tokens
+    assert {point.attributes["gen_ai_token_type"]: point.sum for point in points} == (
+        expected
+    )
+    assert len(points) == len(expected)
+    for point in points:
+        assert point.count == 1
+        assert point.attributes["server_address"] == server_address
+        assert point.attributes["stream"] == (streaming_mode == StreamingMode.SSE)
+        assert point.attributes["gen_ai_response_model"] == "test-model"
+    assert [point.value for point in recorded_metrics["gen_ai.chat.count"]] == [1]
+
+
+def test_missing_usage_does_not_record_token_metrics(fresh_global_meter_provider):
+    reader = InMemoryMetricReader()
+    metrics_api.set_meter_provider(metrics_sdk.MeterProvider(metric_readers=[reader]))
+    recorder = PortalMetricRecorder(name="test-missing-usage")
+    context = SimpleNamespace(agent=SimpleNamespace(), run_config=None)
+
+    recorder.record_call_llm(context, "event", LlmRequest(), LlmResponse())
+
+    assert reader.get_metrics_data() is None
+
+
 def test_apmplus_reuses_preconfigured_global_provider(
     fresh_global_meter_provider,
     monkeypatch,

veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py

@@ -205,29 +205,16 @@ def llm_gen_ai_usage_total_tokens(params: LLMAttributesParams) -> ExtractorRespo
     return ExtractorResponse(content=None)
 
 
-# FIXME
 def llm_gen_ai_usage_cache_creation_input_tokens(
     params: LLMAttributesParams,
 ) -> ExtractorResponse:
-    """Extract the number of tokens used for cache creation.
+    """Omit cache creation usage until response metadata provides it
 
-    Provides the count of tokens used for creating cached content,
-    which affects cost calculation in caching-enabled models.
-
-    Args:
-        params: LLM execution parameters containing response metadata
-
-    Returns:
-        ExtractorResponse: Response containing cache creation token count or None
+    cached_content_token_count measures cache reads, not cache creation
     """
-    if params.llm_response.usage_metadata:
-        return ExtractorResponse(
-            content=params.llm_response.usage_metadata.cached_content_token_count,
-        )
     return ExtractorResponse(content=None)
 
 
-# FIXME
 def llm_gen_ai_usage_cache_read_input_tokens(
     params: LLMAttributesParams,
 ) -> ExtractorResponse:

veadk/tracing/telemetry/portal_metrics.py

@@ -162,7 +162,7 @@ class PortalMetricRecorder:
 
     Metrics Collected:
     - LLM invocation counts and frequencies
-    - Token consumption (input/output) with histogram distribution
+    - Token consumption (input/output/cache_read) with histogram distribution
     - Operation latency with performance bucket analysis
     - Error rates and exception details
     - Span-level performance metrics for APMPlus dashboards
@@ -258,7 +258,7 @@ def record_call_llm(
 
         Metrics Recorded:
         - Invocation count with model and operation attributes
-        - Input/output token usage with separate tracking
+        - Input/output token usage and the cached subset of input tokens
         - Operation duration from span timing data
         - Error counts and exception details
         - Span latency for performance analysis
@@ -290,13 +290,18 @@ def record_call_llm(
             # upload token usage
             input_token = llm_response.usage_metadata.prompt_token_count
             output_token = llm_response.usage_metadata.candidates_token_count
+            cache_read_token = llm_response.usage_metadata.cached_content_token_count
 
             if input_token:
                 token_attributes = {**attributes, "gen_ai_token_type": "input"}
                 self.token_usage.record(input_token, attributes=token_attributes)
             if output_token:
                 token_attributes = {**attributes, "gen_ai_token_type": "output"}
                 self.token_usage.record(output_token, attributes=token_attributes)
+            # Cache reads are already included in input tokens, not extra usage
+            if cache_read_token is not None:
+                token_attributes = {**attributes, "gen_ai_token_type": "cache_read"}
+                self.token_usage.record(cache_read_token, attributes=token_attributes)
 
             # Get llm duration
             span = trace.get_current_span()