diff --git a/core/src/telemetry/token_usage.ts b/core/src/telemetry/token_usage.ts new file mode 100644 index 000000000..b8f55e693 --- /dev/null +++ b/core/src/telemetry/token_usage.ts @@ -0,0 +1,108 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {GenerateContentResponseUsageMetadata} from '@google/genai'; +import {Attributes} from '@opentelemetry/api'; + +/** OpenTelemetry GenAI attribute for the number of input tokens used. */ +const GEN_AI_USAGE_INPUT_TOKENS = 'gen_ai.usage.input_tokens'; + +/** OpenTelemetry GenAI attribute for the number of output tokens used. */ +const GEN_AI_USAGE_OUTPUT_TOKENS = 'gen_ai.usage.output_tokens'; + +/** OpenTelemetry GenAI attribute for input tokens served from cache. */ +const GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS = + 'gen_ai.usage.cache_read.input_tokens'; + +/** OpenTelemetry GenAI attribute for output tokens spent on reasoning. */ +const GEN_AI_USAGE_REASONING_OUTPUT_TOKENS = + 'gen_ai.usage.reasoning.output_tokens'; + +/** + * Not part of the GenAI semantic conventions. The spelling stays snake_case + * and byte-identical to the key adk-python emits so both runtimes land in the + * same dashboard series. + */ +const GEN_AI_USAGE_SYSTEM_INSTRUCTION_TOKENS = + 'gen_ai.usage.experimental.system_instruction_tokens'; + +/** + * Token counts the backend may return but `@google/genai` (2.9.0) does not + * declare. + */ +interface UsageMetadataWithSystemInstructionTokens extends GenerateContentResponseUsageMetadata { + systemInstructionTokens?: number; +} + +/** + * Adds two optional token counts, preserving the distinction between "no + * counts reported" (`undefined`) and "counts reported, and they are zero". + */ +function addTokenCounts(a?: number, b?: number): number | undefined { + if (a === undefined && b === undefined) { + return undefined; + } + return (a ?? 0) + (b ?? 0); +} + +/** Centralized representation and processing of GenAI token usage metadata. */ +export class TokenUsage { + constructor(readonly usageMetadata?: GenerateContentResponseUsageMetadata) {} + + /** + * Prompt and tool-use tokens, which the GenAI semantic conventions bucket + * together as `input`. + */ + get inputTokenCount(): number | undefined { + return addTokenCounts( + this.usageMetadata?.promptTokenCount, + this.usageMetadata?.toolUsePromptTokenCount, + ); + } + + /** + * Candidate and reasoning tokens. The semantic conventions require + * `gen_ai.usage.reasoning.output_tokens` to be included in the output total. + */ + get outputTokenCount(): number | undefined { + return addTokenCounts( + this.usageMetadata?.candidatesTokenCount, + this.usageMetadata?.thoughtsTokenCount, + ); + } + + /** Returns the OpenTelemetry token usage attributes, omitting unknowns. */ + toAttributes(): Attributes { + const attributes: Attributes = {}; + const inputTokenCount = this.inputTokenCount; + if (inputTokenCount !== undefined) { + attributes[GEN_AI_USAGE_INPUT_TOKENS] = inputTokenCount; + } + const outputTokenCount = this.outputTokenCount; + if (outputTokenCount !== undefined) { + attributes[GEN_AI_USAGE_OUTPUT_TOKENS] = outputTokenCount; + } + + const metadata: UsageMetadataWithSystemInstructionTokens | undefined = + this.usageMetadata; + if (metadata === undefined) { + return attributes; + } + if (metadata.cachedContentTokenCount !== undefined) { + attributes[GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS] = + metadata.cachedContentTokenCount; + } + if (metadata.thoughtsTokenCount !== undefined) { + attributes[GEN_AI_USAGE_REASONING_OUTPUT_TOKENS] = + metadata.thoughtsTokenCount; + } + if (metadata.systemInstructionTokens !== undefined) { + attributes[GEN_AI_USAGE_SYSTEM_INSTRUCTION_TOKENS] = + metadata.systemInstructionTokens; + } + return attributes; + } +} diff --git a/core/src/telemetry/tracing.ts b/core/src/telemetry/tracing.ts index 4d97440bc..9b02340e2 100644 --- a/core/src/telemetry/tracing.ts +++ b/core/src/telemetry/tracing.ts @@ -25,6 +25,7 @@ import {LlmRequest} from '../models/llm_request.js'; import {LlmResponse} from '../models/llm_response.js'; import {BaseTool} from '../tools/base_tool.js'; import {version} from '../version.js'; +import {TokenUsage} from './token_usage.js'; const GEN_AI_AGENT_DESCRIPTION = 'gen_ai.agent.description'; const GEN_AI_AGENT_NAME = 'gen_ai.agent.name'; @@ -247,19 +248,7 @@ export function traceCallLlm({ shouldAddRequestResponseToSpans() ? safeJsonSerialize(llmResponse) : '{}', ); - if (llmResponse.usageMetadata) { - span.setAttribute( - 'gen_ai.usage.input_tokens', - llmResponse.usageMetadata.promptTokenCount || 0, - ); - } - - if (llmResponse.usageMetadata?.candidatesTokenCount) { - span.setAttribute( - 'gen_ai.usage.output_tokens', - llmResponse.usageMetadata.candidatesTokenCount, - ); - } + span.setAttributes(new TokenUsage(llmResponse.usageMetadata).toAttributes()); if (llmResponse.finishReason) { // Convert enum to lowercase string array diff --git a/core/src/utils/error_utils.ts b/core/src/utils/error_utils.ts new file mode 100644 index 000000000..60728731b --- /dev/null +++ b/core/src/utils/error_utils.ts @@ -0,0 +1,38 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * Derives the `error.type` label for a failure. + * + * Prefers, in order: a pre-classified `errorType` carried by the error; the + * HTTP status of a `@google/genai` `ApiError` (the SDK reports every failed + * request through that one class, so only the status tells a 429 from a 400); + * and finally the error's name, falling back to its class name when the name + * has been blanked out. + * + * The `ApiError` case is matched on the shape of the error rather than with + * `instanceof`: two copies of `@google/genai` can coexist in one dependency + * tree, and an error raised by one copy is not an `instanceof` the class of + * the other. + * + * @param error The thrown value to classify. Anything can be thrown in + * JavaScript, so this takes `unknown` rather than `Error`. + * @returns The `error.type` attribute value. + */ +export function resolveErrorType(error: unknown): string { + if (typeof error === 'object' && error !== null) { + if ('errorType' in error && typeof error.errorType === 'string') { + return error.errorType; + } + if ('status' in error && typeof error.status === 'number') { + return String(error.status); + } + } + if (error instanceof Error) { + return error.name || error.constructor.name; + } + return String(error); +} diff --git a/core/test/telemetry/token_usage_test.ts b/core/test/telemetry/token_usage_test.ts new file mode 100644 index 000000000..3da0949fc --- /dev/null +++ b/core/test/telemetry/token_usage_test.ts @@ -0,0 +1,220 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {GenerateContentResponseUsageMetadata} from '@google/genai'; +import {describe, expect, it} from 'vitest'; + +import {TokenUsage} from '../../src/telemetry/token_usage.js'; + +// Spelled out rather than imported: these keys are a wire contract shared +// with adk-python, so the test must fail if the module renames one. +const INPUT_TOKENS = 'gen_ai.usage.input_tokens'; +const OUTPUT_TOKENS = 'gen_ai.usage.output_tokens'; +const CACHE_READ_INPUT_TOKENS = 'gen_ai.usage.cache_read.input_tokens'; +const REASONING_OUTPUT_TOKENS = 'gen_ai.usage.reasoning.output_tokens'; +const SYSTEM_INSTRUCTION_TOKENS = + 'gen_ai.usage.experimental.system_instruction_tokens'; + +interface UsageMetadataWithSystemInstructionTokens extends GenerateContentResponseUsageMetadata { + systemInstructionTokens?: number; +} + +describe('TokenUsage', () => { + describe('inputTokenCount', () => { + it('sums prompt and tool use tokens', () => { + const usage = new TokenUsage({ + promptTokenCount: 10, + toolUsePromptTokenCount: 5, + }); + + expect(usage.inputTokenCount).toBe(15); + }); + + it('returns the prompt tokens when tool use tokens are undefined', () => { + const usage = new TokenUsage({ + promptTokenCount: 10, + toolUsePromptTokenCount: undefined, + }); + + expect(usage.inputTokenCount).toBe(10); + }); + + it('returns the tool use tokens when prompt tokens are undefined', () => { + const usage = new TokenUsage({ + promptTokenCount: undefined, + toolUsePromptTokenCount: 5, + }); + + expect(usage.inputTokenCount).toBe(5); + }); + + it('returns undefined when neither count is reported', () => { + const usage = new TokenUsage({ + promptTokenCount: undefined, + toolUsePromptTokenCount: undefined, + }); + + expect(usage.inputTokenCount).toBeUndefined(); + }); + + it('returns 0, not undefined, when both counts are zero', () => { + const usage = new TokenUsage({ + promptTokenCount: 0, + toolUsePromptTokenCount: 0, + }); + + expect(usage.inputTokenCount).toBe(0); + }); + + it('returns undefined when there is no usage metadata', () => { + expect(new TokenUsage(undefined).inputTokenCount).toBeUndefined(); + }); + + it('returns the prompt tokens when the tool use field is absent', () => { + const usage = new TokenUsage({promptTokenCount: 10}); + + expect(usage.inputTokenCount).toBe(10); + }); + }); + + describe('outputTokenCount', () => { + it('sums candidate and reasoning tokens', () => { + const usage = new TokenUsage({ + candidatesTokenCount: 20, + thoughtsTokenCount: 8, + }); + + expect(usage.outputTokenCount).toBe(28); + }); + + it('returns the candidate tokens when reasoning tokens are undefined', () => { + const usage = new TokenUsage({ + candidatesTokenCount: 20, + thoughtsTokenCount: undefined, + }); + + expect(usage.outputTokenCount).toBe(20); + }); + + it('returns the reasoning tokens when candidate tokens are undefined', () => { + const usage = new TokenUsage({ + candidatesTokenCount: undefined, + thoughtsTokenCount: 8, + }); + + expect(usage.outputTokenCount).toBe(8); + }); + + it('returns undefined when neither count is reported', () => { + const usage = new TokenUsage({ + candidatesTokenCount: undefined, + thoughtsTokenCount: undefined, + }); + + expect(usage.outputTokenCount).toBeUndefined(); + }); + + it('returns 0, not undefined, when both counts are zero', () => { + const usage = new TokenUsage({ + candidatesTokenCount: 0, + thoughtsTokenCount: 0, + }); + + expect(usage.outputTokenCount).toBe(0); + }); + + it('returns undefined when there is no usage metadata', () => { + expect(new TokenUsage(undefined).outputTokenCount).toBeUndefined(); + }); + }); + + describe('toAttributes', () => { + it('emits every attribute when every count is reported', () => { + const usage = new TokenUsage({ + promptTokenCount: 10, + toolUsePromptTokenCount: 5, + candidatesTokenCount: 20, + thoughtsTokenCount: 8, + cachedContentTokenCount: 100, + }); + + expect(usage.toAttributes()).toEqual({ + [INPUT_TOKENS]: 15, + [OUTPUT_TOKENS]: 28, + [CACHE_READ_INPUT_TOKENS]: 100, + [REASONING_OUTPUT_TOKENS]: 8, + }); + }); + + it('omits the keys whose counts are undefined', () => { + const usage = new TokenUsage({ + promptTokenCount: 10, + toolUsePromptTokenCount: undefined, + candidatesTokenCount: undefined, + thoughtsTokenCount: undefined, + cachedContentTokenCount: undefined, + }); + + const attributes = usage.toAttributes(); + + expect(attributes[INPUT_TOKENS]).toBe(10); + expect(attributes).not.toHaveProperty(OUTPUT_TOKENS); + expect(attributes).not.toHaveProperty(CACHE_READ_INPUT_TOKENS); + expect(attributes).not.toHaveProperty(REASONING_OUTPUT_TOKENS); + }); + + it('emits nothing when there is no usage metadata', () => { + expect(new TokenUsage(undefined).toAttributes()).toEqual({}); + }); + + it('emits zeros rather than dropping the keys', () => { + const usage = new TokenUsage({ + promptTokenCount: 0, + toolUsePromptTokenCount: 0, + candidatesTokenCount: 0, + thoughtsTokenCount: 0, + cachedContentTokenCount: 0, + }); + + expect(usage.toAttributes()).toEqual({ + [INPUT_TOKENS]: 0, + [OUTPUT_TOKENS]: 0, + [CACHE_READ_INPUT_TOKENS]: 0, + [REASONING_OUTPUT_TOKENS]: 0, + }); + }); + + it('emits the totals when the optional breakdowns are absent', () => { + const usage = new TokenUsage({ + promptTokenCount: 10, + candidatesTokenCount: 20, + }); + + const attributes = usage.toAttributes(); + + expect(attributes[INPUT_TOKENS]).toBe(10); + expect(attributes[OUTPUT_TOKENS]).toBe(20); + }); + + it('emits the system instruction tokens the SDK does not declare', () => { + const metadata: UsageMetadataWithSystemInstructionTokens = { + promptTokenCount: 10, + systemInstructionTokens: 7, + }; + + expect(new TokenUsage(metadata).toAttributes()).toEqual({ + [INPUT_TOKENS]: 10, + [SYSTEM_INSTRUCTION_TOKENS]: 7, + }); + }); + + it('omits the system instruction tokens when the backend omits them', () => { + const attributes = new TokenUsage({promptTokenCount: 10}).toAttributes(); + + expect(attributes).not.toHaveProperty(SYSTEM_INSTRUCTION_TOKENS); + }); + }); +}); diff --git a/core/test/telemetry/tracing_test.ts b/core/test/telemetry/tracing_test.ts index a150f1a63..ae46ce168 100644 --- a/core/test/telemetry/tracing_test.ts +++ b/core/test/telemetry/tracing_test.ts @@ -288,5 +288,46 @@ describe('Telemetry Tracing Functions', () => { expect.anything(), ); }); + + it('should set every token usage attribute the response reports', () => { + vi.mocked(trace.getActiveSpan).mockReturnValue(mockSpan); + const llmResponse: LlmResponse = { + ...mockLlmResponse, + usageMetadata: { + promptTokenCount: 10, + toolUsePromptTokenCount: 5, + candidatesTokenCount: 20, + thoughtsTokenCount: 8, + cachedContentTokenCount: 100, + }, + }; + + traceCallLlm({ + invocationContext: mockInvocationContext, + eventId: 'test-event-id', + llmRequest: mockLlmRequest, + llmResponse, + }); + + expect(mockSpan.setAttributes).toHaveBeenCalledWith({ + 'gen_ai.usage.input_tokens': 15, + 'gen_ai.usage.output_tokens': 28, + 'gen_ai.usage.cache_read.input_tokens': 100, + 'gen_ai.usage.reasoning.output_tokens': 8, + }); + }); + + it('should set no token usage attribute when the response reports none', () => { + vi.mocked(trace.getActiveSpan).mockReturnValue(mockSpan); + + traceCallLlm({ + invocationContext: mockInvocationContext, + eventId: 'test-event-id', + llmRequest: mockLlmRequest, + llmResponse: mockLlmResponse, + }); + + expect(mockSpan.setAttributes).toHaveBeenCalledWith({}); + }); }); }); diff --git a/core/test/utils/error_utils_test.ts b/core/test/utils/error_utils_test.ts new file mode 100644 index 000000000..bfcc22b5c --- /dev/null +++ b/core/test/utils/error_utils_test.ts @@ -0,0 +1,70 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {ApiError} from '@google/genai'; +import {describe, expect, it} from 'vitest'; + +import {resolveErrorType} from '../../src/utils/error_utils.js'; + +class ClassifiedError extends Error { + constructor( + message: string, + readonly errorType: string, + ) { + super(message); + this.name = 'ClassifiedError'; + } +} + +describe('resolveErrorType', () => { + it('prefers an error type the error classified itself with', () => { + expect( + resolveErrorType(new ClassifiedError('boom', 'MCP_TOOL_ERROR')), + ).toBe('MCP_TOOL_ERROR'); + }); + + it('reports the HTTP status of a genai API error', () => { + const error = new ApiError({message: 'rate limited', status: 429}); + + expect(resolveErrorType(error)).toBe('429'); + }); + + it('prefers a classified error type over an HTTP status', () => { + const error: Error & {errorType?: string} = new ApiError({ + message: 'rate limited', + status: 429, + }); + error.errorType = 'QUOTA_EXCEEDED'; + + expect(resolveErrorType(error)).toBe('QUOTA_EXCEEDED'); + }); + + it('ignores an error type that is not a string', () => { + const error: Error & {errorType?: number} = new TypeError('boom'); + error.errorType = 503; + + expect(resolveErrorType(error)).toBe('TypeError'); + }); + + it('falls back to the error name', () => { + expect(resolveErrorType(new TypeError('not a function'))).toBe('TypeError'); + }); + + it('falls back to the class name when the name has been blanked out', () => { + const error = new TypeError('not a function'); + error.name = ''; + + expect(resolveErrorType(error)).toBe('TypeError'); + }); + + it('stringifies a thrown value that is not an object', () => { + expect(resolveErrorType('just a string')).toBe('just a string'); + }); + + it('stringifies a thrown null', () => { + expect(resolveErrorType(null)).toBe('null'); + }); +});