feat(ai): add vercel ai integration

.github/actions/testagent-llmobs/logs/action.yml

@@ -0,0 +1,54 @@
+name: "Get APM Test Agent Logs"
+description: "Get the APM Test Agent logs from the test run, with LLM Observability cassettes attached."
+inputs:
+  container-id:
+    description: "ID of the Docker Container to get logs from (optional)"
+    required: false
+  suffix:
+    description: "suffix of the artifact file name"
+    required: false
+runs:
+  using: composite
+  steps:
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    - name: Create Artifacts Directories
+      run: |
+        mkdir -p "./artifacts/logs"
+        mkdir -p "./artifacts/supported-integrations"
+      shell: bash
+    - name: Save Test Agent Logs
+      id: save_logs
+      if:  runner.debug == '1' # only create test agent log artifacts if the github action has been re-run with debug mode
+      run: |
+        if [ -n "${{inputs.container-id}}" ]; then
+          docker logs "${{inputs.container-id}}" > "artifacts/logs/test_agent_logs_${{ inputs.suffix }}.txt"
+        else
+          docker compose logs testagent-llmobs-vcr > "artifacts/logs/test_agent_logs_${{ inputs.suffix }}.txt"
+        fi
+      shell: bash
+    - name: Archive Test Agent Logs
+      if:  runner.debug == '1' # only create test agent log artifacts if the github action has been re-run with debug mode
+      uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
+      with:
+        name: "test_agent_logs_${{ inputs.suffix }}"
+        path: "./artifacts/logs"
+    - name: Get Tested Integrations from Test Agent
+      run: |
+        # make temporary files to save response data to
+        response=$(mktemp) && headers=$(mktemp)
+
+        # get tested integrations
+        curl -o "$response" -D "$headers" http://127.0.0.1:10126/test/integrations/tested_versions
+
+        # get filename representing the name of the tested integration from headers
+        filename=$(awk -F': ' '/file-name/{print $2}' "$headers" | tr -d '\r\n')
+
+        # copy data to final file and remove temp files
+        mv "$response" "artifacts/supported-integrations/${filename}_supported_versions.csv"
+        rm "$headers"
+      shell: bash
+    - name: Archive Test Agent Tested Versions Artifacts
+      uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2
+      with:
+        name: supported-integrations-${{ inputs.suffix }}
+        path: ./artifacts

.github/actions/testagent-llmobs/start/action.yml

@@ -0,0 +1,8 @@
+name: "Start the APM Test Agent"
+description: "Starts the APM Test Agent image with environment, with LLM Observability cassettes attached."
+runs:
+  using: composite
+  steps:
+    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+    - run: docker compose up -d testagent-llmobs-vcr || docker compose up -d testagent-llmobs-vcr
+      shell: bash

.github/workflows/llmobs.yml

@@ -116,3 +116,26 @@ jobs:
         uses: ./.github/actions/testagent/logs
         with:
           suffix: llmobs-${{ github.job }}
+
+  ai:
+    runs-on: ubuntu-latest
+    env:
+      PLUGINS: ai
+      _DD_LLMOBS_FLUSH_INTERVAL: 0
+    steps:
+      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+      - uses: ./.github/actions/testagent-llmobs/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-llmobs/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,25 +154,24 @@ 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"
+      - "127.0.0.1:9126:9126"
     environment:
-        - LOG_LEVEL=DEBUG
-        - TRACE_LANGUAGE=javascript
-        - ENABLED_CHECKS=trace_stall,meta_tracer_version_header,trace_count_header,trace_peer_service
-        - PORT=9126
+      - LOG_LEVEL=DEBUG
+      - TRACE_LANGUAGE=javascript
+      - ENABLED_CHECKS=trace_stall,meta_tracer_version_header,trace_count_header,trace_peer_service
+      - PORT=9126
 
-  # Use this for local development when making new VCR cassettes to persist to the test agent
-  # Do not use the above testagent service if using this one.
-  # testagent-vcr:
-  #   image: ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.24.1
-  #   ports:
-  #     - "127.0.0.1:9126:9126"
-  #   environment:
-  #     - LOG_LEVEL=DEBUG
-  #     - TRACE_LANGUAGE=javascript
-  #     - ENABLED_CHECKS=trace_stall,meta_tracer_version_header,trace_count_header,trace_peer_service
-  #     - PORT=9126
-  #   volumes:
-  #     - ${VCR_CASSETTES_PATH:-/tmp/empty-vcr}:/vcr-cassettes:delegated
+  testagent-llmobs-vcr:
+    image: ghcr.io/datadog/dd-apm-test-agent/ddapm-test-agent:v1.27.0
+    ports:
+      - "127.0.0.1:10126:10126"
+    environment:
+      - LOG_LEVEL=DEBUG
+      - TRACE_LANGUAGE=javascript
+      - ENABLED_CHECKS=trace_stall,meta_tracer_version_header,trace_count_header,trace_peer_service
+      - PORT=10126
+      - VCR_CASSETTES_DIRECTORY=/vcr-cassettes
+    volumes:
+      - ./packages/dd-trace/test/llmobs/cassettes:/vcr-cassettes

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

@@ -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
+}