@depup/dd-trace 5.90.0-depup.0

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

@@ -0,0 +1,288 @@
+'use strict'
+
+const LLMObsPlugin = require('./base')
+
+const ALLOWED_METADATA_KEYS = new Set([
+  'max_tokens',
+  'stop_sequences',
+  'temperature',
+  'top_k',
+  'top_p',
+])
+
+class AnthropicLLMObsPlugin extends LLMObsPlugin {
+  static integration = 'anthropic' // used for llmobs telemetry
+  static id = 'anthropic'
+  static prefix = 'tracing:apm:anthropic:request'
+
+  constructor () {
+    super(...arguments)
+
+    this.addSub('apm:anthropic:request:chunk', ({ ctx, chunk, done }) => {
+      ctx.chunks ??= []
+      const chunks = ctx.chunks
+      if (chunk) chunks.push(chunk)
+
+      if (!done) return
+
+      const response = { content: [] }
+
+      for (const chunk of chunks) {
+        switch (chunk.type) {
+          case 'message_start': {
+            const { message } = chunk
+            if (!message) continue
+
+            const { role, usage } = message
+            if (role) response.role = role
+            if (usage) response.usage = usage
+            break
+          }
+          case 'content_block_start': {
+            const contentBlock = chunk.content_block
+            if (!contentBlock) continue
+
+            const { type } = contentBlock
+            if (type === 'text') {
+              response.content.push({ type, text: contentBlock.text })
+            } else if (type === 'tool_use') {
+              response.content.push({ type, name: contentBlock.name, input: '', id: contentBlock.id })
+            }
+            break
+          }
+          case 'content_block_delta': {
+            const { delta } = chunk
+            if (!delta) continue
+
+            const { text } = delta
+            if (text) response.content[response.content.length - 1].text += text
+
+            const partialJson = delta.partial_json
+            if (partialJson && delta.type === 'input_json_delta') {
+              response.content[response.content.length - 1].input += partialJson
+            }
+            break
+          }
+          case 'content_block_stop': {
+            const type = response.content[response.content.length - 1].type
+            if (type === 'tool_use') {
+              const input = response.content[response.content.length - 1].input ?? '{}'
+              response.content[response.content.length - 1].input = JSON.parse(input)
+            }
+            break
+          }
+          case 'message_delta': {
+            const { delta } = chunk
+
+            const finishReason = delta?.stop_reason
+            if (finishReason) response.finish_reason = finishReason
+
+            const { usage } = chunk
+            if (usage) {
+              const responseUsage = response.usage ?? (response.usage = { input_tokens: 0, output_tokens: 0 })
+              responseUsage.output_tokens = usage.output_tokens
+
+              const cacheCreationTokens = usage.cache_creation_input_tokens
+              const cacheReadTokens = usage.cache_read_input_tokens
+              if (cacheCreationTokens) responseUsage.cache_creation_input_tokens = cacheCreationTokens
+              if (cacheReadTokens) responseUsage.cache_read_input_tokens = cacheReadTokens
+            }
+
+            break
+          }
+          case 'error': {
+            const { error } = chunk
+            if (!error) continue
+
+            response.error = {}
+            if (error.type) response.error.type = error.type
+            if (error.message) response.error.message = error.message
+
+            break
+          }
+        }
+
+        ctx.result = response
+      }
+    })
+  }
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    const { options } = ctx
+    const { model } = options
+
+    return {
+      kind: 'llm',
+      modelName: model,
+      modelProvider: 'anthropic',
+    }
+  }
+
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    if (!span) return
+
+    const { options, result } = ctx
+
+    this.#tagAnthropicInputMessages(span, options)
+    this.#tagAnthropicOutputMessages(span, result)
+    this.#tagAnthropicMetadata(span, options)
+    this.#tagAnthropicUsage(span, result)
+  }
+
+  #tagAnthropicInputMessages (span, options) {
+    const { system, messages } = options
+    const inputMessages = []
+
+    if (system) {
+      messages.unshift({ content: system, role: 'system' })
+    }
+
+    for (const message of messages) {
+      const { content, role } = message
+
+      if (typeof content === 'string') {
+        inputMessages.push({ content, role })
+        continue
+      }
+
+      for (const block of content) {
+        if (block.type === 'text') {
+          inputMessages.push({ content: block.text, role })
+        } else if (block.type === 'image') {
+          inputMessages.push({ content: '([IMAGE DETECTED])', role })
+        } else if (block.type === 'tool_use') {
+          const { text, name, id, type } = block
+          let input = block.input
+          if (typeof input === 'string') {
+            input = JSON.parse(input)
+          }
+
+          const toolCall = {
+            name,
+            arguments: input,
+            toolId: id,
+            type,
+          }
+
+          inputMessages.push({ content: text ?? '', role, toolCalls: [toolCall] })
+        } else if (block.type === 'tool_result') {
+          const { content } = block
+          const formattedContent = this.#formatAnthropicToolResultContent(content)
+          const toolResult = {
+            result: formattedContent,
+            toolId: block.tool_use_id,
+            type: 'tool_result',
+          }
+
+          inputMessages.push({ content: '', role, toolResults: [toolResult] })
+        } else {
+          inputMessages.push({ content: JSON.stringify(block), role })
+        }
+      }
+    }
+
+    this._tagger.tagLLMIO(span, inputMessages)
+  }
+
+  #tagAnthropicOutputMessages (span, result) {
+    if (!result) return
+
+    const { content, role } = result
+
+    if (typeof content === 'string') {
+      this._tagger.tagLLMIO(span, null, [{ content, role }])
+      return
+    }
+
+    const outputMessages = []
+    for (const block of content) {
+      const { text } = block
+      if (typeof text === 'string') {
+        outputMessages.push({ content: text, role })
+      } else if (block.type === 'tool_use') {
+        let input = block.input
+        if (typeof input === 'string') {
+          input = JSON.parse(input)
+        }
+
+        const toolCall = {
+          name: block.name,
+          arguments: input,
+          toolId: block.id,
+          type: block.type,
+        }
+
+        outputMessages.push({ content: text ?? '', role, toolCalls: [toolCall] })
+      }
+    }
+
+    this._tagger.tagLLMIO(span, null, outputMessages)
+  }
+
+  #tagAnthropicMetadata (span, options) {
+    const metadata = {}
+    for (const [key, value] of Object.entries(options)) {
+      if (ALLOWED_METADATA_KEYS.has(key)) {
+        metadata[key] = value
+      }
+    }
+
+    this._tagger.tagMetadata(span, metadata)
+  }
+
+  #tagAnthropicUsage (span, result) {
+    if (!result) return
+
+    const { usage } = result
+    if (!usage) return
+
+    const inputTokens = usage.input_tokens
+    const outputTokens = usage.output_tokens
+    const cacheWriteTokens = usage.cache_creation_input_tokens
+    const cacheReadTokens = usage.cache_read_input_tokens
+
+    const metrics = {
+      inputTokens: (inputTokens ?? 0) + (cacheWriteTokens ?? 0) + (cacheReadTokens ?? 0),
+    }
+
+    if (outputTokens) metrics.outputTokens = outputTokens
+    const totalTokens = metrics.inputTokens + (outputTokens ?? 0)
+    if (totalTokens) metrics.totalTokens = totalTokens
+
+    if (cacheWriteTokens != null) metrics.cacheWriteTokens = cacheWriteTokens
+    if (cacheReadTokens != null) metrics.cacheReadTokens = cacheReadTokens
+
+    const cacheCreation = usage.cache_creation
+    if (cacheCreation) {
+      metrics.cacheWrite5mTokens = cacheCreation.ephemeral_5m_input_tokens ?? 0
+      metrics.cacheWrite1hTokens = cacheCreation.ephemeral_1h_input_tokens ?? 0
+    } else if (cacheWriteTokens != null) {
+      metrics.cacheWrite5mTokens = cacheWriteTokens
+      metrics.cacheWrite1hTokens = 0
+    }
+
+    this._tagger.tagMetrics(span, metrics)
+  }
+
+  // maybe can make into a util file
+  #formatAnthropicToolResultContent (content) {
+    if (typeof content === 'string') {
+      return content
+    } else if (Array.isArray(content)) {
+      const formattedContent = []
+      for (const toolResultBlock of content) {
+        if (toolResultBlock.text) {
+          formattedContent.push(toolResultBlock.text)
+        } else if (toolResultBlock.type === 'image') {
+          formattedContent.push('([IMAGE DETECTED])')
+        }
+      }
+
+      return formattedContent.join(',')
+    }
+    return JSON.stringify(content)
+  }
+}
+
+module.exports = AnthropicLLMObsPlugin

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

