feat(ai): add vercel ai integration

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

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

packages/datadog-instrumentations/src/vercel-ai.js

@@ -0,0 +1,70 @@
+'use strict'
+
+const { addHook } = require('./helpers/instrument')
+const tracer = require('../../dd-trace')
+const shimmer = require('../../datadog-shimmer')
+
+const { TracerProvider } = tracer
+const provider = new TracerProvider()
+provider.register()
+
+const { channel } = 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
+}
+
+function wrapWithTracer (fn) {
+  return function () {
+    const options = arguments[0]
+    if (options.experimental_telemetry != null) return fn.apply(this, arguments)
+
+    options.experimental_telemetry = {
+      isEnabled: true,
+      // TODO(sabrenner): need to figure out how a user can configure this tracer
+      // maybe we advise they do this manually?
+      tracer: provider.getTracer()
+    }
+
+    return fn.apply(this, arguments)
+  }
+}
+
+function wrapTool (tool) {
+  return function () {
+    const args = arguments[0]
+    toolCreationChannel.publish(args)
+
+    return tool.apply(this, arguments)
+  }
+}
+
+addHook({
+  name: 'ai',
+  versions: ['>=4.0.0'],
+}, exports => {
+  // the exports from this package are not configurable
+  // to return wrapped functions with the tracer provided above,
+  // we need to copy the exports
+  const wrappedExports = {}
+
+  for (const [fnName, patchingFn] of Object.entries(TRACED_FUNCTIONS)) {
+    const original = exports[fnName]
+    wrappedExports[fnName] = shimmer.wrapFunction(original, patchingFn)
+  }
+
+  Object.getOwnPropertyNames(exports).forEach(prop => {
+    if (!Object.keys(TRACED_FUNCTIONS).includes(prop)) {
+      wrappedExports[prop] = exports[prop]
+    }
+  })
+
+  return wrappedExports
+})

packages/datadog-plugin-vercel-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/vercel-ai')
+const VercelAITracingPlugin = require('./tracing')
+
+class VercelAIPlugin extends CompositePlugin {
+  static get id () { return 'vercel-ai' }
+  static get plugins () {
+    return {
+      llmobs: VercelAILLMObsPlugin,
+      tracing: VercelAITracingPlugin
+    }
+  }
+}
+
+module.exports = VercelAIPlugin

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

@@ -0,0 +1,55 @@
+'use strict'
+
+const Plugin = require('../../dd-trace/src/plugins/plugin')
+
+const { channel } = require('dc-polyfill')
+const spanFinishCh = channel('dd-trace:otel:span:finish')
+
+const { isVercelAISpan } = require('./util')
+
+// filter input/output based tags for APM spans
+// this is for data access controls and for sensitive data
+// the LLM Observability feature is recommended in its place for better
+// data access controls and sensitive data scrubbing
+const TAG_PATTERNS_TO_FILTER = [
+  'prompt', // TODO(sabrenner): we need to refine this so it doesn't filter out prompt tokens
+  'messages',
+  'response.text',
+  'toolCalls',
+  'toolCall.args',
+  'toolCall.result',
+  'values',
+  'embedding'
+]
+
+/**
+ * Determines if an OpenTelemetry span tag should be filtered out on the final DD span
+ *
+ * @param {String} key
+ * @returns {Boolean}
+ */
+function shouldFilterTag (key) {
+  return TAG_PATTERNS_TO_FILTER.some(pattern => key.includes(pattern))
+}
+
+class VercelAITracingPlugin extends Plugin {
+  static get id () { return 'ai' }
+
+  constructor (...args) {
+    super(...args)
+
+    spanFinishCh.subscribe(({ ddSpan }) => {
+      if (!isVercelAISpan(ddSpan)) {
+        return
+      }
+
+      for (const key of Object.keys(ddSpan.context()._tags)) {
+        if (shouldFilterTag(key)) {
+          ddSpan.setTag(key, undefined)
+        }
+      }
+    })
+  }
+}
+
+module.exports = VercelAITracingPlugin

packages/datadog-plugin-vercel-ai/src/util.js

@@ -0,0 +1,15 @@
+'use strict'
+
+/**
+ * Determines if an OpenTelemetry span is a Vercel AI span
+ *
+ * @param {import('../../dd-trace/src/opentracing/span') | null} span
+ * @returns {Boolean}
+ */
+function isVercelAISpan (span) {
+  return span._name?.startsWith('ai') && span.context()._tags?.['resource.name']?.startsWith('ai')
+}
+
+module.exports = {
+  isVercelAISpan
+}

packages/dd-trace/src/id.js

@@ -182,6 +182,6 @@ function writeUInt32BE (buffer, value, offset) {
   buffer[0 + offset] = value & 255
 }
 
-module.exports = function createIdentifier (value, radix) {
+module.exports = Object.assign(function createIdentifier (value, radix) {
   return new Identifier(value, radix)
-}
+}, { Identifier })

packages/dd-trace/src/llmobs/plugins/base.js

@@ -22,12 +22,17 @@ class LLMObsPlugin extends TracingPlugin {
     throw new Error('getLLMObsSPanRegisterOptions must be implemented by the subclass')
   }
 
-  start (ctx) {
+  start (ctx, options = {}) {
     // even though llmobs span events won't be enqueued if llmobs is disabled
     // we should avoid doing any computations here (these listeners aren't disabled)
     const enabled = this._tracerConfig.llmobs.enabled
     if (!enabled) return
 
+    const {
+      useLlmObsParent = true,
+      enterIntoLlmObsStorage = true
+    } = options
+
     const parent = this.getLLMObsParent(ctx)
     const apmStore = ctx.currentStore
     const span = apmStore?.span
@@ -40,10 +45,20 @@ class LLMObsPlugin extends TracingPlugin {
       telemetry.incrementLLMObsSpanStartCount({ autoinstrumented: true, integration: this.constructor.integration })
 
       ctx.llmobs = {} // initialize context-based namespace
-      llmobsStorage.enterWith({ span })
       ctx.llmobs.parent = parent
 
-      this._tagger.registerLLMObsSpan(span, { parent, integration: this.constructor.integration, ...registerOptions })
+      if (enterIntoLlmObsStorage) {
+        llmobsStorage.enterWith({ span })
+      }
+
+      const coalescedRegisterOptions = {
+        integration: this.constructor.integration,
+        ...registerOptions
+      }
+
+      coalescedRegisterOptions.parent = useLlmObsParent ? parent : span.context()._parentId
+
+      this._tagger.registerLLMObsSpan(span, coalescedRegisterOptions)
     }
   }