From 262ec2fbb36ef8ea55209b06d982d20059151ff4 Mon Sep 17 00:00:00 2001 From: evanlowe <62918515+evanlowe@users.noreply.github.com> Date: Wed, 23 Sep 2026 11:03:32 +0800 Subject: [PATCH] feat(tracing): report cached input token usage --- tests/test_llm_attributes_extractors.py | 44 +++++++++++ tests/test_portal_metrics.py | 77 +++++++++++++++++++ .../extractors/llm_attributes_extractors.py | 17 +--- veadk/tracing/telemetry/portal_metrics.py | 9 ++- 4 files changed, 130 insertions(+), 17 deletions(-) diff --git a/tests/test_llm_attributes_extractors.py b/tests/test_llm_attributes_extractors.py index 56f3d53cd..c1ab1d6ee 100644 --- a/tests/test_llm_attributes_extractors.py +++ b/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() diff --git a/tests/test_portal_metrics.py b/tests/test_portal_metrics.py index 0b60a2353..45d52e66e 100644 --- a/tests/test_portal_metrics.py +++ b/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, diff --git a/veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py b/veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py index 94645a08d..b437483e5 100644 --- a/veadk/tracing/telemetry/attributes/extractors/llm_attributes_extractors.py +++ b/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: diff --git a/veadk/tracing/telemetry/portal_metrics.py b/veadk/tracing/telemetry/portal_metrics.py index 4842dce2a..a7a2fa5ac 100644 --- a/veadk/tracing/telemetry/portal_metrics.py +++ b/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,6 +290,7 @@ 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"} @@ -297,6 +298,10 @@ def record_call_llm( 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()