diff --git a/packages/core/src/tracing/vercel-ai/index.ts b/packages/core/src/tracing/vercel-ai/index.ts index 93be1ca33423..e64b4b1a9cbf 100644 --- a/packages/core/src/tracing/vercel-ai/index.ts +++ b/packages/core/src/tracing/vercel-ai/index.ts @@ -119,6 +119,15 @@ function processEndedVercelAiSpan(span: SpanJSON): void { renameAttributeKey(attributes, AI_USAGE_PROMPT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE); renameAttributeKey(attributes, AI_USAGE_CACHED_INPUT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_INPUT_TOKENS_CACHED_ATTRIBUTE); + // Input tokens is the sum of prompt tokens and cached input tokens + if ( + typeof attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] === 'number' && + typeof attributes[GEN_AI_USAGE_INPUT_TOKENS_CACHED_ATTRIBUTE] === 'number' + ) { + attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] = + attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] + attributes[GEN_AI_USAGE_INPUT_TOKENS_CACHED_ATTRIBUTE]; + } + if ( typeof attributes[GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE] === 'number' && typeof attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] === 'number' diff --git a/packages/core/test/lib/tracing/vercel-ai-cached-tokens.test.ts b/packages/core/test/lib/tracing/vercel-ai-cached-tokens.test.ts new file mode 100644 index 000000000000..7e85121b9e92 --- /dev/null +++ b/packages/core/test/lib/tracing/vercel-ai-cached-tokens.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it } from 'vitest'; +import { addVercelAiProcessors } from '../../../src/tracing/vercel-ai'; +import type { SpanJSON } from '../../../src/types-hoist/span'; +import { getDefaultTestClientOptions, TestClient } from '../../mocks/client'; + +describe('vercel-ai cached tokens', () => { + it('should add cached input tokens to total input tokens', () => { + const options = getDefaultTestClientOptions({ tracesSampleRate: 1.0 }); + const client = new TestClient(options); + client.init(); + addVercelAiProcessors(client); + + const mockSpan: SpanJSON = { + description: 'test', + span_id: 'test-span-id', + trace_id: 'test-trace-id', + start_timestamp: 1000, + timestamp: 2000, + origin: 'auto.vercelai.otel', + data: { + 'ai.usage.promptTokens': 100, + 'ai.usage.cachedInputTokens': 50, + }, + }; + + const event = { + type: 'transaction' as const, + spans: [mockSpan], + }; + + const eventProcessor = client['_eventProcessors'].find(processor => processor.id === 'VercelAiEventProcessor'); + expect(eventProcessor).toBeDefined(); + + const processedEvent = eventProcessor!(event, {}); + + expect(processedEvent?.spans?.[0]?.data?.['gen_ai.usage.input_tokens']).toBe(150); + expect(processedEvent?.spans?.[0]?.data?.['gen_ai.usage.input_tokens.cached']).toBe(50); + }); +});