dd-trace 5.24.0 → 5.26.0

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

@@ -0,0 +1,65 @@
+'use strict'
+
+const log = require('../../log')
+const { storage } = require('../storage')
+
+const TracingPlugin = require('../../plugins/tracing')
+const LLMObsTagger = require('../tagger')
+
+// we make this a `Plugin` so we don't have to worry about `finish` being called
+class LLMObsPlugin extends TracingPlugin {
+  constructor (...args) {
+    super(...args)
+
+    this._tagger = new LLMObsTagger(this._tracerConfig, true)
+  }
+
+  getName () {}
+
+  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) {
+    const oldStore = storage.getStore()
+    const parent = oldStore?.span
+    const span = ctx.currentStore?.span
+
+    const registerOptions = this.getLLMObsSPanRegisterOptions(ctx)
+
+    this._tagger.registerLLMObsSpan(span, { parent, ...registerOptions })
+  }
+
+  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 span = ctx.currentStore?.span
+    if (!span) {
+      log.debug(
+        `Tried to start an LLMObs span for ${this.constructor.name} without an active APM span.
+        Not starting LLMObs span.`
+      )
+      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/openai.js

@@ -0,0 +1,205 @@
+'use strict'
+
+const LLMObsPlugin = require('./base')
+
+class OpenAiLLMObsPlugin extends LLMObsPlugin {
+  static get prefix () {
+    return 'tracing:apm:openai:request'
+  }
+
+  getLLMObsSPanRegisterOptions (ctx) {
+    const resource = ctx.methodName
+    const methodName = gateResource(normalizeOpenAIResourceName(resource))
+    if (!methodName) return // we will not trace all openai methods for llmobs
+
+    const inputs = ctx.args[0] // completion, chat completion, and embeddings take one argument
+    const operation = getOperation(methodName)
+    const kind = operation === 'embedding' ? 'embedding' : 'llm'
+    const name = `openai.${methodName}`
+
+    return {
+      modelProvider: 'openai',
+      modelName: inputs.model,
+      kind,
+      name
+    }
+  }
+
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    const resource = ctx.methodName
+    const methodName = gateResource(normalizeOpenAIResourceName(resource))
+    if (!methodName) return // we will not trace all openai methods for llmobs
+
+    const inputs = ctx.args[0] // completion, chat completion, and embeddings take one argument
+    const response = ctx.result?.data // no result if error
+    const error = !!span.context()._tags.error
+
+    const operation = getOperation(methodName)
+
+    if (operation === 'completion') {
+      this._tagCompletion(span, inputs, response, error)
+    } else if (operation === 'chat') {
+      this._tagChatCompletion(span, inputs, response, error)
+    } else if (operation === 'embedding') {
+      this._tagEmbedding(span, inputs, response, error)
+    }
+
+    if (!error) {
+      const metrics = this._extractMetrics(response)
+      this._tagger.tagMetrics(span, metrics)
+    }
+  }
+
+  _extractMetrics (response) {
+    const metrics = {}
+    const tokenUsage = response.usage
+
+    if (tokenUsage) {
+      const inputTokens = tokenUsage.prompt_tokens
+      if (inputTokens) metrics.inputTokens = inputTokens
+
+      const outputTokens = tokenUsage.completion_tokens
+      if (outputTokens) metrics.outputTokens = outputTokens
+
+      const totalTokens = tokenUsage.total_toksn || (inputTokens + outputTokens)
+      if (totalTokens) metrics.totalTokens = totalTokens
+    }
+
+    return metrics
+  }
+
+  _tagEmbedding (span, inputs, response, error) {
+    const { model, ...parameters } = inputs
+
+    const metadata = {
+      encoding_format: parameters.encoding_format || 'float'
+    }
+    if (inputs.dimensions) metadata.dimensions = inputs.dimensions
+    this._tagger.tagMetadata(span, metadata)
+
+    let embeddingInputs = inputs.input
+    if (!Array.isArray(embeddingInputs)) embeddingInputs = [embeddingInputs]
+    const embeddingInput = embeddingInputs.map(input => ({ text: input }))
+
+    if (error) {
+      this._tagger.tagEmbeddingIO(span, embeddingInput, undefined)
+      return
+    }
+
+    const float = Array.isArray(response.data[0].embedding)
+    let embeddingOutput
+    if (float) {
+      const embeddingDim = response.data[0].embedding.length
+      embeddingOutput = `[${response.data.length} embedding(s) returned with size ${embeddingDim}]`
+    } else {
+      embeddingOutput = `[${response.data.length} embedding(s) returned]`
+    }
+
+    this._tagger.tagEmbeddingIO(span, embeddingInput, embeddingOutput)
+  }
+
+  _tagCompletion (span, inputs, response, error) {
+    let { prompt, model, ...parameters } = inputs
+    if (!Array.isArray(prompt)) prompt = [prompt]
+
+    const completionInput = prompt.map(p => ({ content: p }))
+
+    const completionOutput = error ? [{ content: '' }] : response.choices.map(choice => ({ content: choice.text }))
+
+    this._tagger.tagLLMIO(span, completionInput, completionOutput)
+    this._tagger.tagMetadata(span, parameters)
+  }
+
+  _tagChatCompletion (span, inputs, response, error) {
+    const { messages, model, ...parameters } = inputs
+
+    if (error) {
+      this._tagger.tagLLMIO(span, messages, [{ content: '' }])
+      return
+    }
+
+    const outputMessages = []
+    const { choices } = response
+    for (const choice of choices) {
+      const message = choice.message || choice.delta
+      const content = message.content || ''
+      const role = message.role
+
+      if (message.function_call) {
+        const functionCallInfo = {
+          name: message.function_call.name,
+          arguments: JSON.parse(message.function_call.arguments)
+        }
+        outputMessages.push({ content, role, toolCalls: [functionCallInfo] })
+      } else if (message.tool_calls) {
+        const toolCallsInfo = []
+        for (const toolCall of message.tool_calls) {
+          const toolCallInfo = {
+            arguments: JSON.parse(toolCall.function.arguments),
+            name: toolCall.function.name,
+            toolId: toolCall.id,
+            type: toolCall.type
+          }
+          toolCallsInfo.push(toolCallInfo)
+        }
+        outputMessages.push({ content, role, toolCalls: toolCallsInfo })
+      } else {
+        outputMessages.push({ content, role })
+      }
+    }
+
+    this._tagger.tagLLMIO(span, messages, outputMessages)
+
+    const metadata = Object.entries(parameters).reduce((obj, [key, value]) => {
+      if (!['tools', 'functions'].includes(key)) {
+        obj[key] = value
+      }
+
+      return obj
+    }, {})
+
+    this._tagger.tagMetadata(span, metadata)
+  }
+}
+
+// TODO: this will be moved to the APM integration
+function normalizeOpenAIResourceName (resource) {
+  switch (resource) {
+    // completions
+    case 'completions.create':
+      return 'createCompletion'
+
+    // chat completions
+    case 'chat.completions.create':
+      return 'createChatCompletion'
+
+    // embeddings
+    case 'embeddings.create':
+      return 'createEmbedding'
+    default:
+      return resource
+  }
+}
+
+function gateResource (resource) {
+  return ['createCompletion', 'createChatCompletion', 'createEmbedding'].includes(resource)
+    ? resource
+    : undefined
+}
+
+function getOperation (resource) {
+  switch (resource) {
+    case 'createCompletion':
+      return 'completion'
+    case 'createChatCompletion':
+      return 'chat'
+    case 'createEmbedding':
+      return 'embedding'
+    default:
+      // should never happen
+      return 'unknown'
+  }
+}
+
+module.exports = OpenAiLLMObsPlugin