@@ -0,0 +1,96 @@
+'use strict'
+
+const log = require('../../log')
+const { storage: llmobsStorage } = require('../storage')
+const telemetry = require('../telemetry')
+
+const TracingPlugin = require('../../plugins/tracing')
+const LLMObsTagger = require('../tagger')
+
+class LLMObsPlugin extends TracingPlugin {
+  constructor (...args) {
+    super(...args)
+
+    this._tagger = new LLMObsTagger(this._tracerConfig, true)
+  }
+
+  setLLMObsTags (ctx) {
+    throw new Error('setLLMObsTags must be implemented by the subclass')
+  }
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    throw new Error('getLLMObsSPanRegisterOptions must be implemented by the subclass')
+  }
+
+  start (ctx) {
+    // 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 parentStore = llmobsStorage.getStore()
+    const apmStore = ctx.currentStore
+    const span = apmStore?.span
+
+    const registerOptions = this.getLLMObsSpanRegisterOptions(ctx)
+
+    // register options may not be set for operations we do not trace with llmobs
+    // ie OpenAI fine tuning jobs, file jobs, etc.
+    if (registerOptions) {
+      telemetry.incrementLLMObsSpanStartCount({ autoinstrumented: true, integration: this.constructor.integration })
+
+      ctx.llmobs = {} // initialize context-based namespace
+      llmobsStorage.enterWith({ ...parentStore, span })
+      ctx.llmobs.parent = parentStore
+
+      this._tagger.registerLLMObsSpan(span, {
+        parent: parentStore?.span,
+        integration: this.constructor.integration,
+        ...registerOptions,
+      })
+    }
+  }
+
+  end (ctx) {
+    const enabled = this._tracerConfig.llmobs.enabled
+    if (!enabled) return
+
+    // only attempt to restore the context if the current span was an LLMObs span
+    const apmStore = ctx.currentStore
+    const span = apmStore?.span
+    if (!LLMObsTagger.tagMap.has(span)) return
+
+    const parentStore = ctx.llmobs.parent
+    llmobsStorage.enterWith(parentStore)
+  }
+
+  asyncEnd (ctx) {
+    // 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 apmStore = ctx.currentStore
+    const span = apmStore?.span
+    if (!span) {
+      log.debug(
+        'Tried to start an LLMObs span for %s without an active APM span. Not starting LLMObs span.',
+        this.constructor.name
+      )
+      return
+    }
+
+    this.setLLMObsTags(ctx)
+  }
+
+  configure (config) {
+    // we do not want to enable any LLMObs plugins if it is disabled on the tracer
+    const llmobsEnabled = this._tracerConfig.llmobs.enabled
+    if (llmobsEnabled === false) {
+      config = typeof config === 'boolean' ? false : { ...config, enabled: false } // override to false
+    }
+    super.configure(config)
+  }
+}
+
+module.exports = LLMObsPlugin

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

