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

Datadog APM tracing client for JavaScript

'use strict'

const { channel } = require('dc-polyfill')

const { isError, isTrue } = require('../util')
const tracerVersion = require('../../../../package.json').version
const logger = require('../log')
const { getValueFromEnvSources } = require('../config/helper')
const Span = require('../opentracing/span')
const {
  SPAN_KIND,
  OUTPUT_VALUE,
  INPUT_VALUE,
  LLMOBS_TRACE_ID_BRIDGE_KEY,
  LLMOBS_PARENT_ID_BRIDGE_KEY,
} = require('./constants/tags')
const {
  getFunctionArguments,
  validateKind,
} = require('./util')
const { storage } = require('./storage')
const telemetry = require('./telemetry')
const LLMObsTagger = require('./tagger')

// communicating with writer
const evalMetricAppendCh = channel('llmobs:eval-metric:append')
const flushCh = channel('llmobs:writers:flush')
const registerUserSpanProcessorCh = channel('llmobs:register-processor')
const NoopLLMObs = require('./noop')

class LLMObs extends NoopLLMObs {
  /**
   * flag representing if a user span processor has been registered
   * @type {boolean}
   */
  #hasUserSpanProcessor = false

  /**
   * @param {import('../tracer')} tracer - Tracer instance
   * @param {import('./index')} llmobsModule - LLMObs module instance
   * @param {import('../config/config-base')} config - Tracer configuration
   */
  constructor (tracer, llmobsModule, config) {
    super(tracer)

    /** @type {import('../config/config-base')} */
    this._config = config

    this._llmobsModule = llmobsModule
    this._tagger = new LLMObsTagger(config)
  }

  get enabled () {
    return this._config.llmobs.enabled ?? false
  }

  enable (options = {}) {
    logger.warn(
      'Enabling LLM Observability via `llmobs.enable()` is deprecated and will be removed in dd-trace@7.0.0. ' +
      'Please instantiate LLM Observability via DD_LLMOBS_ENABLED or `tracer.init({ llmobs: ...options })`.'
    )

    if (this.enabled) {
      logger.debug('LLMObs is already enabled.')
      return
    }

    logger.debug('Enabling LLMObs')

    const DD_LLMOBS_ENABLED = getValueFromEnvSources('DD_LLMOBS_ENABLED')

    if (DD_LLMOBS_ENABLED != null && !isTrue(DD_LLMOBS_ENABLED)) {
      logger.debug('LLMObs.enable() called when DD_LLMOBS_ENABLED is false. No action taken.')
      return
    }

    // TODO: These configs should be passed through directly at construction time instead.
    this._config.llmobs.enabled = true
    this._config.llmobs.mlApp = options.mlApp
    this._config.llmobs.agentlessEnabled = options.agentlessEnabled

    // configure writers and channel subscribers
    this._llmobsModule.enable(this._config)
  }

