fix(ai): properly handle token usage for ai@>=5.0.0 (#6707)

* fix(ai): add support for AI SDK v5 token properties

AI SDK v5 renamed token usage properties from promptTokens/completionTokens
to inputTokens/outputTokens/totalTokens. This was causing token metrics to
not show up for streaming operations.

The getUsage() function now checks for v5 properties first, then falls back
to v4 for backwards compatibility:
- ai.usage.inputTokens (v5) → ai.usage.promptTokens (v4)
- ai.usage.outputTokens (v5) → ai.usage.completionTokens (v4)
- ai.usage.totalTokens (v5) → computed sum (v4)

See: https://ai-sdk.dev/docs/migration-guides/migration-guide-5-0#usage-token-properties

* fix tests

---------

Co-authored-by: Mark Anthony Cianfrani <mark@cianfrani.me>

Author: sabrenner

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

@@ -45,19 +45,24 @@ function getOperation (span) {
 
 /**
  * Get the LLM token usage from the span tags
+ * Supports both AI SDK v4 (promptTokens/completionTokens) and v5 (inputTokens/outputTokens)
  * @template T extends {inputTokens: number, outputTokens: number, totalTokens: number}
  * @param {T} tags
  * @returns {Pick<T, 'inputTokens' | 'outputTokens' | 'totalTokens'>}
  */
 function getUsage (tags) {
   const usage = {}
-  const inputTokens = tags['ai.usage.promptTokens']
-  const outputTokens = tags['ai.usage.completionTokens']
+
+  // AI SDK v5 uses inputTokens/outputTokens, v4 uses promptTokens/completionTokens
+  // Check v5 properties first, fall back to v4
+  const inputTokens = tags['ai.usage.inputTokens'] ?? tags['ai.usage.promptTokens']
+  const outputTokens = tags['ai.usage.outputTokens'] ?? tags['ai.usage.completionTokens']
 
   if (inputTokens != null) usage.inputTokens = inputTokens
   if (outputTokens != null) usage.outputTokens = outputTokens
 
-  const totalTokens = inputTokens + outputTokens
+  // v5 provides totalTokens directly, v4 requires computation
+  const totalTokens = tags['ai.usage.totalTokens'] ?? (inputTokens + outputTokens)
   if (!Number.isNaN(totalTokens)) usage.totalTokens = totalTokens
 
   return usage

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

@@ -17,6 +17,7 @@ const {
   MOCK_NUMBER,
   MOCK_OBJECT
 } = require('../../util')
+const assert = require('node:assert')
 
 chai.Assertion.addMethod('deepEqualWithMockValues', deepEqualWithMockValues)
 
@@ -279,6 +280,15 @@ describe('Plugin', () => {
 
       expect(llmobsSpans[0]).to.deepEqualWithMockValues(expectedWorkflowSpan)
       expect(llmobsSpans[1]).to.deepEqualWithMockValues(expectedLlmSpan)
+
+      // manually asserting the token metrics are set correctly
+      // TODO(MLOB-4234): the llmobs span event assertions are slightly buggy and need to be re-worked
+      assert.ok(typeof llmobsSpans[1].metrics.input_tokens === 'number')
+      assert.ok(llmobsSpans[1].metrics.input_tokens > 0)
+      assert.ok(typeof llmobsSpans[1].metrics.output_tokens === 'number')
+      assert.ok(llmobsSpans[1].metrics.output_tokens > 0)
+      assert.ok(typeof llmobsSpans[1].metrics.total_tokens === 'number')
+      assert.ok(llmobsSpans[1].metrics.total_tokens > 0)
     })
 
     it('creates a span for streamObject', async () => {
@@ -338,6 +348,15 @@ describe('Plugin', () => {
 
       expect(llmobsSpans[0]).to.deepEqualWithMockValues(expectedWorkflowSpan)
       expect(llmobsSpans[1]).to.deepEqualWithMockValues(expectedLlmSpan)
+
+      // manually asserting the token metrics are set correctly
+      // TODO(MLOB-4234): the llmobs span event assertions are slightly buggy and need to be re-worked
+      assert.ok(typeof llmobsSpans[1].metrics.input_tokens === 'number')
+      assert.ok(llmobsSpans[1].metrics.input_tokens > 0)
+      assert.ok(typeof llmobsSpans[1].metrics.output_tokens === 'number')
+      assert.ok(llmobsSpans[1].metrics.output_tokens > 0)
+      assert.ok(typeof llmobsSpans[1].metrics.total_tokens === 'number')
+      assert.ok(llmobsSpans[1].metrics.total_tokens > 0)
     })
 
     it('creates a span for a tool call', async () => {
@@ -692,6 +711,22 @@ describe('Plugin', () => {
       expect(llmSpan).to.deepEqualWithMockValues(expectedLlmSpan)
       expect(toolCallSpan).to.deepEqualWithMockValues(expectedToolCallSpan)
       expect(llmSpan2).to.deepEqualWithMockValues(expectedLlmSpan2)
+
+      // manually asserting the token metrics are set correctly
+      // TODO(MLOB-4234): the llmobs span event assertions are slightly buggy and need to be re-worked
+      assert.ok(typeof llmSpan.metrics.input_tokens === 'number')
+      assert.ok(llmSpan.metrics.input_tokens > 0)
+      assert.ok(typeof llmSpan.metrics.output_tokens === 'number')
+      assert.ok(llmSpan.metrics.output_tokens > 0)
+      assert.ok(typeof llmSpan.metrics.total_tokens === 'number')
+      assert.ok(llmSpan.metrics.total_tokens > 0)
+
+      assert.ok(typeof llmSpan2.metrics.input_tokens === 'number')
+      assert.ok(llmSpan2.metrics.input_tokens > 0)
+      assert.ok(typeof llmSpan2.metrics.output_tokens === 'number')
+      assert.ok(llmSpan2.metrics.output_tokens > 0)
+      assert.ok(typeof llmSpan2.metrics.total_tokens === 'number')
+      assert.ok(llmSpan2.metrics.total_tokens > 0)
     })
 
     it('creates a span that respects the functionId', async () => {