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

Fix JS-1216

Author: isaacs

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

@@ -10,7 +10,12 @@ import {
 import { getTruncatedJsonString } from '../ai/utils';
 import { toolCallSpanMap } from './constants';
 import type { TokenSummary } from './types';
-import { accumulateTokensForParent, applyAccumulatedTokens, convertAvailableToolsToJsonString } from './utils';
+import {
+  accumulateTokensForParent,
+  applyAccumulatedTokens,
+  convertAvailableToolsToJsonString,
+  requestMessagesFromPrompt,
+} from './utils';
 import type { ProviderMetadata } from './vercel-ai-attributes';
 import {
   AI_MODEL_ID_ATTRIBUTE,
@@ -141,6 +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');
 
   addProviderMetadataToAttributes(attributes);
 
@@ -206,6 +212,10 @@ function processGenerateSpan(span: Span, name: string, attributes: SpanAttribute
   if (attributes[AI_PROMPT_ATTRIBUTE]) {
     const truncatedPrompt = getTruncatedJsonString(attributes[AI_PROMPT_ATTRIBUTE] as string | string[]);
     span.setAttribute('gen_ai.prompt', truncatedPrompt);
+
+    if (!attributes['gen_ai.request.messages']) {
+      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,6 +1,7 @@
 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 { getTruncatedJsonString } from '../ai/utils';
 import { toolCallSpanMap } from './constants';
 import type { TokenSummary } from './types';
 
@@ -87,3 +88,39 @@ export function convertAvailableToolsToJsonString(tools: unknown[]): string {
   });
   return JSON.stringify(toolObjects);
 }
+
+/**
+ * Convert the prompt string to messages array
+ */
+export function convertPromptToMessages(prompt: string): { role: string; content: string }[] | undefined {
+  try {
+    const p = JSON.parse(prompt);
+    if (!!p && typeof p === 'object') {
+      const { prompt, system } = p;
+      if (typeof prompt === 'string' || typeof system === 'string') {
+        const messages: { role: string; content: string }[] = [];
+        if (typeof system === 'string') {
+          messages.push({ role: 'system', content: system });
+        }
+        if (typeof prompt === 'string') {
+          messages.push({ role: 'user', content: prompt });
+        }
+        return messages.length ? messages : [];
+      }
+    }
+    // eslint-disable-next-line no-empty
+  } catch {}
+  return undefined;
+}
+
+/**
+ * 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));
+  }
+}

packages/core/test/lib/utils/vercelai-utils.test.ts

@@ -0,0 +1,55 @@
+import { describe, expect, it } from 'vitest';
+import { convertPromptToMessages } from '../../../src/tracing/vercel-ai/utils';
+
+describe('vercel-ai-utils', () => {
+  describe('convertPromptToMessages', () => {
+    it('should convert a prompt with system to a messages array', () => {
+      expect(
+        convertPromptToMessages(
+          JSON.stringify({
+            system: 'You are a friendly robot',
+            prompt: 'Hello, robot',
+          }),
+        ),
+      ).toStrictEqual([
+        { role: 'system', content: 'You are a friendly robot' },
+        { role: 'user', content: 'Hello, robot' },
+      ]);
+    });
+
+    it('should convert a system prompt to a messages array', () => {
+      expect(
+        convertPromptToMessages(
+          JSON.stringify({
+            system: 'You are a friendly robot',
+          }),
+        ),
+      ).toStrictEqual([{ role: 'system', content: 'You are a friendly robot' }]);
+    });
+
+    it('should convert a user only prompt to a messages array', () => {
+      expect(
+        convertPromptToMessages(
+          JSON.stringify({
+            prompt: 'Hello, robot',
+          }),
+        ),
+      ).toStrictEqual([{ role: 'user', content: 'Hello, robot' }]);
+    });
+
+    it('should ignore unexpected data', () => {
+      expect(
+        convertPromptToMessages(
+          JSON.stringify({
+            randomField: 'Hello, robot',
+            nothing: 'that we know how to handle',
+          }),
+        ),
+      ).toBe(undefined);
+    });
+
+    it('should not break on invalid json', () => {
+      expect(convertPromptToMessages('this is not json')).toBe(undefined);
+    });
+  });
+});