feat(ai): add vercel ai integration (#5858)

* add vercel ai integration with otel processing

* add some typedocs and comments

* fix tagger test

* rename to 'ai'

* try doing with a custom tracer

* change up implementation slightly

* codeowners

* undo id changes

* get rid of otel span start/end publishes

* revert llmobs tagger test change

* add better noop default tracer and esm support

* delete util file

* add initial test skeleton

* fix duplicate wrapping

* simplify patching

* apm tests

* write some tests

* add rest of llmobs tests

* add ci job

* fix node version and import issue with check and externally defined version

* fix metadata tagging

* handle tool rolls

* remove import

* add default return and docstring for formatMessage

* fix tool message tests

* add model name and provider tags to apm tracing

* some self review

* address some review comments

* do not stub for tests, instead use dummy test agent

* move cassettes to local directory to fix tests

* configurable flush interval for tests

* use separate image for ai tests that have local cassettes attached

* use different port

* move test flush interval back local

* change in esm test

* Revert "move test flush interval back local"

This reverts commit 39ccf68.

* Revert "use different port"

This reverts commit ac4071c.

* Revert "use separate image for ai tests that have local cassettes attached"

This reverts commit 72eca83.

* Revert "move cassettes to local directory to fix tests"

This reverts commit 20c3c63.

* Revert "change in esm test"

This reverts commit 32d75a6.

* remove env var from llmobs workflow

* fix type hint for test util

* add test cassettes

* re-trigger ci

* more review fixes

* try removing configuration

* revert supported config undoing

* add ai to versions package.json

* add @ai-sdk/openai to versions package.json

* fix tests

* remove unrelated change

* clean up tests & versions more to not use zod directly

* require withVersions directly

* forgotten withversions

* change ci for ai job to use node latest

Author: sabrenner

.github/workflows/llmobs.yml

@@ -119,3 +119,25 @@ jobs:
         uses: ./.github/actions/testagent/logs
         with:
           suffix: llmobs-${{ github.job }}
+
+  ai:
+    runs-on: ubuntu-latest
+    env:
+      PLUGINS: ai
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: ./.github/actions/testagent/start
+      - uses: ./.github/actions/node/oldest-maintenance-lts
+      - uses: ./.github/actions/install
+      - run: yarn test:plugins:ci
+      - run: yarn test:llmobs:plugins:ci
+        shell: bash
+      - uses: ./.github/actions/node/latest
+      - run: yarn test:plugins:ci
+      - run: yarn test:llmobs:plugins:ci
+        shell: bash
+      - uses: codecov/codecov-action@18283e04ce6e62d37312384ff67231eb8fd56d24 # v5.4.3
+      - if: always()
+        uses: ./.github/actions/testagent/logs
+        with:
+          suffix: llmobs-${{ github.job }}

CODEOWNERS

@@ -61,9 +61,11 @@
 /packages/datadog-plugin-openai/ @DataDog/ml-observability
 /packages/datadog-plugin-langchain/ @DataDog/ml-observability
 /packages/datadog-plugin-google-cloud-vertexai/ @DataDog/ml-observability
+/packages/datadog-plugin-ai/ @DataDog/ml-observability
 /packages/datadog-instrumentations/src/openai.js @DataDog/ml-observability
 /packages/datadog-instrumentations/src/langchain.js @DataDog/ml-observability
 /packages/datadog-instrumentations/src/google-cloud-vertexai.js @DataDog/ml-observability
+/packages/datadog-instrumentations/src/ai.js @DataDog/ml-observability
 /packages/datadog-plugin-aws-sdk/src/services/bedrockruntime @DataDog/ml-observability
 /packages/datadog-plugin-aws-sdk/test/bedrockruntime.spec.js @DataDog/ml-observability
 

packages/datadog-instrumentations/src/ai.js

@@ -0,0 +1,140 @@
+'use strict'
+
+const { addHook } = require('./helpers/instrument')
+const shimmer = require('../../datadog-shimmer')
+
+const { channel, tracingChannel } = require('dc-polyfill')
+const toolCreationChannel = channel('dd-trace:vercel-ai:tool')
+
+const TRACED_FUNCTIONS = {
+  generateText: wrapWithTracer,
+  streamText: wrapWithTracer,
+  generateObject: wrapWithTracer,
+  streamObject: wrapWithTracer,
+  embed: wrapWithTracer,
+  embedMany: wrapWithTracer,
+  tool: wrapTool
+}
+
+const vercelAiTracingChannel = tracingChannel('dd-trace:vercel-ai')
+const vercelAiSpanSetAttributesChannel = channel('dd-trace:vercel-ai:span:setAttributes')
+
+const noopTracer = {
+  startActiveSpan () {
+    const fn = arguments[arguments.length - 1]
+
+    const span = {
+      spanContext () { return { traceId: '', spanId: '', traceFlags: 0 } },
+      setAttribute () { return this },
+      setAttributes () { return this },
+      addEvent () { return this },
+      addLink () { return this },
+      addLinks () { return this },
+      setStatus () { return this },
+      updateName () { return this },
+      end () { return this },
+      isRecording () { return false },
+      recordException () { return this }
+    }
+
+    return fn(span)
+  }
+}
+
+function wrapTracer (tracer) {
+  if (Object.hasOwn(tracer, Symbol.for('_dd.wrapped'))) return
+
+  shimmer.wrap(tracer, 'startActiveSpan', function (startActiveSpan) {
+    return function () {
+      const name = arguments[0]
+      const options = arguments.length > 2 ? (arguments[1] ?? {}) : {} // startActiveSpan(name, fn)
+      const cb = arguments[arguments.length - 1]
+
+      const ctx = {
+        name,
+        attributes: options.attributes ?? {}
+      }
+
+      arguments[arguments.length - 1] = shimmer.wrapFunction(cb, function (originalCb) {
+        return function (span) {
+          shimmer.wrap(span, 'end', function (spanEnd) {
+            return function () {
+              vercelAiTracingChannel.asyncEnd.publish(ctx)
+              return spanEnd.apply(this, arguments)
+            }
+          })
+
+          shimmer.wrap(span, 'setAttributes', function (setAttributes) {
+            return function (attributes) {
+              vercelAiSpanSetAttributesChannel.publish({ ctx, attributes })
+              return setAttributes.apply(this, arguments)
+            }
+          })
+
+          shimmer.wrap(span, 'recordException', function (recordException) {
+            return function (exception) {
+              ctx.error = exception
+              vercelAiTracingChannel.error.publish(ctx)
+              return recordException.apply(this, arguments)
+            }
+          })
+
+          return originalCb.apply(this, arguments)
+        }
+      })
+
+      return vercelAiTracingChannel.start.runStores(ctx, () => {
+        const result = startActiveSpan.apply(this, arguments)
+        vercelAiTracingChannel.end.publish(ctx)
+        return result
+      })
+    }
+  })
+
+  Object.defineProperty(tracer, Symbol.for('_dd.wrapped'), { value: true })
+}
+
+function wrapWithTracer (fn) {
+  return function () {
+    const options = arguments[0]
+
+    options.experimental_telemetry ??= { isEnabled: true, tracer: noopTracer }
+    wrapTracer(options.experimental_telemetry.tracer)
+
+    return fn.apply(this, arguments)
+  }
+}
+
+function wrapTool (tool) {
+  return function () {
+    const args = arguments[0]
+    toolCreationChannel.publish(args)
+
+    return tool.apply(this, arguments)
+  }
+}
+
+// CJS exports
+addHook({
+  name: 'ai',
+  versions: ['>=4.0.0'],
+}, exports => {
+  for (const [fnName, patchingFn] of Object.entries(TRACED_FUNCTIONS)) {
+    exports = shimmer.wrap(exports, fnName, patchingFn, { replaceGetter: true })
+  }
+
+  return exports
+})
+
+// ESM exports
+addHook({
+  name: 'ai',
+  versions: ['>=4.0.0'],
+  file: 'dist/index.mjs'
+}, exports => {
+  for (const [fnName, patchingFn] of Object.entries(TRACED_FUNCTIONS)) {
+    exports = shimmer.wrap(exports, fnName, patchingFn, { replaceGetter: true })
+  }
+
+  return exports
+})

packages/datadog-instrumentations/src/helpers/hooks.js

@@ -30,6 +30,7 @@ module.exports = {
   '@smithy/smithy-client': () => require('../aws-sdk'),
   '@vitest/runner': { esmFirst: true, fn: () => require('../vitest') },
   aerospike: () => require('../aerospike'),
+  ai: () => require('../ai'),
   amqp10: () => require('../amqp10'),
   amqplib: () => require('../amqplib'),
   avsc: () => require('../avsc'),

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

@@ -0,0 +1,17 @@
+'use strict'
+
+const CompositePlugin = require('../../dd-trace/src/plugins/composite')
+const VercelAILLMObsPlugin = require('../../dd-trace/src/llmobs/plugins/ai')
+const VercelAITracingPlugin = require('./tracing')
+
+class VercelAIPlugin extends CompositePlugin {
+  static get id () { return 'ai' }
+  static get plugins () {
+    return {
+      llmobs: VercelAILLMObsPlugin,
+      tracing: VercelAITracingPlugin
+    }
+  }
+}
+
+module.exports = VercelAIPlugin

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

@@ -0,0 +1,33 @@
+'use strict'
+
+const TracingPlugin = require('../../dd-trace/src/plugins/tracing')
+const { getModelProvider } = require('./utils')
+
+class VercelAITracingPlugin extends TracingPlugin {
+  static id = 'ai'
+  static prefix = 'tracing:dd-trace:vercel-ai'
+
+  bindStart (ctx) {
+    const attributes = ctx.attributes
+
+    const model = attributes['ai.model.id']
+    const modelProvider = getModelProvider(attributes)
+
+    this.startSpan(ctx.name, {
+      meta: {
+        'resource.name': ctx.name,
+        'ai.request.model': model,
+        'ai.request.model_provider': modelProvider
+      }
+    }, ctx)
+
+    return ctx.currentStore
+  }
+
+  asyncEnd (ctx) {
+    const span = ctx.currentStore?.span
+    span?.finish()
+  }
+}
+
+module.exports = VercelAITracingPlugin

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

@@ -0,0 +1,28 @@
+'use strict'
+
+const { parseModelId } = require('../../datadog-plugin-aws-sdk/src/services/bedrockruntime/utils')
+
+/**
+ * Get the model provider from the span tags or attributes.
+ * This is normalized to LLM Observability model provider standards.
+ *
+ * @param {Record<string, string>} tags
+ * @returns {string}
+ */
+function getModelProvider (tags) {
+  const modelProviderTag = tags['ai.model.provider']
+  const providerParts = modelProviderTag?.split('.')
+  const provider = providerParts?.[0]
+
+  if (provider === 'amazon-bedrock') {
+    const modelId = tags['ai.model.id']
+    const model = modelId && parseModelId(modelId)
+    return model?.modelProvider ?? provider
+  }
+
+  return provider
+}
+
+module.exports = {
+  getModelProvider
+}