fix(ai, llmobs): respect the functionId argument for vercel-ai operations (#6524)

* fix for ai apm spans to respect functionId as resource name

* add functionId to llmobs span name

Author: sabrenner

packages/datadog-plugin-ai/src/tracing.js

@@ -8,14 +8,14 @@ class VercelAITracingPlugin extends TracingPlugin {
   static prefix = 'tracing:dd-trace:vercel-ai'
 
   bindStart (ctx) {
-    const attributes = ctx.attributes
+    const { attributes, name } = ctx
 
     const model = attributes['ai.model.id']
     const modelProvider = getModelProvider(attributes)
 
-    this.startSpan(ctx.name, {
+    this.startSpan(name, {
       meta: {
-        'resource.name': ctx.name,
+        'resource.name': attributes['resource.name'] ?? name,
         'ai.request.model': model,
         'ai.request.model_provider': modelProvider
       }

packages/datadog-plugin-ai/test/index.spec.js

@@ -451,5 +451,37 @@ describe('Plugin', () => {
 
       await checkTraces
     })
+
+    it('creates a span that respects the functionId', async () => {
+      const checkTraces = agent.assertSomeTraces(traces => {
+        const generateTextSpan = traces[0][0]
+        const doGenerateSpan = traces[0][1]
+
+        assert.strictEqual(generateTextSpan.name, 'ai.generateText')
+        assert.strictEqual(generateTextSpan.resource, 'test')
+        assert.strictEqual(generateTextSpan.meta['ai.request.model'], 'gpt-4o-mini')
+        assert.strictEqual(generateTextSpan.meta['ai.request.model_provider'], 'openai')
+
+        assert.strictEqual(doGenerateSpan.name, 'ai.generateText.doGenerate')
+        assert.strictEqual(doGenerateSpan.resource, 'test')
+        assert.strictEqual(doGenerateSpan.meta['ai.request.model'], 'gpt-4o-mini')
+        assert.strictEqual(doGenerateSpan.meta['ai.request.model_provider'], 'openai')
+      })
+
+      const result = await ai.generateText({
+        model: openai('gpt-4o-mini'),
+        system: 'You are a helpful assistant',
+        prompt: 'Hello, OpenAI!',
+        maxTokens: 100,
+        temperature: 0.5,
+        experimental_telemetry: {
+          functionId: 'test'
+        }
+      })
+
+      assert.ok(result.text, 'Expected result to be truthy')
+
+      await checkTraces
+    })
   })
 })

packages/dd-trace/src/llmobs/plugins/ai/index.js

@@ -17,7 +17,8 @@ const {
   getModelMetadata,
   getGenerationMetadata,
   getToolNameFromTags,
-  getToolCallResultContent
+  getToolCallResultContent,
+  getLlmObsSpanName
 } = require('./util')
 
 const SPAN_NAME_TO_KIND_MAPPING = {
@@ -96,7 +97,7 @@ class VercelAILLMObsPlugin extends BaseLLMObsPlugin {
     const kind = SPAN_NAME_TO_KIND_MAPPING[operation]
     if (!kind) return
 
-    return { kind, name: operation }
+    return { kind, name: getLlmObsSpanName(operation, ctx.attributes['ai.telemetry.functionId']) }
   }
 
   setLLMObsTags (ctx) {

packages/dd-trace/src/llmobs/plugins/ai/util.js

@@ -167,6 +167,16 @@ function getToolCallResultContent (content) {
   }
 }
 
+/**
+ * Computes the LLM Observability `ai` span name
+ * @param {string} operation
+ * @param {string} functionId
+ * @returns {string}
+ */
+function getLlmObsSpanName (operation, functionId) {
+  return functionId ? `${functionId}.${operation}` : operation
+}
+
 module.exports = {
   getSpanTags,
   getOperation,
@@ -175,5 +185,6 @@ module.exports = {
   getModelMetadata,
   getGenerationMetadata,
   getToolNameFromTags,
-  getToolCallResultContent
+  getToolCallResultContent,
+  getLlmObsSpanName
 }

packages/dd-trace/test/llmobs/plugins/ai/index.spec.js

@@ -659,5 +659,59 @@ describe('Plugin', () => {
       expect(toolCallSpan).to.deepEqualWithMockValues(expectedToolCallSpan)
       expect(llmSpan2).to.deepEqualWithMockValues(expectedLlmSpan2)
     })
+
+    it('creates a span that respects the functionId', async () => {
+      await ai.generateText({
+        model: openai('gpt-4o-mini'),
+        system: 'You are a helpful assistant',
+        prompt: 'Hello, OpenAI!',
+        maxTokens: 100,
+        temperature: 0.5,
+        experimental_telemetry: {
+          functionId: 'test'
+        }
+      })
+
+      const { apmSpans, llmobsSpans } = await getEvents()
+
+      const expectedWorkflowSpan = expectedLLMObsNonLLMSpanEvent({
+        span: apmSpans[0],
+        name: 'test.generateText',
+        spanKind: 'workflow',
+        inputValue: 'Hello, OpenAI!',
+        outputValue: MOCK_STRING,
+        metadata: {
+          maxTokens: 100,
+          temperature: 0.5,
+          maxSteps: MOCK_NUMBER,
+          maxRetries: MOCK_NUMBER,
+        },
+        tokenMetrics: { input_tokens: MOCK_NUMBER, output_tokens: MOCK_NUMBER, total_tokens: MOCK_NUMBER },
+        tags: { ml_app: 'test', language: 'javascript', integration: 'ai' },
+      })
+
+      const expectedLlmSpan = expectedLLMObsLLMSpanEvent({
+        span: apmSpans[1],
+        parentId: llmobsSpans[0].span_id,
+        spanKind: 'llm',
+        modelName: 'gpt-4o-mini',
+        modelProvider: 'openai',
+        name: 'test.doGenerate',
+        inputMessages: [
+          { content: 'You are a helpful assistant', role: 'system' },
+          { content: 'Hello, OpenAI!', role: 'user' }
+        ],
+        outputMessages: [{ content: MOCK_STRING, role: 'assistant' }],
+        metadata: {
+          max_tokens: 100,
+          temperature: 0.5,
+        },
+        tokenMetrics: { input_tokens: MOCK_NUMBER, output_tokens: MOCK_NUMBER, total_tokens: MOCK_NUMBER },
+        tags: { ml_app: 'test', language: 'javascript', integration: 'ai' },
+      })
+
+      expect(llmobsSpans[0]).to.deepEqualWithMockValues(expectedWorkflowSpan)
+      expect(llmobsSpans[1]).to.deepEqualWithMockValues(expectedLlmSpan)
+    })
   })
 })