dd-trace 6.0.0 → 6.2.0

package/packages/dd-trace/src/llmobs/plugins/ai/index.js

@@ -1,402 +1,15 @@
 'use strict'
 
-const { channel } = require('dc-polyfill')
-const BaseLLMObsPlugin = require('../base')
-const { getModelProvider } = require('../../../../../datadog-plugin-ai/src/utils')
+const CompositePlugin = require('../../../plugins/composite')
+const DdTelemetryPlugin = require('./ddTelemetry')
+const VercelAiTelemetryPlugin = require('./vercelTelemetry')
 
-const toolCreationCh = channel('tracing:orchestrion:ai:tool:start')
-const setAttributesCh = channel('dd-trace:vercel-ai:span:setAttributes')
-
-const { MODEL_NAME, MODEL_PROVIDER, NAME } = require('../../constants/tags')
-const {
-  getSpanTags,
-  getOperation,
-  getUsage,
-  getJsonStringValue,
-  getModelMetadata,
-  getGenerationMetadata,
-  getToolNameFromTags,
-  getToolCallResultContent,
-  getLlmObsSpanName,
-  getTelemetryMetadata,
-} = require('./util')
-
-/**
- * @typedef {Record<string, unknown> & { description?: string, id?: string }} AvailableToolArgs
- */
-
-/**
- * @typedef {string | number | boolean | null | undefined | string[] | number[] | boolean[]} TagValue
- * @typedef {Record<string, TagValue>} SpanTags
- *
- * @typedef {{ name?: string, description?: string }} ToolForModel
- *
- * @typedef {{ type: 'text' | 'reasoning' | 'redacted-reasoning', text?: string, data?: string }} TextPart
- * @typedef {{ type: 'tool-call', toolName: string, toolCallId: string, args?: unknown, input?: unknown }} ToolCallPart
- * @typedef {(
- *   { type: 'tool-result', toolCallId: string, output?: { type: string, value?: unknown }, result?: unknown } &
- *   Record<string, unknown>
- * )} ToolResultPart
- *
- * @typedef {{
- *   role: 'system',
- *   content: string
- * } | {
- *   role: 'user',
- *   content: TextPart[]
- * } | {
- *   role: 'assistant',
- *   content: Array<TextPart | ToolCallPart>
- * } | {
- *   role: 'tool',
- *   content: ToolResultPart[]
- * }} AiSdkMessage
- */
-
-const SPAN_NAME_TO_KIND_MAPPING = {
-  // embeddings
-  embed: 'workflow',
-  embedMany: 'workflow',
-  doEmbed: 'embedding',
-  // object generation
-  generateObject: 'workflow',
-  streamObject: 'workflow',
-  // text generation
-  generateText: 'workflow',
-  streamText: 'workflow',
-  // llm operations
-  doGenerate: 'llm',
-  doStream: 'llm',
-  // tools
-  toolCall: 'tool',
-}
-
-class VercelAILLMObsPlugin extends BaseLLMObsPlugin {
-  static id = 'ai'
+class VercelAILLMObsPlugin extends CompositePlugin {
+  static id = 'ai_llmobs'
   static integration = 'ai'
-  static prefix = 'tracing:dd-trace:vercel-ai'
-
-  /**
-   * The available tools within the runtime scope of this integration.
-   * This essentially acts as a global registry for all tools made through the Vercel AI SDK.
-   * @type {Set<AvailableToolArgs>}
-   */
-  #availableTools
-
-  /**
-   * A mapping of tool call IDs to tool names.
-   * This is used to map the tool call ID to the tool name for the output message.
-   * @type {Record<string, string>}
-   */
-  #toolCallIdsToName
-
-  constructor (...args) {
-    super(...args)
-
-    this.#toolCallIdsToName = {}
-    this.#availableTools = new Set()
-    toolCreationCh.subscribe(ctx => {
-      const toolArgs = ctx.arguments
-      const tool = toolArgs[0] ?? {}
-      this.#availableTools.add(tool)
-    })
-
-    setAttributesCh.subscribe(({ ctx, attributes }) => {
-      Object.assign(ctx.attributes, attributes)
-    })
-  }
-
-  /**
-   * Does a best-effort attempt to find the right tool name for the given tool description.
-   * This is because the Vercel AI SDK does not tag tools by name properly, but
-   * rather by the index they were passed in. Tool names appear nowhere in the span tags.
-   *
-   * We use the tool description as the next best identifier for a tool.
-   *
-   * @param {string} toolName
-   * @param {string | undefined} toolDescription
-   * @returns {string | undefined}
-   */
-  findToolName (toolName, toolDescription) {
-    if (Number.isNaN(Number.parseInt(toolName))) return toolName
-
-    for (const availableTool of this.#availableTools) {
-      const description = availableTool.description
-      if (description === toolDescription && availableTool.id) {
-        return availableTool.id
-      }
-    }
-  }
-
-  /**
-   * @override
-   */
-  getLLMObsSpanRegisterOptions (ctx) {
-    const span = ctx.currentStore?.span
-    const operation = getOperation(span)
-    const kind = SPAN_NAME_TO_KIND_MAPPING[operation]
-    if (!kind) return
-
-    return { kind, name: getLlmObsSpanName(operation, ctx.attributes['ai.telemetry.functionId']) }
-  }
-
-  /**
-   * @override
-   */
-  setLLMObsTags (ctx) {
-    const span = ctx.currentStore?.span
-    if (!span) return
-
-    const operation = getOperation(span)
-    const kind = SPAN_NAME_TO_KIND_MAPPING[operation]
-    if (!kind) return
-
-    const tags = getSpanTags(ctx)
-
-    if (['embedding', 'llm'].includes(kind)) {
-      this._tagger._setTag(span, MODEL_NAME, tags['ai.model.id'])
-      this._tagger._setTag(span, MODEL_PROVIDER, getModelProvider(tags))
-    }
-
-    switch (operation) {
-      case 'embed':
-      case 'embedMany':
-        this.setEmbeddingWorkflowTags(span, tags)
-        break
-      case 'doEmbed':
-        this.setEmbeddingTags(span, tags)
-        break
-      case 'generateObject':
-      case 'streamObject':
-        this.setObjectGenerationTags(span, tags)
-        break
-      case 'generateText':
-      case 'streamText':
-        this.setTextGenerationTags(span, tags)
-        break
-      case 'doGenerate':
-      case 'doStream':
-        this.setLLMOperationTags(span, tags)
-        break
-      case 'toolCall':
-        this.setToolTags(span, tags)
-        break
-      default:
-        break
-    }
-  }
-
-  setEmbeddingWorkflowTags (span, tags) {
-    const inputs = tags['ai.value'] ?? tags['ai.values']
-    const parsedInputs = Array.isArray(inputs)
-      ? inputs.map(input => getJsonStringValue(input, ''))
-      : getJsonStringValue(inputs, '')
-
-    const embeddingsOutput = tags['ai.embedding'] ?? tags['ai.embeddings']
-    const isSingleEmbedding = !Array.isArray(embeddingsOutput)
-    const numberOfEmbeddings = isSingleEmbedding ? 1 : embeddingsOutput.length
-    const embeddingsLength = getJsonStringValue(isSingleEmbedding ? embeddingsOutput : embeddingsOutput?.[0], []).length
-    const output = `[${numberOfEmbeddings} embedding(s) returned with size ${embeddingsLength}]`
-
-    this._tagger.tagTextIO(span, parsedInputs, output)
-
-    const metadata = getGenerationMetadata(tags)
-    this._tagger.tagMetadata(span, metadata)
-  }
-
-  setEmbeddingTags (span, tags) {
-    const inputs = tags['ai.values']
-    if (!Array.isArray(inputs)) return
-
-    const parsedInputs = inputs.map(input => getJsonStringValue(input, ''))
-
-    const embeddingsOutput = tags['ai.embeddings']
-    const numberOfEmbeddings = embeddingsOutput?.length
-    const embeddingsLength = getJsonStringValue(embeddingsOutput?.[0], []).length
-    const output = `[${numberOfEmbeddings} embedding(s) returned with size ${embeddingsLength}]`
-
-    this._tagger.tagEmbeddingIO(span, parsedInputs, output)
-
-    const metadata = getTelemetryMetadata(tags)
-    this._tagger.tagMetadata(span, metadata)
-
-    const usage = tags['ai.usage.tokens']
-    this._tagger.tagMetrics(span, {
-      inputTokens: usage,
-      totalTokens: usage,
-    })
-  }
-
-  setObjectGenerationTags (span, tags) {
-    const promptInfo = getJsonStringValue(tags['ai.prompt'], {})
-    const lastUserPrompt =
-      promptInfo.prompt ??
-      promptInfo.messages.reverse().find(message => message.role === 'user')?.content
-    const prompt = Array.isArray(lastUserPrompt) ? lastUserPrompt.map(part => part.text ?? '').join('') : lastUserPrompt
-
-    const output = tags['ai.response.object']
-
-    this._tagger.tagTextIO(span, prompt, output)
-
-    const metadata = getGenerationMetadata(tags)
-    metadata.schema = getJsonStringValue(tags['ai.schema'], {})
-    this._tagger.tagMetadata(span, metadata)
-  }
-
-  setTextGenerationTags (span, tags) {
-    const promptInfo = getJsonStringValue(tags['ai.prompt'], {})
-    const lastUserPrompt =
-      promptInfo.prompt ??
-      promptInfo.messages.reverse().find(message => message.role === 'user')?.content
-    const prompt = Array.isArray(lastUserPrompt) ? lastUserPrompt.map(part => part.text ?? '').join('') : lastUserPrompt
-
-    const output = tags['ai.response.text']
-
-    this._tagger.tagTextIO(span, prompt, output)
-
-    const metadata = getGenerationMetadata(tags)
-    this._tagger.tagMetadata(span, metadata)
-  }
-
-  /**
-   * @param {import('../../../opentracing/span')} span
-   * @param {SpanTags} tags
-   */
-  setLLMOperationTags (span, tags) {
-    const toolsForModel = tags['ai.prompt.tools']?.map(getJsonStringValue)
-
-    const inputMessages = getJsonStringValue(tags['ai.prompt.messages'], [])
-    const parsedInputMessages = []
-    for (const message of inputMessages) {
-      const formattedMessages = this.formatMessage(message, toolsForModel)
-      parsedInputMessages.push(...formattedMessages)
-    }
-
-    const outputMessage = this.formatOutputMessage(tags, toolsForModel)
-
-    this._tagger.tagLLMIO(span, parsedInputMessages, outputMessage)
-
-    const metadata = getModelMetadata(tags)
-    this._tagger.tagMetadata(span, metadata)
-
-    const usage = getUsage(tags)
-    this._tagger.tagMetrics(span, usage)
-  }
-
-  setToolTags (span, tags) {
-    const toolCallId = tags['ai.toolCall.id']
-    const name = getToolNameFromTags(tags) ?? this.#toolCallIdsToName[toolCallId]
-    if (name) this._tagger._setTag(span, NAME, name)
-
-    const input = tags['ai.toolCall.args']
-    const output = tags['ai.toolCall.result']
-
-    this._tagger.tagTextIO(span, input, output)
-  }
-
-  formatOutputMessage (tags, toolsForModel) {
-    const outputMessageText = tags['ai.response.text'] ?? tags['ai.response.object']
-    const outputMessageToolCalls = getJsonStringValue(tags['ai.response.toolCalls'], [])
-
-    const formattedToolCalls = []
-    for (const toolCall of outputMessageToolCalls) {
-      const toolArgs = toolCall.args ?? toolCall.input
-      const toolCallArgs = typeof toolArgs === 'string' ? getJsonStringValue(toolArgs, {}) : toolArgs
-      const toolDescription = toolsForModel?.find(tool => toolCall.toolName === tool.name)?.description
-      const name = this.findToolName(toolCall.toolName, toolDescription)
-      this.#toolCallIdsToName[toolCall.toolCallId] = name
-
-      formattedToolCalls.push({
-        arguments: toolCallArgs,
-        name,
-        toolId: toolCall.toolCallId,
-        type: toolCall.toolCallType ?? 'function',
-      })
-    }
-
-    return {
-      role: 'assistant',
-      content: outputMessageText,
-      toolCalls: formattedToolCalls,
-    }
-  }
-
-  /**
-   * Returns a list of formatted messages from a message object.
-   * Most of these will just be one entry, but in the case of a "tool" role,
-   * it is possible to have multiple tool call results in a single message that we
-   * need to split into multiple messages.
-   *
-   * @param {AiSdkMessage} message
-   * @param {ToolForModel[] | null | undefined} toolsForModel
-   * @returns {Array<{role: string, content: string, toolId?: string,
-   *   toolCalls?: Array<{arguments: string, name: string, toolId: string, type: string}>}>}
-   */
-  formatMessage (message, toolsForModel) {
-    const { role, content } = message
-
-    if (role === 'system') {
-      return [{ role, content }]
-    } else if (role === 'user') {
-      let finalContent = ''
-      for (const part of content) {
-        const { type } = part
-        if (type === 'text') {
-          finalContent += part.text
-        }
-      }
-
-      return [{ role, content: finalContent }]
-    } else if (role === 'assistant') {
-      const toolCalls = []
-      let finalContent = ''
-
-      for (const part of content) {
-        const { type } = part
-        // TODO(sabrenner): do we want to include reasoning?
-        if (['text', 'reasoning', 'redacted-reasoning'].includes(type)) {
-          finalContent += part.text ?? part.data
-        } else if (type === 'tool-call') {
-          const toolDescription = toolsForModel?.find(tool => part.toolName === tool.name)?.description
-          const name = this.findToolName(part.toolName, toolDescription)
-
-          toolCalls.push({
-            arguments: part.args ?? part.input,
-            name,
-            toolId: part.toolCallId,
-            type: 'function',
-          })
-        }
-      }
-
-      const finalMessage = {
-        role,
-        content: finalContent,
-      }
-
-      if (toolCalls.length) {
-        finalMessage.toolCalls = toolCalls.length ? toolCalls : undefined
-      }
-
-      return [finalMessage]
-    } else if (role === 'tool') {
-      const finalMessages = []
-      for (const part of content) {
-        if (part.type === 'tool-result') {
-          const safeResult = getToolCallResultContent(part)
-
-          finalMessages.push({
-            role,
-            content: safeResult,
-            toolId: part.toolCallId,
-          })
-        }
-      }
-
-      return finalMessages
-    }
-
-    return []
+  static plugins = {
+    dd: DdTelemetryPlugin,
+    ai: VercelAiTelemetryPlugin,
   }
 }
 

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

@@ -259,6 +259,32 @@ function getGenerationMetadata (tags) {
   return Object.keys(metadata).length ? metadata : null
 }
 
+/**
+ * Get the generation metadata from the span tags (maxSteps, maxRetries, etc.)
+ * Additionally, set telemetry metadata from manual telemetry tags.
+ * @param {Record<string, unknown>} event
+ * @returns {Record<string, unknown> | null}
+ */
+function getGenerationMetadataFromEvent (event) {
+  /** @type {Record<string, unknown>} */
+  const metadata = {}
+
+  for (const [key, value] of Object.entries(event)) {
+    const transformedKey = key.replaceAll(/[A-Z]/g, letter => '_' + letter.toLowerCase())
+    if (!MODEL_METADATA_KEYS.has(transformedKey)) {
+      if (key === 'runtimeContext') { // custom telemetry metadata
+        Object.assign(metadata, value)
+      }
+
+      continue
+    }
+
+    metadata[transformedKey] = value
+  }
+
+  return Object.keys(metadata).length ? metadata : null
+}
+
 /**
  * Get the tool name from the span tags.
  * If the tool name is a parsable number, or is not found, null is returned.
@@ -348,4 +374,5 @@ module.exports = {
   getToolCallResultContent,
   getLlmObsSpanName,
   getTelemetryMetadata,
+  getGenerationMetadataFromEvent,
 }