feat(vercel-ai): Add rerank support and fix token attribute mapping (#19144)

Author: sergical

packages/core/src/tracing/ai/gen-ai-attributes.ts

@@ -234,6 +234,11 @@ export const GEN_AI_EMBED_DO_EMBED_OPERATION_ATTRIBUTE = 'gen_ai.embed';
  */
 export const GEN_AI_EMBED_MANY_DO_EMBED_OPERATION_ATTRIBUTE = 'gen_ai.embed_many';
 
+/**
+ * The span operation name for reranking
+ */
+export const GEN_AI_RERANK_DO_RERANK_OPERATION_ATTRIBUTE = 'gen_ai.rerank';
+
 /**
  * The span operation name for executing a tool
  */

packages/core/src/tracing/vercel-ai/constants.ts

@@ -12,6 +12,7 @@ export const INVOKE_AGENT_OPS = new Set([
   'ai.streamObject',
   'ai.embed',
   'ai.embedMany',
+  'ai.rerank',
 ]);
 
 export const GENERATE_CONTENT_OPS = new Set([
@@ -22,3 +23,5 @@ export const GENERATE_CONTENT_OPS = new Set([
 ]);
 
 export const EMBEDDINGS_OPS = new Set(['ai.embed.doEmbed', 'ai.embedMany.doEmbed']);
+
+export const RERANK_OPS = new Set(['ai.rerank.doRerank']);

packages/core/src/tracing/vercel-ai/index.ts

@@ -19,7 +19,7 @@ import {
   GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
   GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
-import { EMBEDDINGS_OPS, GENERATE_CONTENT_OPS, INVOKE_AGENT_OPS, toolCallSpanMap } from './constants';
+import { EMBEDDINGS_OPS, GENERATE_CONTENT_OPS, INVOKE_AGENT_OPS, RERANK_OPS, toolCallSpanMap } from './constants';
 import type { TokenSummary } from './types';
 import {
   accumulateTokensForParent,
@@ -70,6 +70,9 @@ function mapVercelAiOperationName(operationName: string): string {
   if (EMBEDDINGS_OPS.has(operationName)) {
     return 'embeddings';
   }
+  if (RERANK_OPS.has(operationName)) {
+    return 'rerank';
+  }
   if (operationName === 'ai.toolCall') {
     return 'execute_tool';
   }
@@ -149,6 +152,13 @@ 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);
 
+  // Parent spans (ai.streamText, ai.streamObject, etc.) use inputTokens/outputTokens instead of promptTokens/completionTokens
+  renameAttributeKey(attributes, 'ai.usage.inputTokens', GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE);
+  renameAttributeKey(attributes, 'ai.usage.outputTokens', GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE);
+
+  // AI SDK uses avgOutputTokensPerSecond, map to our expected attribute name
+  renameAttributeKey(attributes, 'ai.response.avgOutputTokensPerSecond', 'ai.response.avgCompletionTokensPerSecond');
+
   // Input tokens is the sum of prompt tokens and cached input tokens
   if (
     typeof attributes[GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE] === 'number' &&
@@ -290,6 +300,9 @@ function processGenerateSpan(span: Span, name: string, attributes: SpanAttribute
       case 'ai.embedMany.doEmbed':
         span.updateName(`embed_many ${modelId}`);
         break;
+      case 'ai.rerank.doRerank':
+        span.updateName(`rerank ${modelId}`);
+        break;
     }
   }
 }

packages/core/src/tracing/vercel-ai/utils.ts

@@ -9,6 +9,7 @@ import {
   GEN_AI_INPUT_MESSAGES_ATTRIBUTE,
   GEN_AI_INPUT_MESSAGES_ORIGINAL_LENGTH_ATTRIBUTE,
   GEN_AI_INVOKE_AGENT_OPERATION_ATTRIBUTE,
+  GEN_AI_RERANK_DO_RERANK_OPERATION_ATTRIBUTE,
   GEN_AI_STREAM_OBJECT_DO_STREAM_OPERATION_ATTRIBUTE,
   GEN_AI_STREAM_TEXT_DO_STREAM_OPERATION_ATTRIBUTE,
   GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE,
@@ -190,6 +191,7 @@ export function getSpanOpFromName(name: string): string | undefined {
     case 'ai.streamObject':
     case 'ai.embed':
     case 'ai.embedMany':
+    case 'ai.rerank':
       return GEN_AI_INVOKE_AGENT_OPERATION_ATTRIBUTE;
     case 'ai.generateText.doGenerate':
       return GEN_AI_GENERATE_TEXT_DO_GENERATE_OPERATION_ATTRIBUTE;
@@ -203,6 +205,8 @@ export function getSpanOpFromName(name: string): string | undefined {
       return GEN_AI_EMBED_DO_EMBED_OPERATION_ATTRIBUTE;
     case 'ai.embedMany.doEmbed':
       return GEN_AI_EMBED_MANY_DO_EMBED_OPERATION_ATTRIBUTE;
+    case 'ai.rerank.doRerank':
+      return GEN_AI_RERANK_DO_RERANK_OPERATION_ATTRIBUTE;
     case 'ai.toolCall':
       return GEN_AI_EXECUTE_TOOL_OPERATION_ATTRIBUTE;
     default:

packages/core/test/lib/tracing/vercel-ai-parent-tokens.test.ts

@@ -0,0 +1,75 @@
+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 parent span token attributes', () => {
+  it('should map ai.usage.inputTokens to gen_ai.usage.input_tokens', () => {
+    const options = getDefaultTestClientOptions({ tracesSampleRate: 1.0 });
+    const client = new TestClient(options);
+    client.init();
+    addVercelAiProcessors(client);
+
+    const mockSpan: SpanJSON = {
+      description: 'ai.streamText',
+      span_id: 'test-span-id',
+      trace_id: 'test-trace-id',
+      start_timestamp: 1000,
+      timestamp: 2000,
+      origin: 'auto.vercelai.otel',
+      data: {
+        'ai.usage.inputTokens': 100,
+        'ai.usage.outputTokens': 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(100);
+    expect(processedEvent?.spans?.[0]?.data?.['gen_ai.usage.output_tokens']).toBe(50);
+    // Original attributes should be renamed to vercel.ai.* namespace
+    expect(processedEvent?.spans?.[0]?.data?.['ai.usage.inputTokens']).toBeUndefined();
+    expect(processedEvent?.spans?.[0]?.data?.['ai.usage.outputTokens']).toBeUndefined();
+  });
+
+  it('should map ai.response.avgOutputTokensPerSecond to ai.response.avgCompletionTokensPerSecond', () => {
+    const options = getDefaultTestClientOptions({ tracesSampleRate: 1.0 });
+    const client = new TestClient(options);
+    client.init();
+    addVercelAiProcessors(client);
+
+    const mockSpan: SpanJSON = {
+      description: 'ai.streamText.doStream',
+      span_id: 'test-span-id',
+      trace_id: 'test-trace-id',
+      start_timestamp: 1000,
+      timestamp: 2000,
+      origin: 'auto.vercelai.otel',
+      data: {
+        'ai.response.avgOutputTokensPerSecond': 25.5,
+      },
+    };
+
+    const event = {
+      type: 'transaction' as const,
+      spans: [mockSpan],
+    };
+
+    const eventProcessor = client['_eventProcessors'].find(processor => processor.id === 'VercelAiEventProcessor');
+    expect(eventProcessor).toBeDefined();
+
+    const processedEvent = eventProcessor!(event, {});
+
+    // Should be renamed to match the expected attribute name
+    expect(processedEvent?.spans?.[0]?.data?.['vercel.ai.response.avgCompletionTokensPerSecond']).toBe(25.5);
+    expect(processedEvent?.spans?.[0]?.data?.['ai.response.avgOutputTokensPerSecond']).toBeUndefined();
+  });
+});

packages/core/test/lib/tracing/vercel-ai-rerank.test.ts

@@ -0,0 +1,14 @@
+import { describe, expect, it } from 'vitest';
+import { getSpanOpFromName } from '../../../src/tracing/vercel-ai/utils';
+
+describe('vercel-ai rerank support', () => {
+  describe('getSpanOpFromName', () => {
+    it('should map ai.rerank to gen_ai.invoke_agent', () => {
+      expect(getSpanOpFromName('ai.rerank')).toBe('gen_ai.invoke_agent');
+    });
+
+    it('should map ai.rerank.doRerank to gen_ai.rerank', () => {
+      expect(getSpanOpFromName('ai.rerank.doRerank')).toBe('gen_ai.rerank');
+    });
+  });
+});

packages/node/src/integrations/tracing/vercelai/instrumentation.ts

@@ -26,6 +26,7 @@ const INSTRUMENTED_METHODS = [
   'streamObject',
   'embed',
   'embedMany',
+  'rerank',
 ] as const;
 
 interface MethodFirstArg extends Record<string, unknown> {
@@ -263,15 +264,21 @@ export class SentryVercelAiInstrumentation extends InstrumentationBase {
     if (Object.prototype.toString.call(moduleExports) === '[object Module]') {
       // In ESM we take the usual route and just replace the exports we want to instrument
       for (const method of INSTRUMENTED_METHODS) {
-        moduleExports[method] = generatePatch(moduleExports[method]);
+        // Skip methods that don't exist in this version of the AI SDK (e.g., rerank was added in v6)
+        if (moduleExports[method] != null) {
+          moduleExports[method] = generatePatch(moduleExports[method]);
+        }
       }
 
       return moduleExports;
     } else {
       // In CJS we can't replace the exports in the original module because they
       // don't have setters, so we create a new object with the same properties
       const patchedModuleExports = INSTRUMENTED_METHODS.reduce((acc, curr) => {
-        acc[curr] = generatePatch(moduleExports[curr]);
+        // Skip methods that don't exist in this version of the AI SDK (e.g., rerank was added in v6)
+        if (moduleExports[curr] != null) {
+          acc[curr] = generatePatch(moduleExports[curr]);
+        }
         return acc;
       }, {} as PatchedModuleExports);