diff --git a/packages/server-utils/src/ai/google-genai/index.ts b/packages/server-utils/src/ai/google-genai/index.ts index e8c002a921f8..38d640a56386 100644 --- a/packages/server-utils/src/ai/google-genai/index.ts +++ b/packages/server-utils/src/ai/google-genai/index.ts @@ -39,6 +39,7 @@ import type { ContentListUnion, Message, PartListUnion } from './utils'; import { candidatesToMessageParts, contentUnionToMessages, + getInputTokens, setOutputMessagesAttribute, systemInstructionToText, } from './utils'; @@ -206,7 +207,7 @@ export function addResponseAttributes(span: Span, response: GoogleGenAIResponse, const hasOutput = typeof usage.candidatesTokenCount === 'number' || typeof usage.thoughtsTokenCount === 'number'; span.setAttributes( getTokenUsageAttributes( - usage.promptTokenCount, + getInputTokens(usage.promptTokenCount, usage.toolUsePromptTokenCount), hasOutput ? (usage.candidatesTokenCount ?? 0) + (usage.thoughtsTokenCount ?? 0) : undefined, undefined, usage.cachedContentTokenCount, diff --git a/packages/server-utils/src/ai/google-genai/streaming.ts b/packages/server-utils/src/ai/google-genai/streaming.ts index 3a2f720f2628..13df6febb1ef 100644 --- a/packages/server-utils/src/ai/google-genai/streaming.ts +++ b/packages/server-utils/src/ai/google-genai/streaming.ts @@ -3,7 +3,7 @@ import type { Span } from '@sentry/core'; import { endStreamSpan } from '../core/utils'; import type { GoogleGenAIResponse } from './types'; import type { MessagePart } from './utils'; -import { candidatesToMessageParts, setOutputMessagesAttribute } from './utils'; +import { candidatesToMessageParts, getInputTokens, setOutputMessagesAttribute } from './utils'; /** * State object used to accumulate information from a stream of Google GenAI events. @@ -17,8 +17,10 @@ interface StreamingState { responseId?: string; /** The model name. */ responseModel?: string; - /** Number of prompt/input tokens used. */ + /** Number of prompt/input tokens used, including the tool-use prompt tokens. */ promptTokens?: number; + promptTokenCount?: number; + toolUsePromptTokenCount?: number; /** Number of completion/output tokens used. */ completionTokens?: number; candidateTokens?: number; @@ -62,7 +64,11 @@ function handleResponseMetadata(chunk: GoogleGenAIResponse, state: StreamingStat const usage = chunk.usageMetadata; if (usage) { - if (typeof usage.promptTokenCount === 'number') state.promptTokens = usage.promptTokenCount; + if (typeof usage.promptTokenCount === 'number') state.promptTokenCount = usage.promptTokenCount; + if (typeof usage.toolUsePromptTokenCount === 'number') { + state.toolUsePromptTokenCount = usage.toolUsePromptTokenCount; + } + state.promptTokens = getInputTokens(state.promptTokenCount, state.toolUsePromptTokenCount); if (typeof usage.candidatesTokenCount === 'number') state.candidateTokens = usage.candidatesTokenCount; if (typeof usage.thoughtsTokenCount === 'number') state.reasoningOutputTokens = usage.thoughtsTokenCount; if (state.candidateTokens !== undefined || state.reasoningOutputTokens !== undefined) { diff --git a/packages/server-utils/src/ai/google-genai/utils.ts b/packages/server-utils/src/ai/google-genai/utils.ts index 72d1a2c143f7..d7b893e86951 100644 --- a/packages/server-utils/src/ai/google-genai/utils.ts +++ b/packages/server-utils/src/ai/google-genai/utils.ts @@ -205,3 +205,16 @@ export function systemInstructionToText(systemInstruction: unknown): string | un return texts.length ? texts.join('\n') : undefined; } + +/** + * Input tokens of a Gemini response. `toolUsePromptTokenCount` (results of built-in tools such as Google Search, + * URL context or code execution, fed back to the model) is reported apart from `promptTokenCount`, but + * `totalTokenCount` includes it, so it is part of the input. + */ +export function getInputTokens(promptTokenCount?: number, toolUsePromptTokenCount?: number): number | undefined { + if (typeof promptTokenCount !== 'number' && typeof toolUsePromptTokenCount !== 'number') return undefined; + return ( + (typeof promptTokenCount === 'number' ? promptTokenCount : 0) + + (typeof toolUsePromptTokenCount === 'number' ? toolUsePromptTokenCount : 0) + ); +} diff --git a/packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts b/packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts index 924dc7b8df18..6fc3a2994dcc 100644 --- a/packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts +++ b/packages/server-utils/test/ai/lib/tracing/google-genai-streaming.test.ts @@ -1,6 +1,11 @@ import { describe, expect, it } from 'vitest'; import type { Span } from '@sentry/core'; -import { GEN_AI_RESPONSE_TOOL_CALLS } from '@sentry/conventions/attributes'; +import { + GEN_AI_RESPONSE_TOOL_CALLS, + GEN_AI_USAGE_INPUT_TOKENS, + GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_USAGE_TOTAL_TOKENS, +} from '@sentry/conventions/attributes'; import { addResponseAttributes } from '../../../../src/ai/google-genai/index'; import { instrumentStream } from '../../../../src/ai/google-genai/streaming'; import type { ContentPart, GoogleGenAIResponse } from '../../../../src/ai/google-genai/types'; @@ -126,3 +131,56 @@ describe('instrumentStream (google-genai tool calls)', () => { expect(toolCalls).toEqual([functionCall]); }); }); + +describe('google-genai tool-use prompt tokens', () => { + // Google Search, URL context and code execution results are fed back to the model and counted in + // `toolUsePromptTokenCount`, apart from `promptTokenCount`; `totalTokenCount` includes both. + const usageMetadata = { + promptTokenCount: 100, + toolUsePromptTokenCount: 40, + candidatesTokenCount: 20, + totalTokenCount: 160, + }; + + it('counts them as input tokens on a response', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes(span, { usageMetadata } as GoogleGenAIResponse, false); + + expect(attributes[GEN_AI_USAGE_INPUT_TOKENS]).toBe(140); + expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(20); + expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(160); + }); + + it('leaves the input tokens alone when there are no tool-use prompt tokens', () => { + const { span, attributes } = createMockSpan(); + + addResponseAttributes( + span, + { + usageMetadata: { promptTokenCount: 100, candidatesTokenCount: 20, totalTokenCount: 120 }, + } as GoogleGenAIResponse, + false, + ); + + expect(attributes[GEN_AI_USAGE_INPUT_TOKENS]).toBe(100); + expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(120); + }); + + it('counts them as input tokens on a stream, also when a later chunk only repeats one of the counts', async () => { + const { span, attributes } = createMockSpan(); + + const chunks = [ + chunkWithParts([{ text: 'Hi' }], { usageMetadata: { promptTokenCount: 100, toolUsePromptTokenCount: 40 } }), + chunkWithParts([{ text: ' there' }], { + usageMetadata: { candidatesTokenCount: 20, totalTokenCount: 160 }, + }), + ]; + + await drain(instrumentStream(streamOf(chunks), span, false)); + + expect(attributes[GEN_AI_USAGE_INPUT_TOKENS]).toBe(140); + expect(attributes[GEN_AI_USAGE_OUTPUT_TOKENS]).toBe(20); + expect(attributes[GEN_AI_USAGE_TOTAL_TOKENS]).toBe(160); + }); +});