fixup! fix(tracing): Add missing attributes in vercel-ai spans (#18333)

Author: isaacs

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

@@ -4,10 +4,10 @@ import type { Event } from '../../types-hoist/event';
 import type { Span, SpanAttributes, SpanAttributeValue, SpanJSON, SpanOrigin } from '../../types-hoist/span';
 import { spanToJSON } from '../../utils/spanUtils';
 import {
+  GEN_AI_REQUEST_MESSAGES_ATTRIBUTE,
   GEN_AI_USAGE_INPUT_TOKENS_CACHE_WRITE_ATTRIBUTE,
   GEN_AI_USAGE_INPUT_TOKENS_CACHED_ATTRIBUTE,
 } from '../ai/gen-ai-attributes';
-import { getTruncatedJsonString } from '../ai/utils';
 import { toolCallSpanMap } from './constants';
 import type { TokenSummary } from './types';
 import {
@@ -20,7 +20,6 @@ import type { ProviderMetadata } from './vercel-ai-attributes';
 import {
   AI_MODEL_ID_ATTRIBUTE,
   AI_MODEL_PROVIDER_ATTRIBUTE,
-  AI_PROMPT_ATTRIBUTE,
   AI_PROMPT_MESSAGES_ATTRIBUTE,
   AI_PROMPT_TOOLS_ATTRIBUTE,
   AI_RESPONSE_OBJECT_ATTRIBUTE,
@@ -36,6 +35,7 @@ import {
   AI_USAGE_CACHED_INPUT_TOKENS_ATTRIBUTE,
   AI_USAGE_COMPLETION_TOKENS_ATTRIBUTE,
   AI_USAGE_PROMPT_TOKENS_ATTRIBUTE,
+  GEN_AI_REQUEST_MODEL_ATTRIBUTE,
   GEN_AI_RESPONSE_MODEL_ATTRIBUTE,
   GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE,
   GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
@@ -136,7 +136,7 @@ function processEndedVercelAiSpan(span: SpanJSON): void {
   }
 
   // Rename AI SDK attributes to standardized gen_ai attributes
-  renameAttributeKey(attributes, AI_PROMPT_MESSAGES_ATTRIBUTE, 'gen_ai.request.messages');
+  renameAttributeKey(attributes, AI_PROMPT_MESSAGES_ATTRIBUTE, GEN_AI_REQUEST_MESSAGES_ATTRIBUTE);
   renameAttributeKey(attributes, AI_RESPONSE_TEXT_ATTRIBUTE, 'gen_ai.response.text');
   renameAttributeKey(attributes, AI_RESPONSE_TOOL_CALLS_ATTRIBUTE, 'gen_ai.response.tool_calls');
   renameAttributeKey(attributes, AI_RESPONSE_OBJECT_ATTRIBUTE, 'gen_ai.response.object');
@@ -146,7 +146,7 @@ function processEndedVercelAiSpan(span: SpanJSON): void {
   renameAttributeKey(attributes, AI_TOOL_CALL_RESULT_ATTRIBUTE, 'gen_ai.tool.output');
 
   renameAttributeKey(attributes, AI_SCHEMA_ATTRIBUTE, 'gen_ai.request.schema');
-  renameAttributeKey(attributes, AI_MODEL_ID_ATTRIBUTE, 'gen_ai.request.model');
+  renameAttributeKey(attributes, AI_MODEL_ID_ATTRIBUTE, GEN_AI_REQUEST_MODEL_ATTRIBUTE);
 
   addProviderMetadataToAttributes(attributes);
 
@@ -209,14 +209,8 @@ function processGenerateSpan(span: Span, name: string, attributes: SpanAttribute
     span.setAttribute('gen_ai.function_id', functionId);
   }
 
-  if (attributes[AI_PROMPT_ATTRIBUTE]) {
-    const truncatedPrompt = getTruncatedJsonString(attributes[AI_PROMPT_ATTRIBUTE] as string | string[]);
-    span.setAttribute('gen_ai.prompt', truncatedPrompt);
+  requestMessagesFromPrompt(span, attributes);
 
-    if (!attributes['gen_ai.request.messages'] && !attributes[AI_PROMPT_MESSAGES_ATTRIBUTE]) {
-      requestMessagesFromPrompt(span, attributes[AI_PROMPT_ATTRIBUTE]);
-    }
-  }
   if (attributes[AI_MODEL_ID_ATTRIBUTE] && !attributes[GEN_AI_RESPONSE_MODEL_ATTRIBUTE]) {
     span.setAttribute(GEN_AI_RESPONSE_MODEL_ATTRIBUTE, attributes[AI_MODEL_ID_ATTRIBUTE]);
   }

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

@@ -1,9 +1,14 @@
 import type { TraceContext } from '../../types-hoist/context';
-import type { Span, SpanJSON } from '../../types-hoist/span';
-import { GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE, GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE } from '../ai/gen-ai-attributes';
+import type { Span, SpanAttributes, SpanJSON } from '../../types-hoist/span';
+import {
+  GEN_AI_REQUEST_MESSAGES_ATTRIBUTE,
+  GEN_AI_USAGE_INPUT_TOKENS_ATTRIBUTE,
+  GEN_AI_USAGE_OUTPUT_TOKENS_ATTRIBUTE,
+} from '../ai/gen-ai-attributes';
 import { getTruncatedJsonString } from '../ai/utils';
 import { toolCallSpanMap } from './constants';
 import type { TokenSummary } from './types';
+import { AI_PROMPT_ATTRIBUTE, AI_PROMPT_MESSAGES_ATTRIBUTE } from './vercel-ai-attributes';
 
 /**
  * Accumulates token data from a span to its parent in the token accumulator map.
@@ -92,7 +97,7 @@ export function convertAvailableToolsToJsonString(tools: unknown[]): string {
 /**
  * Convert the prompt string to messages array
  */
-export function convertPromptToMessages(prompt: string): { role: string; content: string }[] | undefined {
+export function convertPromptToMessages(prompt: string): { role: string; content: string }[] {
   try {
     const p = JSON.parse(prompt);
     if (!!p && typeof p === 'object') {
@@ -105,22 +110,30 @@ export function convertPromptToMessages(prompt: string): { role: string; content
         if (typeof prompt === 'string') {
           messages.push({ role: 'user', content: prompt });
         }
-        return messages.length ? messages : [];
+        return messages;
       }
     }
     // eslint-disable-next-line no-empty
   } catch {}
-  return undefined;
+  return [];
 }
 
 /**
  * Generate a request.messages JSON array from the prompt field in the
  * invoke_agent op
  */
-export function requestMessagesFromPrompt(span: Span, prompt: unknown): void {
-  if (typeof prompt !== 'string') return;
-  const maybeMessages = convertPromptToMessages(prompt);
-  if (maybeMessages !== undefined) {
-    span.setAttribute('gen_ai.request.messages', getTruncatedJsonString(maybeMessages));
+export function requestMessagesFromPrompt(span: Span, attributes: SpanAttributes): void {
+  if (attributes[AI_PROMPT_ATTRIBUTE]) {
+    const truncatedPrompt = getTruncatedJsonString(attributes[AI_PROMPT_ATTRIBUTE] as string | string[]);
+    span.setAttribute('gen_ai.prompt', truncatedPrompt);
+  }
+  const prompt = attributes[AI_PROMPT_ATTRIBUTE];
+  if (
+    typeof prompt === 'string' &&
+    !attributes[GEN_AI_REQUEST_MESSAGES_ATTRIBUTE] &&
+    !attributes[AI_PROMPT_MESSAGES_ATTRIBUTE]
+  ) {
+    const messages = convertPromptToMessages(prompt);
+    if (messages.length) span.setAttribute(GEN_AI_REQUEST_MESSAGES_ATTRIBUTE, getTruncatedJsonString(messages));
   }
 }