Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions packages/server-utils/src/ai/vercel-ai/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
14 changes: 14 additions & 0 deletions packages/server-utils/src/ai/vercel-ai/vercel-ai-attributes.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
};
}

/**
Expand Down
Original file line number Diff line number Diff line change
@@ -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();
});
});