diff --git a/packages/server-utils/src/ai/vercel-ai/index.ts b/packages/server-utils/src/ai/vercel-ai/index.ts index fe7b63c6a64b..3633c40699ce 100644 --- a/packages/server-utils/src/ai/vercel-ai/index.ts +++ b/packages/server-utils/src/ai/vercel-ai/index.ts @@ -574,6 +574,21 @@ export function getProviderMetadataAttributes(providerMetadata: unknown): Record setAttributeIfDefined(attributes, 'gen_ai.usage.input_tokens.cache_miss', metadata.deepseek.promptCacheMissTokens); } + // Google (v5 uses 'google', v6 Vertex AI uses 'vertex'). Gemini reports its reasoning ("thoughts") + // tokens separately from the candidate output count, so the SDK's `outputTokens` covers only the + // visible answer. `gen_ai.usage.output_tokens` must include reasoning tokens, so recompute it (and + // the total) from the raw usageMetadata. Deriving output from `candidatesTokenCount + thoughtsTokenCount` + // rather than adding onto the existing value keeps this correct even if a future SDK version already + // folds reasoning into `outputTokens`. + const googleUsage = (metadata.google ?? metadata.vertex)?.usageMetadata; + if (googleUsage && typeof googleUsage.thoughtsTokenCount === 'number' && googleUsage.thoughtsTokenCount > 0) { + setAttributeIfDefined(attributes, GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, googleUsage.thoughtsTokenCount); + if (typeof googleUsage.candidatesTokenCount === 'number') { + attributes[GEN_AI_USAGE_OUTPUT_TOKENS] = googleUsage.candidatesTokenCount + googleUsage.thoughtsTokenCount; + } + setAttributeIfDefined(attributes, GEN_AI_USAGE_TOTAL_TOKENS, googleUsage.totalTokenCount); + } + return attributes; } diff --git a/packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts b/packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts index 62d89f50c17c..63547b1f058c 100644 --- a/packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts +++ b/packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts @@ -446,6 +446,20 @@ export interface GoogleGenerativeAIProviderMetadata { * @see https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/configure-safety-filters */ safetyRatings?: null | unknown; + + /** + * Raw token usage returned by the Gemini API. Reasoning ("thoughts") tokens are reported here in + * `thoughtsTokenCount`, separately from the candidate output count, so they have to be added back + * into `gen_ai.usage.output_tokens`. + * @see https://ai.google.dev/api/generate-content#UsageMetadata + * @see https://github.com/vercel/ai/blob/main/packages/google/src/google-language-model.ts + */ + usageMetadata?: null | { + promptTokenCount?: number; + candidatesTokenCount?: number; + thoughtsTokenCount?: number; + totalTokenCount?: number; + }; } /** diff --git a/packages/server-utils/test/ai/lib/tracing/vercel-ai-reasoning-tokens.test.ts b/packages/server-utils/test/ai/lib/tracing/vercel-ai-reasoning-tokens.test.ts new file mode 100644 index 000000000000..7f9666da607b --- /dev/null +++ b/packages/server-utils/test/ai/lib/tracing/vercel-ai-reasoning-tokens.test.ts @@ -0,0 +1,110 @@ +import { describe, expect, it } from 'vitest'; +import type { SpanJSON } from '@sentry/core'; +import { addVercelAiProcessors, getProviderMetadataAttributes } from '../../../../src/ai/vercel-ai'; +import { getDefaultTestClientOptions, TestClient } from '../../../mocks/client'; + +function processSpan(data: SpanJSON['data']): SpanJSON { + const options = getDefaultTestClientOptions({ tracesSampleRate: 1.0 }); + const client = new TestClient(options); + client.init(); + addVercelAiProcessors(client); + + const mockSpan: SpanJSON = { + description: 'ai.generateText.doGenerate', + span_id: 'test-span-id', + trace_id: 'test-trace-id', + start_timestamp: 1000, + timestamp: 2000, + origin: 'auto.vercelai.otel', + data, + }; + + const event = { + type: 'transaction' as const, + spans: [mockSpan], + }; + + const eventProcessor = client['_eventProcessors'].find(processor => processor.id === 'VercelAiEventProcessor'); + expect(eventProcessor).toBeDefined(); + + return eventProcessor!(event, {})!.spans![0]!; +} + +// Real usage seen from a Gemini reasoning model: the candidate output is small but the model spent +// most of its budget on hidden reasoning ("thoughts"). The AI SDK reports the candidate count as +// `outputTokens` and exposes the reasoning count only through `providerMetadata.google.usageMetadata`. +const GEMINI_REASONING_METADATA = { + google: { + groundingMetadata: null, + safetyRatings: null, + usageMetadata: { + promptTokenCount: 14, + candidatesTokenCount: 1, + thoughtsTokenCount: 100, + totalTokenCount: 115, + }, + }, +}; + +describe('vercel-ai Gemini reasoning tokens', () => { + it('includes reasoning (thoughts) tokens in output and total for a Gemini doGenerate span', () => { + const span = processSpan({ + 'ai.usage.promptTokens': 14, + 'ai.usage.completionTokens': 1, + 'ai.response.providerMetadata': JSON.stringify(GEMINI_REASONING_METADATA), + }); + + expect(span.data?.['gen_ai.usage.input_tokens']).toBe(14); + // output must include the 100 reasoning tokens, not just the single candidate token + expect(span.data?.['gen_ai.usage.output_tokens']).toBe(101); + expect(span.data?.['gen_ai.usage.reasoning.output_tokens']).toBe(100); + // total is the real Gemini total, not input + candidate-only output (which would be 15) + expect(span.data?.['gen_ai.usage.total_tokens']).toBe(115); + }); + + it('derives reasoning-inclusive output/total from google usageMetadata', () => { + const attributes = getProviderMetadataAttributes(GEMINI_REASONING_METADATA); + + expect(attributes['gen_ai.usage.output_tokens']).toBe(101); + expect(attributes['gen_ai.usage.reasoning.output_tokens']).toBe(100); + expect(attributes['gen_ai.usage.total_tokens']).toBe(115); + }); + + it('reads the v6 vertex provider metadata key too', () => { + const attributes = getProviderMetadataAttributes({ + vertex: { + usageMetadata: { + promptTokenCount: 20, + candidatesTokenCount: 5, + thoughtsTokenCount: 40, + totalTokenCount: 65, + }, + }, + }); + + expect(attributes['gen_ai.usage.output_tokens']).toBe(45); + expect(attributes['gen_ai.usage.reasoning.output_tokens']).toBe(40); + expect(attributes['gen_ai.usage.total_tokens']).toBe(65); + }); + + it('leaves non-reasoning Gemini responses untouched', () => { + const span = processSpan({ + 'ai.usage.promptTokens': 30, + 'ai.usage.completionTokens': 12, + 'ai.response.providerMetadata': JSON.stringify({ + google: { + usageMetadata: { + promptTokenCount: 30, + candidatesTokenCount: 12, + totalTokenCount: 42, + }, + }, + }), + }); + + expect(span.data?.['gen_ai.usage.input_tokens']).toBe(30); + expect(span.data?.['gen_ai.usage.output_tokens']).toBe(12); + expect(span.data?.['gen_ai.usage.total_tokens']).toBe(42); + expect(span.data?.['gen_ai.usage.reasoning.output_tokens']).toBeUndefined(); + }); +});