feat(ai): add vercel ai integration

.github/workflows/llmobs.yml

@@ -116,3 +116,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/active-lts
+      - 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

@@ -58,9 +58,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
 

docker-compose.yml

@@ -154,7 +154,7 @@ services:
       - LDAP_PASSWORDS=password1,password2
 
   testagent:
-    image: ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.24.1
+    image: ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.27.0
     ports:
         - "127.0.0.1:9126:9126"
     environment:

packages/datadog-instrumentations/src/ai.js

@@ -0,0 +1,162 @@
+'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

@@ -29,6 +29,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 get id () { return 'ai' }
+  static get prefix () { return '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
+}