diff --git a/core/src/main/java/com/google/adk/telemetry/README.md b/core/src/main/java/com/google/adk/telemetry/README.md index 8665b3352..f3dfd212c 100644 --- a/core/src/main/java/com/google/adk/telemetry/README.md +++ b/core/src/main/java/com/google/adk/telemetry/README.md @@ -37,7 +37,8 @@ Calls to Large Language Models (LLMs) are traced within a `call_llm` span. The * Model name (`gen_ai.request.model`). * Token usage (`gen_ai.usage.input_tokens`, `gen_ai.usage.output_tokens`). * Configuration parameters (`gen_ai.request.top_p`, - `gen_ai.request.max_tokens`). + `gen_ai.request.max_tokens`, `gen_ai.request.reasoning.level`, and the + thinking budget as `gen_ai.usage.experimental.reasoning_tokens_limit`). * Response finish reason (`gen_ai.response.finish_reasons`). ### Tool Calls and Responses 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 f284f574f..97c063fc2 100644 --- a/core/src/main/java/com/google/adk/telemetry/Tracing.java +++ b/core/src/main/java/com/google/adk/telemetry/Tracing.java @@ -30,6 +30,7 @@ import com.google.genai.types.Content; import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; +import com.google.genai.types.ThinkingConfig; import io.opentelemetry.api.GlobalOpenTelemetry; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.trace.Span; @@ -110,6 +111,8 @@ public class Tracing { AttributeKey.doubleKey("gen_ai.request.top_p"); private static final AttributeKey GEN_AI_REQUEST_MAX_TOKENS = AttributeKey.longKey("gen_ai.request.max_tokens"); + private static final AttributeKey GEN_AI_REQUEST_REASONING_LEVEL = + AttributeKey.stringKey("gen_ai.request.reasoning.level"); private static final AttributeKey GEN_AI_USAGE_INPUT_TOKENS = AttributeKey.longKey("gen_ai.usage.input_tokens"); private static final AttributeKey GEN_AI_USAGE_OUTPUT_TOKENS = @@ -118,6 +121,8 @@ public class Tracing { AttributeKey.longKey("gen_ai.usage.cache_read.input_tokens"); private static final AttributeKey GEN_AI_USAGE_REASONING_OUTPUT_TOKENS = AttributeKey.longKey("gen_ai.usage.reasoning.output_tokens"); + private static final AttributeKey GEN_AI_USAGE_REASONING_TOKENS_LIMIT = + AttributeKey.longKey("gen_ai.usage.experimental.reasoning_tokens_limit"); private static final AttributeKey ADK_TOOL_CALL_ARGS = AttributeKey.stringKey("gcp.vertex.agent.tool_call_args"); @@ -331,6 +336,19 @@ public static void traceCallLlm( .ifPresent( maxTokens -> span.setAttribute(GEN_AI_REQUEST_MAX_TOKENS, maxTokens.longValue())); + config + .thinkingConfig() + .flatMap(ThinkingConfig::thinkingBudget) + .ifPresent( + budget -> + span.setAttribute( + GEN_AI_USAGE_REASONING_TOKENS_LIMIT, budget.longValue())); + // OTel wants the exact string sent to the provider, which toString() returns. + config + .thinkingConfig() + .flatMap(ThinkingConfig::thinkingLevel) + .ifPresent( + level -> span.setAttribute(GEN_AI_REQUEST_REASONING_LEVEL, level.toString())); }); llmResponse .usageMetadata() 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 e641cab18..dc4831eb8 100644 --- a/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java +++ b/core/src/test/java/com/google/adk/telemetry/ContextPropagationTest.java @@ -44,6 +44,7 @@ import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; +import com.google.genai.types.ThinkingConfig; import io.opentelemetry.api.common.AttributeKey; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; @@ -533,6 +534,69 @@ public void testTraceCallLlm_withOnlyToolUsePromptTokens() { assertEquals(20L, (long) attrs.get(AttributeKey.longKey("gen_ai.usage.output_tokens"))); } + @Test + public void testTraceCallLlm_withThinkingBudget() { + Span span = tracer.spanBuilder("test-thinking-budget").startSpan(); + try (Scope scope = span.makeCurrent()) { + LlmRequest llmRequest = + LlmRequest.builder() + .model("gemini-pro") + .contents(ImmutableList.of(Content.fromParts(Part.fromText("hello")))) + .config( + GenerateContentConfig.builder() + .thinkingConfig(ThinkingConfig.builder().thinkingBudget(2048).build()) + .build()) + .build(); + Tracing.traceCallLlm( + span, + buildInvocationContext(), + "event-1", + llmRequest, + LlmResponse.builder().build(), + null); + } finally { + span.end(); + } + List spans = openTelemetryRule.getSpans(); + assertThat(spans).hasSize(1); + Attributes attrs = spans.get(0).getAttributes(); + assertThat(attrs.get(AttributeKey.longKey("gen_ai.usage.experimental.reasoning_tokens_limit"))) + .isEqualTo(2048L); + assertThat(attrs.get(AttributeKey.stringKey("gen_ai.request.reasoning.level"))).isNull(); + } + + @Test + public void testTraceCallLlm_withThinkingLevel_recordsValueSentToProvider() { + Span span = tracer.spanBuilder("test-thinking-level").startSpan(); + try (Scope scope = span.makeCurrent()) { + LlmRequest llmRequest = + LlmRequest.builder() + .model("gemini-pro") + .contents(ImmutableList.of(Content.fromParts(Part.fromText("hello")))) + .config( + GenerateContentConfig.builder() + .thinkingConfig(ThinkingConfig.builder().thinkingLevel("low").build()) + .build()) + .build(); + Tracing.traceCallLlm( + span, + buildInvocationContext(), + "event-1", + llmRequest, + LlmResponse.builder().build(), + null); + } finally { + span.end(); + } + List spans = openTelemetryRule.getSpans(); + assertThat(spans).hasSize(1); + Attributes attrs = spans.get(0).getAttributes(); + assertThat(attrs.get(AttributeKey.stringKey("gen_ai.request.reasoning.level"))) + .isEqualTo("low"); + assertThat(attrs.get(AttributeKey.longKey("gen_ai.usage.experimental.reasoning_tokens_limit"))) + .isNull(); + } + @Test public void testTraceSendData() { Span span = tracer.spanBuilder("test").startSpan();