dd-trace 5.80.0 → 5.82.0

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

@@ -10,16 +10,35 @@ const MODEL_METADATA_KEYS = new Set([
   'stop_sequences'
 ])
 
+/**
+ * @typedef {import('../../../opentracing/span')} Span
+ *
+ * @typedef {string | number | boolean | null | undefined | string[] | number[] | boolean[]} TagValue
+ * @typedef {Record<string, TagValue>} SpanTags
+ *
+ * @typedef {{ span?: Span }} CurrentStore
+ * @typedef {{ currentStore?: CurrentStore, attributes?: SpanTags }} AiPluginContext
+ */
+
+/**
+ * @typedef {{
+ *   type: string,
+ *   value?: unknown
+ * }} ToolCallOutput
+ *
+ * @typedef {{ output?: ToolCallOutput, result?: unknown } & Record<string, unknown>} ToolCallResultContent
+ */
+
 /**
  * Get the span tags from the context (either the attributes or the span tags).
  *
- * @param {Record<string, any>} ctx
- * @returns {Record<string, any>}
+ * @param {AiPluginContext} ctx
+ * @returns {SpanTags}
  */
 function getSpanTags (ctx) {
   const span = ctx.currentStore?.span
   const carrier = ctx.attributes ?? span?.context()._tags ?? {}
-  return carrier
+  return /** @type {SpanTags} */ (carrier)
 }
 
 /**
@@ -88,10 +107,11 @@ function getJsonStringValue (str, defaultValue) {
 
 /**
  * Get the model metadata from the span tags (top_p, top_k, temperature, etc.)
- * @param {Record<string, unknown>} tags
- * @returns {Record<string, string> | null}
+ * @param {SpanTags} tags
+ * @returns {Record<string, unknown> | null}
  */
 function getModelMetadata (tags) {
+  /** @type {Record<string, unknown>} */
   const modelMetadata = {}
   for (const metadata of MODEL_METADATA_KEYS) {
     const metadataTagKey = `gen_ai.request.${metadata}`
@@ -106,10 +126,11 @@ function getModelMetadata (tags) {
 
 /**
  * Get the generation metadata from the span tags (maxSteps, maxRetries, etc.)
- * @param {Record<string, string>} tags
- * @returns {Record<string, string> | null}
+ * @param {SpanTags} tags
+ * @returns {Record<string, unknown> | null}
  */
 function getGenerationMetadata (tags) {
+  /** @type {Record<string, unknown>} */
   const metadata = {}
 
   for (const tag of Object.keys(tags)) {
@@ -131,7 +152,7 @@ function getGenerationMetadata (tags) {
  * If the tool name is a parsable number, or is not found, null is returned.
  * Older versions of the ai sdk would tag the tool name as its index in the tools array.
  *
- * @param {Record<string, string>} tags
+ * @param {SpanTags} tags
  * @returns {string | null}
  */
 function getToolNameFromTags (tags) {
@@ -147,7 +168,7 @@ function getToolNameFromTags (tags) {
 /**
  * Get the content of a tool call result.
  * Version 5 of the ai sdk sets this tag as `content.output`, with a `
- * @param {Record<string, any>} content
+ * @param {{ output?: { type: string, value: unknown }, result?: unknown }} content
  * @returns {string}
  */
 function getToolCallResultContent (content) {

package/packages/dd-trace/src/llmobs/plugins/bedrockruntime.js

@@ -1,16 +1,16 @@
 'use strict'
 
-const BaseLLMObsPlugin = require('./base')
 const { storage } = require('../../../../datadog-core')
-const llmobsStore = storage('llmobs')
 const telemetry = require('../telemetry')
-
 const {
   extractRequestParams,
   extractTextAndResponseReason,
   parseModelId,
   extractTextAndResponseReasonFromStream
 } = require('../../../../datadog-plugin-aws-sdk/src/services/bedrockruntime/utils')
+const BaseLLMObsPlugin = require('./base')
+
+const llmobsStore = storage('llmobs')
 
 const ENABLED_OPERATIONS = new Set(['invokeModel', 'invokeModelWithResponseStream'])
 

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

@@ -0,0 +1,104 @@
+'use strict'
+
+const LLMObsPlugin = require('../base')
+const {
+  getOperation,
+  extractMetrics,
+  extractMetadata,
+  aggregateStreamingChunks,
+  formatInputMessages,
+  formatEmbeddingInput,
+  formatOutputMessages,
+  formatEmbeddingOutput
+} = require('./util')
+
+class GenAiLLMObsPlugin extends LLMObsPlugin {
+  static id = 'google-genai'
+  static integration = 'google_genai'
+  static prefix = 'tracing:apm:google:genai:request'
+
+  constructor () {
+    super(...arguments)
+
+    // Subscribe to streaming chunk events
+    this.addSub('apm:google:genai:request:chunk', ({ ctx, chunk, done }) => {
+      ctx.isStreaming = true
+      ctx.chunks = ctx.chunks || []
+
+      if (chunk) ctx.chunks.push(chunk)
+      if (!done) return
+
+      // Aggregate streaming chunks into a single response
+      ctx.result = aggregateStreamingChunks(ctx.chunks)
+    })
+  }
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    const { args, methodName } = ctx
+    if (!methodName) return
+
+    const inputs = args[0]
+    const operation = getOperation(methodName)
+
+    return {
+      modelProvider: 'google',
+      modelName: inputs.model,
+      kind: operation,
+      name: 'google_genai.request'
+    }
+  }
+
+  setLLMObsTags (ctx) {
+    const { args, methodName } = ctx
+    const span = ctx.currentStore?.span
+    if (!methodName) return
+
+    const inputs = args[0]
+    const response = ctx.result
+    const error = !!span.context()._tags.error
+
+    const operation = getOperation(methodName)
+
+    if (operation === 'llm') {
+      this.#tagGenerateContent(span, inputs, response, error, ctx.isStreaming)
+    } else if (operation === 'embedding') {
+      this.#tagEmbedding(span, inputs, response, error)
+    }
+
+    if (!error && response) {
+      const metrics = extractMetrics(response)
+      this._tagger.tagMetrics(span, metrics)
+    }
+  }
+
+  #tagGenerateContent (span, inputs, response, error, isStreaming = false) {
+    const { config = {} } = inputs
+
+    const inputMessages = formatInputMessages(inputs.contents)
+
+    const metadata = extractMetadata(config)
+    this._tagger.tagMetadata(span, metadata)
+
+    if (error) {
+      this._tagger.tagLLMIO(span, inputMessages, [{ content: '' }])
+      return
+    }
+
+    const outputMessages = formatOutputMessages(response, isStreaming)
+    this._tagger.tagLLMIO(span, inputMessages, outputMessages)
+  }
+
+  #tagEmbedding (span, inputs, response, error) {
+    const embeddingInput = formatEmbeddingInput(inputs.contents)
+
+    if (error) {
+      this._tagger.tagEmbeddingIO(span, embeddingInput)
+      return
+    }
+
+    const embeddingOutput = formatEmbeddingOutput(response)
+    this._tagger.tagEmbeddingIO(span, embeddingInput, embeddingOutput)
+  }
+}
+
+module.exports = GenAiLLMObsPlugin

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

@@ -0,0 +1,486 @@
+'use strict'
+
+// Constants for role mapping
+const ROLES = {
+  MODEL: 'model',
+  ASSISTANT: 'assistant',
+  USER: 'user',
+  REASONING: 'reasoning'
+}
+
+/**
+ * Get the operation type from the method name
+ * @param {string} methodName
+ * @returns {'embedding' | 'llm'}
+ */
+function getOperation (methodName) {
+  return methodName.includes('embed') ? 'embedding' : 'llm'
+}
+
+/**
+ * Extract text parts from an array of parts
+ * @param {Array<{text?: string}>} parts
+ * @returns {string[]}
+ */
+function extractTextParts (parts) {
+  return parts
+    .filter(part => part.text)
+    .map(part => part.text)
+}
+
+/**
+ * Group parts by role (reasoning vs assistant)
+ * @param {Array<{text?: string, thought?: boolean}>} parts
+ * @returns {{reasoning: string, assistant: string}}
+ */
+function groupPartsByRole (parts) {
+  const grouped = {
+    reasoning: '',
+    assistant: ''
+  }
+
+  for (const part of parts) {
+    if (!part.text) continue
+
+    if (part.thought === true) {
+      grouped.reasoning += part.text
+    } else {
+      grouped.assistant += part.text
+    }
+  }
+
+  return grouped
+}
+
+/**
+ * Check if parts contain thought/reasoning content
+ * @param {Array<{thought?: boolean}>} parts
+ * @returns {boolean}
+ */
+function hasThoughtParts (parts) {
+  return parts.some(part => part.thought === true)
+}
+
+/**
+ * Determine the role from a candidate and its parts
+ * @param {object} candidate
+ * @param {Array<{thought?: boolean}>} parts
+ * @returns {string}
+ */
+function determineRole (candidate, parts = []) {
+  // Check parts for thought indicators
+  if (hasThoughtParts(parts)) {
+    return ROLES.REASONING
+  }
+
+  // Extract role from various possible locations
+  const rawRole = candidate.role ||
+                  candidate.content?.role ||
+                  candidate[0]?.content?.role
+
+  return normalizeRole(rawRole)
+}
+
+/**
+ * Normalize role to standard values
+ * @param {string} role
+ * @returns {string}
+ */
+function normalizeRole (role) {
+  if (role === ROLES.MODEL) return ROLES.ASSISTANT
+  if (role === ROLES.ASSISTANT) return ROLES.ASSISTANT
+  if (role === ROLES.USER) return ROLES.USER
+  if (role === ROLES.REASONING) return ROLES.REASONING
+  return ROLES.USER // default
+}
+
+/**
+ * Extract metrics from response
+ * @param {object} response
+ * @returns {object}
+ */
+function extractMetrics (response) {
+  const metrics = {}
+  const tokenUsage = response.usageMetadata
+
+  if (!tokenUsage) return metrics
+
+  if (tokenUsage.promptTokenCount) {
+    metrics.inputTokens = tokenUsage.promptTokenCount
+  }
+
+  if (tokenUsage.candidatesTokenCount) {
+    metrics.outputTokens = tokenUsage.candidatesTokenCount
+  }
+
+  const totalTokens = tokenUsage.totalTokenCount ||
+                     (tokenUsage.promptTokenCount || 0) + (tokenUsage.candidatesTokenCount || 0)
+  if (totalTokens) {
+    metrics.totalTokens = totalTokens
+  }
+
+  return metrics
+}
+
+/**
+ * Extract metadata from config
+ * @param {object} config
+ * @returns {object}
+ */
+function extractMetadata (config) {
+  if (!config) return {}
+
+  const fieldMap = {
+    temperature: 'temperature',
+    top_p: 'topP',
+    top_k: 'topK',
+    candidate_count: 'candidateCount',
+    max_output_tokens: 'maxOutputTokens',
+    stop_sequences: 'stopSequences',
+    response_logprobs: 'responseLogprobs',
+    logprobs: 'logprobs',
+    presence_penalty: 'presencePenalty',
+    frequency_penalty: 'frequencyPenalty',
+    seed: 'seed',
+    response_mime_type: 'responseMimeType',
+    safety_settings: 'safetySettings',
+    automatic_function_calling: 'automaticFunctionCalling'
+  }
+
+  const metadata = {}
+  for (const [metadataKey, configKey] of Object.entries(fieldMap)) {
+    metadata[metadataKey] = config[configKey] ?? null
+  }
+
+  return metadata
+}
+
+/**
+ * Format function call message
+ * @param {Array} parts
+ * @param {Array} functionCalls
+ * @param {string} role
+ * @returns {object}
+ */
+function formatFunctionCallMessage (parts, functionCalls, role) {
+  const toolCalls = functionCalls.map(part => ({
+    name: part.functionCall.name,
+    arguments: part.functionCall.args,
+    toolId: part.functionCall.id || '',
+    type: 'function_call'
+  }))
+
+  const textParts = extractTextParts(parts)
+  const content = textParts.length > 0 ? textParts.join('\n') : undefined
+  const message = { role, toolCalls }
+
+  if (content) message.content = content
+
+  return message
+}
+
+/**
+ * Format function response message
+ * @param {Array} functionResponses
+ * @param {string} role
+ * @returns {object}
+ */
+function formatFunctionResponseMessage (functionResponses, role) {
+  const toolResults = functionResponses.map(part => ({
+    name: part.functionResponse.name,
+    result: JSON.stringify(part.functionResponse.response),
+    toolId: part.functionResponse.id,
+    type: 'function_response'
+  }))
+
+  return {
+    role,
+    toolResults
+  }
+}
+
+/**
+ * Aggregate streaming chunks into a single response
+ * @param {Array} chunks
+ * @returns {object}
+ */
+function aggregateStreamingChunks (chunks) {
+  const response = { candidates: [] }
+
+  for (const chunk of chunks) {
+    if (chunk.candidates) {
+      // Flatten candidates array
+      response.candidates.push(...chunk.candidates)
+    }
+    if (chunk.usageMetadata) {
+      response.usageMetadata = chunk.usageMetadata
+    }
+  }
+
+  return response
+}
+
+/**
+ * Format a content object into a message
+ * @param {object} content
+ * @returns {object}
+ */
+function formatContentObject (content) {
+  const parts = content.parts || []
+  const role = determineRole(content, parts)
+
+  // Check if this is a thought/reasoning part
+  if (hasThoughtParts(parts)) {
+    return {
+      role: ROLES.REASONING,
+      content: extractTextParts(parts).join('\n')
+    }
+  }
+
+  // Check for function calls
+  const functionCalls = parts.filter(part => part.functionCall)
+  if (functionCalls.length > 0) {
+    return formatFunctionCallMessage(parts, functionCalls, role)
+  }
+
+  // Check for function responses
+  const functionResponses = parts.filter(part => part.functionResponse)
+  if (functionResponses.length > 0) {
+    return formatFunctionResponseMessage(functionResponses, role)
+  }
+
+  // Regular text content
+  return {
+    role,
+    content: extractTextParts(parts).join('\n')
+  }
+}
+
+/**
+ * Format input messages from contents
+ * @param {[string|{text: string}|{parts: Array<string|{text: string}>}]} [contents]
+ * @returns {Array}
+ */
+function formatInputMessages (contents) {
+  if (!contents) return []
+
+  const contentArray = Array.isArray(contents) ? contents : [contents]
+  const messages = []
+
+  for (const content of contentArray) {
+    if (typeof content === 'string') {
+      messages.push({ role: ROLES.USER, content })
+    } else if (content.text) {
+      messages.push({ role: ROLES.USER, content: content.text })
+    } else if (content.parts) {
+      const message = formatContentObject(content)
+      if (message) messages.push(message)
+    } else {
+      messages.push({ role: ROLES.USER, content: JSON.stringify(content) })
+    }
+  }
+
+  return messages
+}
+
+/**
+ * Format embedding input from contents
+ * @param {[string|{text: string}|{parts: Array<string|{text: string}>}]} [contents]
+ * @returns {Array}
+ */
+function formatEmbeddingInput (contents) {
+  if (!contents) return []
+
+  const contentArray = Array.isArray(contents) ? contents : [contents]
+  const documents = []
+
+  for (const content of contentArray) {
+    if (typeof content === 'string') {
+      documents.push({ text: content })
+    } else if (content.text) {
+      documents.push({ text: content.text })
+    } else if (content.parts) {
+      for (const part of content.parts) {
+        if (typeof part === 'string') {
+          documents.push({ text: part })
+        } else if (part.text) {
+          documents.push({ text: part.text })
+        }
+      }
+    }
+  }
+
+  return documents
+}
+
+/**
+ * Format a non-streaming candidate into messages
+ * @param {object} candidate
+ * @returns {Array}
+ */
+function formatNonStreamingCandidate (candidate) {
+  const messages = []
+  const content = Array.isArray(candidate) ? candidate[0].content : candidate.content
+
+  if (!content?.parts) return messages
+
+  const { parts } = content
+
+  // Check for function calls
+  const functionCalls = parts.filter(part => part.functionCall)
+  if (functionCalls.length > 0) {
+    messages.push(formatFunctionCallMessage(parts, functionCalls, ROLES.ASSISTANT))
+    return messages
+  }
+
+  // Check for executable code
+  const executableCode = parts.find(part => part.executableCode)
+  if (executableCode) {
+    messages.push({
+      role: ROLES.ASSISTANT,
+      content: JSON.stringify({
+        language: executableCode.executableCode.language,
+        code: executableCode.executableCode.code
+      })
+    })
+    return messages
+  }
+
+  // Check for code execution result
+  const codeExecutionResult = parts.find(part => part.codeExecutionResult)
+  if (codeExecutionResult) {
+    messages.push({
+      role: ROLES.ASSISTANT,
+      content: JSON.stringify({
+        outcome: codeExecutionResult.codeExecutionResult.outcome,
+        output: codeExecutionResult.codeExecutionResult.output
+      })
+    })
+    return messages
+  }
+
+  // Regular text content - may contain both reasoning and assistant parts
+  const partsByRole = groupPartsByRole(parts)
+
+  if (partsByRole.reasoning) {
+    messages.push({
+      role: ROLES.REASONING,
+      content: partsByRole.reasoning
+    })
+  }
+
+  if (partsByRole.assistant) {
+    messages.push({
+      role: ROLES.ASSISTANT,
+      content: partsByRole.assistant
+    })
+  }
+
+  return messages
+}
+
+/**
+ * Format streaming output from response
+ * @param {object} response
+ * @returns {Array}
+ */
+function formatStreamingOutput (response) {
+  const messages = []
+  const messagesByRole = new Map()
+
+  for (const candidate of response.candidates) {
+    const content = Array.isArray(candidate) ? candidate[0].content : candidate.content
+    if (!content?.parts) continue
+
+    // Skip special cases in streaming (handle them as non-streaming)
+    if (content.parts.some(part => part.functionCall ||
+                                   part.executableCode ||
+                                   part.codeExecutionResult)) {
+      messages.push(...formatNonStreamingCandidate(candidate))
+      continue
+    }
+
+    // Accumulate text parts by role
+    const partsByRole = groupPartsByRole(content.parts)
+
+    for (const [partRole, textContent] of Object.entries(partsByRole)) {
+      if (!textContent) continue
+
+      if (messagesByRole.has(partRole)) {
+        const index = messagesByRole.get(partRole)
+        messages[index].content += textContent
+      } else {
+        const messageIndex = messages.length
+        messages.push({ role: partRole, content: textContent })
+        messagesByRole.set(partRole, messageIndex)
+      }
+    }
+  }
+
+  return messages.length > 0 ? messages : [{ content: '' }]
+}
+
+/**
+ * Format non-streaming output from response
+ * @param {object} response
+ * @returns {Array}
+ */
+function formatNonStreamingOutput (response) {
+  const messages = []
+
+  for (const candidate of response.candidates) {
+    messages.push(...formatNonStreamingCandidate(candidate))
+  }
+
+  return messages.length > 0 ? messages : [{ content: '' }]
+}
+
+/**
+ * Format output messages from response
+ * @param {object} response
+ * @param {boolean} isStreaming
+ * @returns {Array}
+ */
+function formatOutputMessages (response, isStreaming = false) {
+  if (!response?.candidates?.length) {
+    return [{ content: '' }]
+  }
+
+  if (isStreaming) {
+    return formatStreamingOutput(response)
+  }
+
+  return formatNonStreamingOutput(response)
+}
+
+/**
+ * Format embedding output from response
+ * @param {object} response
+ * @returns {string}
+ */
+function formatEmbeddingOutput (response) {
+  if (!response?.embeddings?.length) {
+    return ''
+  }
+
+  const embeddingCount = response.embeddings.length
+  const firstEmbedding = response.embeddings[0]
+
+  if (firstEmbedding.values && Array.isArray(firstEmbedding.values)) {
+    const embeddingDim = firstEmbedding.values.length
+    return `[${embeddingCount} embedding(s) returned with size ${embeddingDim}]`
+  }
+
+  return `[${embeddingCount} embedding(s) returned]`
+}
+
+module.exports = {
+  getOperation,
+  extractMetrics,
+  extractMetadata,
+  aggregateStreamingChunks,
+  formatInputMessages,
+  formatEmbeddingInput,
+  formatOutputMessages,
+  formatEmbeddingOutput
+}

package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chain.js

@@ -1,7 +1,7 @@
 'use strict'
 
-const LangChainLLMObsHandler = require('.')
 const { spanHasError } = require('../../../util')
+const LangChainLLMObsHandler = require('.')
 
 class LangChainLLMObsChainHandler extends LangChainLLMObsHandler {
   setMetaTags ({ span, inputs, results }) {

package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chat_model.js

@@ -1,8 +1,8 @@
 'use strict'
 
-const LangChainLLMObsHandler = require('.')
 const LLMObsTagger = require('../../../tagger')
 const { spanHasError } = require('../../../util')
+const LangChainLLMObsHandler = require('.')
 
 const LLM = 'llm'