  disable () {
    logger.warn(
      'Disabling LLM Observability via `llmobs.disable()` is deprecated and will be removed in dd-trace@7.0.0. ' +
      'Set DD_LLMOBS_ENABLED=false to disable LLM Observability.'
    )

    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

    telemetry.incrementLLMObsSpanStartCount({ autoinstrumented: false, 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, ...llmobsOptions }, () => fn(span, cb))
      )
    }

    return this._tracer.trace(name, spanOptions, span =>
      this.#activate(span, { kind, ...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 || 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 (...args) {
      telemetry.incrementLLMObsSpanStartCount({ autoinstrumented: false, kind })

      const span = llmobs._tracer.scope().active()
      const fnArgs = args

      const lastArgId = fnArgs.length - 1
      const cb = fnArgs[lastArgId]
      const hasCallback = typeof cb === 'function'

      if (hasCallback) {
        const scopeBoundCb = llmobs.#bind(cb)
        fnArgs[lastArgId] = function (...args) {
          // it is standard practice to follow the callback signature (err, result)
          // however, we try to parse the arguments to determine if the first argument is an error
          // if it is not, and is not undefined, we will use that for the output value
          const maybeError = args[0]
          const maybeResult = args[1]

          llmobs.#autoAnnotate(
            span,
            kind,
            getFunctionArguments(fn, fnArgs),
            isError(maybeError) || maybeError == null ? maybeResult : maybeError
          )

          return scopeBoundCb.apply(this, args)
        }
      }

      try {
        const result = llmobs.#activate(span, { kind, ...llmobsOptions }, () => fn.apply(this, fnArgs))

        if (result && typeof result.then === 'function') {
          return result.then(
            value => {
              if (!hasCallback) {
                llmobs.#autoAnnotate(span, kind, getFunctionArguments(fn, fnArgs), value)
              }
              return value
            },
            err => {
              llmobs.#autoAnnotate(span, kind, getFunctionArguments(fn, fnArgs))
              throw err
            }
          )
        }

        // it is possible to return a value and have a callback
        // however, since the span finishes when the callback is called, it is possible that
        // the callback is called before the function returns (although unlikely)
        // we do not want to throw for "annotating a finished span" in this case
        if (!hasCallback) {
          llmobs.#autoAnnotate(span, kind, getFunctionArguments(fn, fnArgs), result)
        }

        return result
      } catch (e) {
        llmobs.#autoAnnotate(span, kind, getFunctionArguments(fn, fnArgs))
        throw e
      }
    }

    return this._tracer.wrap(name, spanOptions, wrapped)
  }

  annotate (span, options, autoinstrumented = false) {
    if (!this.enabled) return

    if (!span) {
      span = this._active()
    }

    if ((span && !options) && !(span instanceof Span)) {
      options = span
      span = this._active()
    }

    let err = ''

    try {
      if (!span) {
        err = 'invalid_span_no_active_spans'
        throw new Error('No span provided and no active LLMObs-generated span found')
      }
      if (!options) {
        err = 'invalid_options'
        throw new Error('No options provided for annotation.')
      }

      if (!LLMObsTagger.tagMap.has(span)) {
        err = 'invalid_span_type'
        throw new Error('Span must be an LLMObs-generated span')
      }
      if (span._duration !== undefined) {
        err = 'invalid_finished_span'
        throw new Error('Cannot annotate a finished span')
      }

      const spanKind = LLMObsTagger.tagMap.get(span)[SPAN_KIND]
      if (!spanKind) {
        err = 'invalid_no_span_kind'
        throw new Error('LLMObs span must have a span kind specified')
      }

      const { inputData, outputData, metadata, metrics, tags, prompt, costTags } = 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)
      }
      // Apply tags before costTags so costTags can reference tags from the same annotation.
      if (tags) {
        this._tagger.tagSpanTags(span, tags)
      }
      if (costTags != null) {
        this._tagger.tagCostTags(span, costTags, 'annotate')
      }
      if (prompt) {
        this._tagger.tagPrompt(span, prompt)
      }
    } catch (e) {
      if (e.ddErrorTag) {
        err = e.ddErrorTag
      }
      throw e
    } finally {
      if (autoinstrumented === false) {
        telemetry.recordLLMObsAnnotate(span, err)
      }
    }
  }

  exportSpan (span) {
    span = span || this._active()
    let err = ''
    try {
      if (!span) {
        err = 'no_active_span'
        throw new Error('No span provided and no active LLMObs-generated span found')
      }
      if (!(span instanceof Span)) {
        err = 'invalid_span'
        throw new TypeError('Span must be a valid Span object.')
      }
      if (!LLMObsTagger.tagMap.has(span)) {
        err = 'invalid_span'
        throw new Error('Span must be an LLMObs-generated span')
      }
    } catch (e) {
      telemetry.recordExportSpan(span, err)
      throw e
    }
    try {
      return {
        traceId: span.context().toTraceId(true),
        spanId: span.context().toSpanId(),
      }
    } catch {
      err = 'invalid_span'
      logger.warn('Failed to export span. Span must be a valid Span object.')
    } finally {
      telemetry.recordExportSpan(span, err)
    }
  }

  registerProcessor (processor) {
    if (!this.enabled) return

    if (this.#hasUserSpanProcessor) {
      throw new Error(
        '[LLMObs] Only one user span processor can be registered. ' +
        'To register a new processor, deregister the existing processor first using `llmobs.deregisterProcessor()`.'
      )
    }

    this.#hasUserSpanProcessor = true
    registerUserSpanProcessorCh.publish(processor)
  }

  deregisterProcessor () {
    if (!this.enabled) return

    this.#hasUserSpanProcessor = false
    registerUserSpanProcessorCh.publish(null)
  }

  submitEvaluation (llmobsSpanContext, options = {}) {
    if (!this.enabled) return

    let err = ''
    const { traceId, spanId } = llmobsSpanContext
    try {
      if (!traceId || !spanId) {
        err = 'invalid_span'
        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) {
        err = 'missing_ml_app'
        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) {
        err = 'invalid_timestamp'
        throw new Error('timestampMs must be a non-negative integer. Evaluation metric data will not be sent')
      }

      const { label, value, tags, reasoning, assessment, metadata } = options
      const metricType = options.metricType?.toLowerCase()
      if (!label) {
        err = 'invalid_metric_label'
        throw new Error('label must be the specified name of the evaluation metric')
      }
      if (!metricType || !['categorical', 'score', 'boolean', 'json'].includes(metricType)) {
        err = 'invalid_metric_type'
        throw new Error('metricType must be one of "categorical", "score", "boolean" or "json"')
      }
      if (metricType === 'categorical' && typeof value !== 'string') {
        err = 'invalid_metric_value'
        throw new Error('value must be a string for a categorical metric.')
      }
      if (metricType === 'score' && typeof value !== 'number') {
        err = 'invalid_metric_value'
        throw new Error('value must be a number for a score metric.')
      }
      if (metricType === 'boolean' && typeof value !== 'boolean') {
        err = 'invalid_metric_value'
        throw new Error('value must be a boolean for a boolean metric')
      }
      if (metricType === 'json' && !(typeof value === 'object' && value != null && !Array.isArray(value))) {
        err = 'invalid_metric_value'
        throw new Error('value must be a JSON object for a json metric')
      }
      if (assessment != null && assessment !== 'pass' && assessment !== 'fail') {
        err = 'invalid_assessment'
        throw new Error('assessment must be pass or fail')
      }
      if (reasoning != null && typeof reasoning !== 'string') {
        err = 'invalid_reasoning'
        throw new Error('reasoning must be a string')
      }
      if (metadata != null && (typeof metadata !== 'object' || Array.isArray(metadata))) {
        err = 'invalid_metadata'
        throw new Error('metadata must be a JSON object')
      }

      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
            err = 'invalid_tags'
            throw new Error('Failed to parse tags. Tags for evaluation metrics must be strings')
          }
        }
      }

      // When OTel tracing is enabled, add source:otel tag to allow backend to wait for OTel span conversion
      if (this._config.DD_TRACE_OTEL_ENABLED) {
        evaluationTags.source = 'otel'
      }

      const payload = {
        join_on: {
          span: {
            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}`),
      }
      if (reasoning != null) {
        payload.reasoning = reasoning
      }
      if (metadata != null) {
        payload.metadata = metadata
      }
      if (assessment != null) {
        payload.assessment = assessment
      }
      const currentStore = storage.getStore()
      const routing = currentStore?.routingContext
      evalMetricAppendCh.publish({ payload, routing })
    } finally {
      telemetry.recordSubmitEvaluation(options, err)
    }
  }

  annotationContext (options, fn) {
    if (!this.enabled) return fn()

    const currentStore = storage.getStore()

    const store = {
      ...currentStore,
      annotationContext: {
        ...currentStore?.annotationContext,
        ...options,
      },
    }

    return storage.run(store, fn)
  }

  routingContext (options, fn) {
    if (!this.enabled) return fn()
    if (!options?.ddApiKey) {
      throw new Error('ddApiKey is required for routing context')
    }
    const currentStore = storage.getStore()
    if (currentStore?.routingContext) {
      logger.warn(
        '[LLM Observability] Nested routing context detected. Inner context will override outer context. ' +
        'Spans created in the inner context will only be sent to the inner context.'
      )
    }
    const store = {
      ...currentStore,
      routingContext: {
        apiKey: options.ddApiKey,
        site: options.ddSite,
      },
    }
    return storage.run(store, fn)
  }

  flush () {
    if (!this.enabled) return

    flushCh.publish()
  }

  #autoAnnotate (span, kind, input, output) {
    const annotations = {}
    if (input && !['llm', 'embedding'].includes(kind) && !LLMObsTagger.tagMap.get(span)?.[INPUT_VALUE]) {
      annotations.inputData = input
    }

    if (output && !['llm', 'retrieval'].includes(kind) && !LLMObsTagger.tagMap.get(span)?.[OUTPUT_VALUE]) {
      annotations.outputData = output
    }

    this.annotate(span, annotations, true)
  }

  _active () {
    const store = storage.getStore()
    return store?.span
  }

  #activate (span, options, fn) {
    const parentStore = storage.getStore()
    if (this.enabled) storage.enterWith({ ...parentStore, span })

    if (options) {
      this._tagger.registerLLMObsSpan(span, {
        ...options,
        parent: parentStore?.span,
      })

      // Bridge tags read by the dd-go LLMObs trace-indexer to correlate OTel
      // gen_ai.* spans with SDK LLMObs spans. Written once per local trace,
      // on the first successful SDK LLMObs span registration. The shared
      // _trace.tags bag is serialized to the first span in every flushed
      // chunk's meta, so partial flush is covered automatically without a
      // separate flush-time processor. Writing only after registerLLMObsSpan
      // succeeds avoids poisoning _trace.tags with bridge tags pointing at a
      // span that will never produce an LLMObs event.
      const traceTags = span?.context?.()._trace?.tags
      if (this.enabled && traceTags && !traceTags[LLMOBS_TRACE_ID_BRIDGE_KEY]) {
        traceTags[LLMOBS_TRACE_ID_BRIDGE_KEY] = span.context().toTraceId(true)
        traceTags[LLMOBS_PARENT_ID_BRIDGE_KEY] = span.context().toSpanId()
      }
    }

    try {
      return fn()
    } finally {
      if (this.enabled) storage.enterWith(parentStore)
    }
  }

  // bind function to active LLMObs span
  #bind (fn) {
    if (typeof fn !== 'function') return fn

    const llmobs = this
    const activeSpan = llmobs._active()

    return function (...args) {
      return llmobs.#activate(activeSpan, null, () => {
        return fn.apply(this, args)
      })
    }
  }

  #extractOptions (options) {
    const {
      modelName,
      modelProvider,
      sessionId,
      mlApp,
      _decorator,
      ...spanOptions
    } = options

    return {
      mlApp,
      modelName,
      modelProvider,
      sessionId,
      _decorator,
      spanOptions,
    }
  }
}

module.exports = LLMObs