@@ -0,0 +1,137 @@
+'use strict'
+
+const { storage } = require('../../../../datadog-core')
+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'])
+
+const requestIdsToTokens = {}
+
+class BedrockRuntimeLLMObsPlugin extends BaseLLMObsPlugin {
+  constructor () {
+    super(...arguments)
+
+    this.addSub('apm:aws:request:complete:bedrockruntime', (ctx) => {
+      const { response } = ctx
+      const request = response.request
+      const operation = request.operation
+      // avoids instrumenting other non supported runtime operations
+      if (!ENABLED_OPERATIONS.has(operation)) {
+        return
+      }
+      const { modelProvider, modelName } = parseModelId(request.params.modelId)
+
+      // avoids instrumenting non llm type
+      if (modelName.includes('embed')) {
+        return
+      }
+      const span = ctx.currentStore?.span
+      this.setLLMObsTags({ ctx, request, span, response, modelProvider, modelName })
+    })
+
+    this.addSub('apm:aws:response:deserialize:bedrockruntime', ({ headers }) => {
+      const requestId = headers['x-amzn-requestid']
+      const inputTokenCount = headers['x-amzn-bedrock-input-token-count']
+      const outputTokenCount = headers['x-amzn-bedrock-output-token-count']
+      const cacheReadTokenCount = headers['x-amzn-bedrock-cache-read-input-token-count']
+      const cacheWriteTokenCount = headers['x-amzn-bedrock-cache-write-input-token-count']
+
+      requestIdsToTokens[requestId] = {
+        inputTokensFromHeaders: inputTokenCount && Number.parseInt(inputTokenCount),
+        outputTokensFromHeaders: outputTokenCount && Number.parseInt(outputTokenCount),
+        cacheReadTokensFromHeaders: cacheReadTokenCount && Number.parseInt(cacheReadTokenCount),
+        cacheWriteTokensFromHeaders: cacheWriteTokenCount && Number.parseInt(cacheWriteTokenCount),
+      }
+    })
+
+    this.addSub('apm:aws:response:streamed-chunk:bedrockruntime', ({ ctx, chunk }) => {
+      if (!ctx.chunks) ctx.chunks = []
+
+      if (chunk) ctx.chunks.push(chunk)
+    })
+  }
+
+  setLLMObsTags ({ ctx, request, span, response, modelProvider, modelName }) {
+    const isStream = request?.operation?.toLowerCase().includes('stream')
+    telemetry.incrementLLMObsSpanStartCount({ autoinstrumented: true, integration: 'bedrock' })
+
+    const parent = llmobsStore.getStore()?.span
+    this._tagger.registerLLMObsSpan(span, {
+      parent,
+      modelName: modelName.toLowerCase(),
+      modelProvider: modelProvider.toLowerCase(),
+      kind: 'llm',
+      name: 'bedrock-runtime.command',
+      integration: 'bedrock',
+    })
+
+    const requestParams = extractRequestParams(request.params, modelProvider)
+    // for streamed responses, we'll use the coerced response object we formed in the stream handler
+    const textAndResponseReason = isStream
+      ? extractTextAndResponseReasonFromStream(ctx.chunks, modelProvider, modelName)
+      : extractTextAndResponseReason(response, modelProvider, modelName)
+
+    // add metadata tags
+    this._tagger.tagMetadata(span, {
+      temperature: Number.parseFloat(requestParams.temperature) || 0,
+      max_tokens: Number.parseInt(requestParams.maxTokens) || 0,
+    })
+
+    // add I/O tags
+    this._tagger.tagLLMIO(
+      span,
+      requestParams.prompt,
+      [{ content: textAndResponseReason.message, role: textAndResponseReason.role }]
+    )
+
+    // add token metrics
+    const { inputTokens, outputTokens, totalTokens, cacheReadTokens, cacheWriteTokens } = extractTokens({
+      requestId: response.$metadata.requestId,
+      usage: textAndResponseReason.usage,
+    })
+    this._tagger.tagMetrics(span, {
+      inputTokens,
+      outputTokens,
+      totalTokens,
+      cacheReadTokens,
+      cacheWriteTokens,
+    })
+  }
+}
+
+function extractTokens ({ requestId, usage }) {
+  const {
+    inputTokensFromHeaders,
+    outputTokensFromHeaders,
+    cacheReadTokensFromHeaders,
+    cacheWriteTokensFromHeaders,
+  } = requestIdsToTokens[requestId] || {}
+  delete requestIdsToTokens[requestId]
+
+  const inputTokens = usage.inputTokens || inputTokensFromHeaders || 0
+  const outputTokens = usage.outputTokens || outputTokensFromHeaders || 0
+  const cacheReadTokens = usage.cacheReadTokens || cacheReadTokensFromHeaders || 0
+  const cacheWriteTokens = usage.cacheWriteTokens || cacheWriteTokensFromHeaders || 0
+
+  // adjust for the fact that bedrock input tokens only count non-cached tokens
+  const normalizedInputTokens = inputTokens + cacheReadTokens + cacheWriteTokens
+
+  return {
+    inputTokens: normalizedInputTokens,
+    outputTokens,
+    totalTokens: normalizedInputTokens + outputTokens,
+    cacheReadTokens,
+    cacheWriteTokens,
+  }
+}
+
+module.exports = BedrockRuntimeLLMObsPlugin

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