make simple

Author: nicohrubec

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

@@ -5,22 +5,18 @@ import type { ToolCallSpanContext } from './types';
 // without keeping full Span objects (and their potentially large attributes) alive.
 export const toolCallSpanContextMap = new Map<string, ToolCallSpanContext>();
 
-// Operation sets for efficient mapping to OpenTelemetry semantic convention values
-export const INVOKE_AGENT_OPS = new Set(['ai.generateText', 'ai.streamText', 'ai.generateObject', 'ai.streamObject']);
-
-export const GENERATE_CONTENT_OPS = new Set([
-  'ai.generateText.doGenerate',
-  'ai.streamText.doStream',
-  'ai.generateObject.doGenerate',
-  'ai.streamObject.doStream',
+/** Maps Vercel AI span names to standardized OpenTelemetry operation names. */
+export const SPAN_TO_OPERATION_NAME = new Map<string, string>([
+  ['ai.generateText', 'invoke_agent'],
+  ['ai.streamText', 'invoke_agent'],
+  ['ai.generateObject', 'invoke_agent'],
+  ['ai.streamObject', 'invoke_agent'],
+  ['ai.generateText.doGenerate', 'generate_content'],
+  ['ai.streamText.doStream', 'generate_content'],
+  ['ai.generateObject.doGenerate', 'generate_content'],
+  ['ai.streamObject.doStream', 'generate_content'],
+  ['ai.embed.doEmbed', 'embeddings'],
+  ['ai.embedMany.doEmbed', 'embeddings'],
+  ['ai.rerank.doRerank', 'rerank'],
+  ['ai.toolCall', 'execute_tool'],
 ]);
-
-export const EMBEDDINGS_OPS = new Set(['ai.embed.doEmbed', 'ai.embedMany.doEmbed']);
-
-export const RERANK_OPS = new Set(['ai.rerank.doRerank']);
-
-export const DO_SPAN_NAME_PREFIX: Record<string, string> = {
-  'ai.embed.doEmbed': 'embeddings',
-  'ai.embedMany.doEmbed': 'embeddings',
-  'ai.rerank.doRerank': 'rerank',
-};

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

@@ -22,21 +22,13 @@ import {
   GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
   GEN_AI_USAGE_TOTAL_TOKENS_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
-import {
-  DO_SPAN_NAME_PREFIX,
-  EMBEDDINGS_OPS,
-  GENERATE_CONTENT_OPS,
-  INVOKE_AGENT_OPS,
-  RERANK_OPS,
-  toolCallSpanContextMap,
-} from './constants';
+import { SPAN_TO_OPERATION_NAME, toolCallSpanContextMap } from './constants';
 import type { TokenSummary } from './types';
 import {
   accumulateTokensForParent,
   applyAccumulatedTokens,
   applyToolDescriptionsAndTokens,
   convertAvailableToolsToJsonString,
-  getSpanOpFromName,
   requestMessagesFromPrompt,
 } from './utils';
 import type { OpenAiProviderMetadata, ProviderMetadata } from './vercel-ai-attributes';
@@ -64,32 +56,6 @@ import {
   OPERATION_NAME_ATTRIBUTE,
 } from './vercel-ai-attributes';
 
-/**
- * Maps Vercel AI SDK operation names to OpenTelemetry semantic convention values
- * @see https://opentelemetry.io/docs/specs/semconv/gen-ai/gen-ai-spans/#llm-request-spans
- */
-function mapVercelAiOperationName(operationName: string): string {
-  // Top-level pipeline operations map to invoke_agent
-  if (INVOKE_AGENT_OPS.has(operationName)) {
-    return 'invoke_agent';
-  }
-  // .do* operations are the actual LLM calls
-  if (GENERATE_CONTENT_OPS.has(operationName)) {
-    return 'generate_content';
-  }
-  if (EMBEDDINGS_OPS.has(operationName)) {
-    return 'embeddings';
-  }
-  if (RERANK_OPS.has(operationName)) {
-    return 'rerank';
-  }
-  if (operationName === 'ai.toolCall') {
-    return 'execute_tool';
-  }
-  // Return the original value for unknown operations
-  return operationName;
-}
-
 /**
  * Post-process spans emitted by the Vercel AI SDK.
  * This is supposed to be used in `client.on('spanStart', ...)
@@ -314,7 +280,9 @@ function processEndedVercelAiSpan(span: SpanJSON): void {
   // Rename AI SDK attributes to standardized gen_ai attributes
   // Map operation.name to OpenTelemetry semantic convention values
   if (attributes[OPERATION_NAME_ATTRIBUTE]) {
-    const operationName = mapVercelAiOperationName(attributes[OPERATION_NAME_ATTRIBUTE] as string);
+    const operationName =
+      SPAN_TO_OPERATION_NAME.get(attributes[OPERATION_NAME_ATTRIBUTE] as string) ??
+      (attributes[OPERATION_NAME_ATTRIBUTE] as string);
     attributes[GEN_AI_OPERATION_NAME_ATTRIBUTE] = operationName;
     // eslint-disable-next-line @typescript-eslint/no-dynamic-delete
     delete attributes[OPERATION_NAME_ATTRIBUTE];
@@ -415,15 +383,17 @@ function processGenerateSpan(span: Span, name: string, attributes: SpanAttribute
   }
   span.setAttribute('ai.streaming', name.includes('stream'));
 
-  // Set the op based on the span name
-  const op = getSpanOpFromName(name);
-  if (op) {
-    span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, op);
+  // Set the op based on the operation name registry
+  const operationName = SPAN_TO_OPERATION_NAME.get(name);
+  if (operationName) {
+    span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, `gen_ai.${operationName}`);
+  } else if (name.startsWith('ai.stream')) {
+    span.setAttribute(SEMANTIC_ATTRIBUTE_SENTRY_OP, 'ai.run');
   }
 
   // For invoke_agent pipeline spans, use 'invoke_agent' as the description
   // to be consistent with other AI integrations (e.g. LangGraph)
-  if (INVOKE_AGENT_OPS.has(name)) {
+  if (operationName === 'invoke_agent') {
     if (functionId && typeof functionId === 'string') {
       span.updateName(`invoke_agent ${functionId}`);
     } else {
@@ -433,11 +403,8 @@ function processGenerateSpan(span: Span, name: string, attributes: SpanAttribute
   }
 
   const modelId = attributes[AI_MODEL_ID_ATTRIBUTE];
-  if (modelId) {
-    const doSpanPrefix = GENERATE_CONTENT_OPS.has(name) ? 'generate_content' : DO_SPAN_NAME_PREFIX[name];
-    if (doSpanPrefix) {
-      span.updateName(`${doSpanPrefix} ${modelId}`);
-    }
+  if (modelId && operationName) {
+    span.updateName(`${operationName} ${modelId}`);
   }
 }
 

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

@@ -1,13 +1,8 @@
 import type { TraceContext } from '../../types-hoist/context';
 import type { Span, SpanAttributes, SpanJSON } from '../../types-hoist/span';
 import {
-  GEN_AI_EMBEDDINGS_OPERATION_ATTRIBUTE,
-  GEN_AI_EXECUTE_TOOL_OPERATION_ATTRIBUTE,
-  GEN_AI_GENERATE_CONTENT_OPERATION_ATTRIBUTE,
   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_REQUEST_AVAILABLE_TOOLS_ATTRIBUTE,
   GEN_AI_SYSTEM_INSTRUCTIONS_ATTRIBUTE,
   GEN_AI_TOOL_DESCRIPTION_ATTRIBUTE,
@@ -279,33 +274,3 @@ export function requestMessagesFromPrompt(span: Span, attributes: SpanAttributes
     } catch {}
   }
 }
-
-/**
- * Maps a Vercel AI span name to the corresponding Sentry op.
- */
-export function getSpanOpFromName(name: string): string | undefined {
-  switch (name) {
-    case 'ai.generateText':
-    case 'ai.streamText':
-    case 'ai.generateObject':
-    case 'ai.streamObject':
-      return GEN_AI_INVOKE_AGENT_OPERATION_ATTRIBUTE;
-    case 'ai.generateText.doGenerate':
-    case 'ai.streamText.doStream':
-    case 'ai.generateObject.doGenerate':
-    case 'ai.streamObject.doStream':
-      return GEN_AI_GENERATE_CONTENT_OPERATION_ATTRIBUTE;
-    case 'ai.embed.doEmbed':
-    case 'ai.embedMany.doEmbed':
-      return GEN_AI_EMBEDDINGS_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:
-      if (name.startsWith('ai.stream')) {
-        return 'ai.run';
-      }
-      return undefined;
-  }
-}

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

@@ -1,14 +1,14 @@
 import { describe, expect, it } from 'vitest';
-import { getSpanOpFromName } from '../../../src/tracing/vercel-ai/utils';
+import { SPAN_TO_OPERATION_NAME } from '../../../src/tracing/vercel-ai/constants';
 
 describe('vercel-ai rerank support', () => {
-  describe('getSpanOpFromName', () => {
-    it('should not assign a gen_ai op to ai.rerank pipeline span', () => {
-      expect(getSpanOpFromName('ai.rerank')).toBeUndefined();
+  describe('SPAN_TO_OPERATION_NAME', () => {
+    it('should not have a mapping for ai.rerank pipeline span', () => {
+      expect(SPAN_TO_OPERATION_NAME.get('ai.rerank')).toBeUndefined();
     });
 
-    it('should map ai.rerank.doRerank to gen_ai.rerank', () => {
-      expect(getSpanOpFromName('ai.rerank.doRerank')).toBe('gen_ai.rerank');
+    it('should map ai.rerank.doRerank to rerank', () => {
+      expect(SPAN_TO_OPERATION_NAME.get('ai.rerank.doRerank')).toBe('rerank');
     });
   });
 });