package/packages/dd-trace/src/llmobs/sdk.js

@@ -0,0 +1,377 @@
+'use strict'
+
+const { SPAN_KIND, OUTPUT_VALUE } = require('./constants/tags')
+
+const {
+  getFunctionArguments,
+  validateKind
+} = require('./util')
+const { isTrue } = require('../util')
+
+const { storage } = require('./storage')
+
+const Span = require('../opentracing/span')
+
+const tracerVersion = require('../../../../package.json').version
+const logger = require('../log')
+
+const LLMObsTagger = require('./tagger')
+
+// communicating with writer
+const { channel } = require('dc-polyfill')
+const evalMetricAppendCh = channel('llmobs:eval-metric:append')
+const flushCh = channel('llmobs:writers:flush')
+const NoopLLMObs = require('./noop')
+
+class LLMObs extends NoopLLMObs {
+  constructor (tracer, llmobsModule, config) {
+    super(tracer)
+
+    this._config = config
+    this._llmobsModule = llmobsModule
+    this._tagger = new LLMObsTagger(config)
+  }
+
+  get enabled () {
+    return this._config.llmobs.enabled
+  }
+
+  enable (options = {}) {
+    if (this.enabled) {
+      logger.debug('LLMObs is already enabled.')
+      return
+    }
+
+    logger.debug('Enabling LLMObs')
+
+    const { mlApp, agentlessEnabled } = options
+
+    const { DD_LLMOBS_ENABLED } = process.env
+
+    const llmobsConfig = {
+      mlApp,
+      agentlessEnabled
+    }
+
+    const enabled = DD_LLMOBS_ENABLED == null || isTrue(DD_LLMOBS_ENABLED)
+    if (!enabled) {
+      logger.debug('LLMObs.enable() called when DD_LLMOBS_ENABLED is false. No action taken.')
+      return
+    }
+
+    this._config.llmobs.enabled = true
+    this._config.configure({ ...this._config, llmobs: llmobsConfig })
+
+    // configure writers and channel subscribers
+    this._llmobsModule.enable(this._config)
+  }
+
+  disable () {
+    if (!this.enabled) {
+      logger.debug('LLMObs is already disabled.')
+      return
+    }
+
+    logger.debug('Disabling LLMObs')
+
+    this._config.llmobs.enabled = false
+
+    // disable writers and channel subscribers
+    this._llmobsModule.disable()
+  }
+
+  trace (options = {}, fn) {
+    if (typeof options === 'function') {
+      fn = options
+      options = {}
+    }
+
+    const kind = validateKind(options.kind) // will throw if kind is undefined or not an expected kind
+
+    // name is required for spans generated with `trace`
+    // while `kind` is required, this should never throw (as otherwise it would have thrown above)
+    const name = options.name || kind
+    if (!name) {
+      throw new Error('No span name provided for `trace`.')
+    }
+
+    const {
+      spanOptions,
+      ...llmobsOptions
+    } = this._extractOptions(options)
+
+    if (fn.length > 1) {
+      return this._tracer.trace(name, spanOptions, (span, cb) =>
+        this._activate(span, { kind, options: llmobsOptions }, () => fn(span, cb))
+      )
+    }
+
+    return this._tracer.trace(name, spanOptions, span =>
+      this._activate(span, { kind, options: llmobsOptions }, () => fn(span))
+    )
+  }
+
+  wrap (options = {}, fn) {
+    if (typeof options === 'function') {
+      fn = options
+      options = {}
+    }
+
+    const kind = validateKind(options.kind) // will throw if kind is undefined or not an expected kind
+    let name = options.name || (fn?.name ? fn.name : undefined) || kind
+
+    if (!name) {
+      logger.warn('No span name provided for `wrap`. Defaulting to "unnamed-anonymous-function".')
+      name = 'unnamed-anonymous-function'
+    }
+
+    const {
+      spanOptions,
+      ...llmobsOptions
+    } = this._extractOptions(options)
+
+    const llmobs = this
+
+    function wrapped () {
+      const span = llmobs._tracer.scope().active()
+
+      const result = llmobs._activate(span, { kind, options: llmobsOptions }, () => {
+        if (!['llm', 'embedding'].includes(kind)) {
+          llmobs.annotate(span, { inputData: getFunctionArguments(fn, arguments) })
+        }
+
+        return fn.apply(this, arguments)
+      })
+
+      if (result && typeof result.then === 'function') {
+        return result.then(value => {
+          if (value && !['llm', 'retrieval'].includes(kind) && !LLMObsTagger.tagMap.get(span)?.[OUTPUT_VALUE]) {
+            llmobs.annotate(span, { outputData: value })
+          }
+          return value
+        })
+      }
+
+      if (result && !['llm', 'retrieval'].includes(kind) && !LLMObsTagger.tagMap.get(span)?.[OUTPUT_VALUE]) {
+        llmobs.annotate(span, { outputData: result })
+      }
+
+      return result
+    }
+
+    return this._tracer.wrap(name, spanOptions, wrapped)
+  }
+
+  annotate (span, options) {
+    if (!this.enabled) return
+
+    if (!span) {
+      span = this._active()
+    }
+
+    if ((span && !options) && !(span instanceof Span)) {
+      options = span
+      span = this._active()
+    }
+
+    if (!span) {
+      throw new Error('No span provided and no active LLMObs-generated span found')
+    }
+    if (!options) {
+      throw new Error('No options provided for annotation.')
+    }
+
+    if (!LLMObsTagger.tagMap.has(span)) {
+      throw new Error('Span must be an LLMObs-generated span')
+    }
+    if (span._duration !== undefined) {
+      throw new Error('Cannot annotate a finished span')
+    }
+
+    const spanKind = LLMObsTagger.tagMap.get(span)[SPAN_KIND]
+    if (!spanKind) {
+      throw new Error('LLMObs span must have a span kind specified')
+    }
+
+    const { inputData, outputData, metadata, metrics, tags } = options
+
+    if (inputData || outputData) {
+      if (spanKind === 'llm') {
+        this._tagger.tagLLMIO(span, inputData, outputData)
+      } else if (spanKind === 'embedding') {
+        this._tagger.tagEmbeddingIO(span, inputData, outputData)
+      } else if (spanKind === 'retrieval') {
+        this._tagger.tagRetrievalIO(span, inputData, outputData)
+      } else {
+        this._tagger.tagTextIO(span, inputData, outputData)
+      }
+    }
+
+    if (metadata) {
+      this._tagger.tagMetadata(span, metadata)
+    }
+
+    if (metrics) {
+      this._tagger.tagMetrics(span, metrics)
+    }
+
+    if (tags) {
+      this._tagger.tagSpanTags(span, tags)
+    }
+  }
+
+  exportSpan (span) {
+    span = span || this._active()
+
+    if (!span) {
+      throw new Error('No span provided and no active LLMObs-generated span found')
+    }
+
+    if (!(span instanceof Span)) {
+      throw new Error('Span must be a valid Span object.')
+    }
+
+    if (!LLMObsTagger.tagMap.has(span)) {
+      throw new Error('Span must be an LLMObs-generated span')
+    }
+
+    try {
+      return {
+        traceId: span.context().toTraceId(true),
+        spanId: span.context().toSpanId()
+      }
+    } catch {
+      logger.warn('Faild to export span. Span must be a valid Span object.')
+    }
+  }
+
+  submitEvaluation (llmobsSpanContext, options = {}) {
+    if (!this.enabled) return
+
+    if (!this._config.apiKey) {
+      throw new Error(
+        'DD_API_KEY is required for sending evaluation metrics. Evaluation metric data will not be sent.\n' +
+        'Ensure this configuration is set before running your application.'
+      )
+    }
+
+    const { traceId, spanId } = llmobsSpanContext
+    if (!traceId || !spanId) {
+      throw new Error(
+        'spanId and traceId must both be specified for the given evaluation metric to be submitted.'
+      )
+    }
+
+    const mlApp = options.mlApp || this._config.llmobs.mlApp
+    if (!mlApp) {
+      throw new Error(
+        'ML App name is required for sending evaluation metrics. Evaluation metric data will not be sent.'
+      )
+    }
+
+    const timestampMs = options.timestampMs || Date.now()
+    if (typeof timestampMs !== 'number' || timestampMs < 0) {
+      throw new Error('timestampMs must be a non-negative integer. Evaluation metric data will not be sent')
+    }
+
+    const { label, value, tags } = options
+    const metricType = options.metricType?.toLowerCase()
+    if (!label) {
+      throw new Error('label must be the specified name of the evaluation metric')
+    }
+    if (!metricType || !['categorical', 'score'].includes(metricType)) {
+      throw new Error('metricType must be one of "categorical" or "score"')
+    }
+
+    if (metricType === 'categorical' && typeof value !== 'string') {
+      throw new Error('value must be a string for a categorical metric.')
+    }
+    if (metricType === 'score' && typeof value !== 'number') {
+      throw new Error('value must be a number for a score metric.')
+    }
+
+    const evaluationTags = {
+      'ddtrace.version': tracerVersion,
+      ml_app: mlApp
+    }
+
+    if (tags) {
+      for (const key in tags) {
+        const tag = tags[key]
+        if (typeof tag === 'string') {
+          evaluationTags[key] = tag
+        } else if (typeof tag.toString === 'function') {
+          evaluationTags[key] = tag.toString()
+        } else if (tag == null) {
+          evaluationTags[key] = Object.prototype.toString.call(tag)
+        } else {
+          // should be a rare case
+          // every object in JS has a toString, otherwise every primitive has its own toString
+          // null and undefined are handled above
+          throw new Error('Failed to parse tags. Tags for evaluation metrics must be strings')
+        }
+      }
+    }
+
+    const payload = {
+      span_id: spanId,
+      trace_id: traceId,
+      label,
+      metric_type: metricType,
+      ml_app: mlApp,
+      [`${metricType}_value`]: value,
+      timestamp_ms: timestampMs,
+      tags: Object.entries(evaluationTags).map(([key, value]) => `${key}:${value}`)
+    }
+
+    evalMetricAppendCh.publish(payload)
+  }
+
+  flush () {
+    if (!this.enabled) return
+
+    flushCh.publish()
+  }
+
+  _active () {
+    const store = storage.getStore()
+    return store?.span
+  }
+
+  _activate (span, { kind, options } = {}, fn) {
+    const parent = this._active()
+    if (this.enabled) storage.enterWith({ span })
+
+    this._tagger.registerLLMObsSpan(span, {
+      ...options,
+      parent,
+      kind
+    })
+
+    try {
+      return fn()
+    } finally {
+      if (this.enabled) storage.enterWith({ span: parent })
+    }
+  }
+
+  _extractOptions (options) {
+    const {
+      modelName,
+      modelProvider,
+      sessionId,
+      mlApp,
+      ...spanOptions
+    } = options
+
+    return {
+      mlApp,
+      modelName,
+      modelProvider,
+      sessionId,
+      spanOptions
+    }
+  }
+}
+
+module.exports = LLMObs