diff --git a/core/src/main/java/com/google/adk/telemetry/Tracing.java b/core/src/main/java/com/google/adk/telemetry/Tracing.java index d57e46c54..226d7011c 100644 --- a/core/src/main/java/com/google/adk/telemetry/Tracing.java +++ b/core/src/main/java/com/google/adk/telemetry/Tracing.java @@ -336,9 +336,13 @@ public static void traceCallLlm( .usageMetadata() .ifPresent( usage -> { - usage - .promptTokenCount() - .ifPresent(tokens -> span.setAttribute(GEN_AI_USAGE_INPUT_TOKENS, (long) tokens)); + if (usage.promptTokenCount().isPresent() + || usage.toolUsePromptTokenCount().isPresent()) { + span.setAttribute( + GEN_AI_USAGE_INPUT_TOKENS, + (long) usage.promptTokenCount().orElse(0) + + usage.toolUsePromptTokenCount().orElse(0)); + } // According to OpenTelemetry Semantic Conventions: // https://github.com/open-telemetry/semantic-conventions/blob/v1.41.0/docs/registry/attributes/gen-ai.md // gen_ai.usage.reasoning.output_tokens (thoughts_token_count) SHOULD be included in diff --git a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java index ecfba4995..33810f081 100644 --- a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java +++ b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java @@ -463,6 +463,75 @@ public void testTraceCallLlm_withReasoningAndCacheTokens() { 15L, (long) attrs.get(AttributeKey.longKey("gen_ai.usage.reasoning.output_tokens"))); } + @Test + public void testTraceCallLlm_withToolUsePromptTokens() { + Span span = tracer.spanBuilder("test-tool-use-prompt").startSpan(); + try (Scope scope = span.makeCurrent()) { + LlmRequest llmRequest = + LlmRequest.builder() + .model("gemini-pro") + .contents(ImmutableList.of(Content.fromParts(Part.fromText("hello")))) + .config(GenerateContentConfig.builder().topP(0.9f).maxOutputTokens(100).build()) + .build(); + LlmResponse llmResponse = + LlmResponse.builder() + .content(Content.builder().parts(Part.fromText("world")).build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .usageMetadata( + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(10) + .toolUsePromptTokenCount(8) + .candidatesTokenCount(20) + .totalTokenCount(38) + .build()) + .build(); + Tracing.traceCallLlm( + span, buildInvocationContext(), "event-1", llmRequest, llmResponse, null); + } finally { + span.end(); + } + List spans = openTelemetryRule.getSpans(); + assertThat(spans).hasSize(1); + SpanData spanData = spans.get(0); + Attributes attrs = spanData.getAttributes(); + assertEquals(18L, (long) attrs.get(AttributeKey.longKey("gen_ai.usage.input_tokens"))); + assertEquals(20L, (long) attrs.get(AttributeKey.longKey("gen_ai.usage.output_tokens"))); + } + + @Test + public void testTraceCallLlm_withOnlyToolUsePromptTokens() { + Span span = tracer.spanBuilder("test-only-tool-use-prompt").startSpan(); + try (Scope scope = span.makeCurrent()) { + LlmRequest llmRequest = + LlmRequest.builder() + .model("gemini-pro") + .contents(ImmutableList.of(Content.fromParts(Part.fromText("hello")))) + .config(GenerateContentConfig.builder().topP(0.9f).maxOutputTokens(100).build()) + .build(); + LlmResponse llmResponse = + LlmResponse.builder() + .content(Content.builder().parts(Part.fromText("world")).build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .usageMetadata( + GenerateContentResponseUsageMetadata.builder() + .toolUsePromptTokenCount(8) + .candidatesTokenCount(20) + .totalTokenCount(28) + .build()) + .build(); + Tracing.traceCallLlm( + span, buildInvocationContext(), "event-1", llmRequest, llmResponse, null); + } finally { + span.end(); + } + List spans = openTelemetryRule.getSpans(); + assertThat(spans).hasSize(1); + SpanData spanData = spans.get(0); + Attributes attrs = spanData.getAttributes(); + assertEquals(8L, (long) attrs.get(AttributeKey.longKey("gen_ai.usage.input_tokens"))); + assertEquals(20L, (long) attrs.get(AttributeKey.longKey("gen_ai.usage.output_tokens"))); + } + @Test public void testTraceSendData() { Span span = tracer.spanBuilder("test").startSpan();