diff --git a/src/google/adk/telemetry/_metrics.py b/src/google/adk/telemetry/_metrics.py index 47cd2bc5b0d..10a7f757a39 100644 --- a/src/google/adk/telemetry/_metrics.py +++ b/src/google/adk/telemetry/_metrics.py @@ -580,11 +580,11 @@ def record_client_token_usage( if not responses: return - # The assumption is that token usage in streaming responses is cumulative. - # The last response chunk contains the total usage for the entire request. - # Summing them up across all response chunks would result in overcounting. - last_response = responses[-1] - if not last_response.usage_metadata: + # Usage in streaming responses is cumulative, so summing it would overcount. + # The newest report is the whole figure. It is not always on the last + # response: a stream cut off inside a tool call ends with one that has none. + token_usage = TokenUsage.from_llm_responses(responses) + if token_usage is None: logger.warning( "Skipping missing token usage metadata for agent %s and model %s", agent_name, @@ -598,10 +598,9 @@ def record_client_token_usage( # thoughts tokens for "output". # `cached_content_token_count` is omitted as it's already included in prompt tokens. # `total_token_count` is omitted as SemConv expects input/output breakdown. - token_usage = TokenUsage.from_usage_metadata(last_response.usage_metadata) input_token_count = token_usage.input_tokens or 0 output_token_count = token_usage.output_tokens or 0 - response_model = last_response.model_version or llm_request.model + response_model = responses[-1].model_version or llm_request.model base_attrs = { gen_ai_attributes.GEN_AI_AGENT_NAME: agent_name, gen_ai_attributes.GEN_AI_OPERATION_NAME: "generate_content", diff --git a/tests/unittests/telemetry/test_metrics.py b/tests/unittests/telemetry/test_metrics.py index fa7bd960d1e..eb46a3d34fd 100644 --- a/tests/unittests/telemetry/test_metrics.py +++ b/tests/unittests/telemetry/test_metrics.py @@ -17,6 +17,7 @@ import os from unittest import mock +from google.adk.models.llm_response import LlmResponse from google.adk.telemetry import _hallucination from google.adk.telemetry import _metrics from google.adk.telemetry import _token_usage @@ -458,6 +459,81 @@ def test_record_client_token_usage(mock_meter_setup): } +def test_record_client_token_usage_keeps_usage_before_trailing_response( + mock_meter_setup, +): + """Usage reported before the last response of a call is still recorded. + + A LiteLlm stream cut off by max_tokens inside a tool call yields the text + with the usage, then an error response that carries none. + """ + llm_request = mock.MagicMock(model="test-model") + responses = [ + LlmResponse( + content=types.ModelContent("Let me check the weather."), + finish_reason=types.FinishReason.MAX_TOKENS, + model_version="test-model-v1", + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=120, candidates_token_count=40 + ), + ), + LlmResponse( + finish_reason=types.FinishReason.MAX_TOKENS, + model_version="test-model-v1", + ), + ] + + _metrics.record_client_token_usage( + agent_name="test_agent", + llm_request=llm_request, + responses=responses, + ) + + recorded = { + kwargs["attributes"]["gen_ai.token.type"]: args[0] + for args, kwargs in mock_meter_setup[ + "client_token_usage" + ].record.call_args_list + } + assert recorded == {"input": 120, "output": 40} + + +@pytest.mark.parametrize("trailing_without_usage", [False, True]) +def test_record_client_token_usage_takes_newest_cumulative_report( + mock_meter_setup, trailing_without_usage +): + """Streaming usage is cumulative, so the newest report is recorded, not a sum.""" + llm_request = mock.MagicMock(model="test-model") + responses = [ + LlmResponse( + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=120, candidates_token_count=10 + ), + ), + LlmResponse( + usage_metadata=types.GenerateContentResponseUsageMetadata( + prompt_token_count=120, candidates_token_count=40 + ), + ), + ] + if trailing_without_usage: + responses.append(LlmResponse(finish_reason=types.FinishReason.MAX_TOKENS)) + + _metrics.record_client_token_usage( + agent_name="test_agent", + llm_request=llm_request, + responses=responses, + ) + + recorded = { + kwargs["attributes"]["gen_ai.token.type"]: args[0] + for args, kwargs in mock_meter_setup[ + "client_token_usage" + ].record.call_args_list + } + assert recorded == {"input": 120, "output": 40} + + @pytest.fixture(name="call_count_histograms") def _call_count_histograms(monkeypatch): """Redirects the two per-invocation call-count histograms."""