dd-trace 4.47.1 → 4.49.0

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 = {
+      'dd-trace.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

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

@@ -0,0 +1,195 @@
+'use strict'
+
+const {
+  SPAN_KIND,
+  MODEL_NAME,
+  MODEL_PROVIDER,
+  METADATA,
+  INPUT_MESSAGES,
+  INPUT_VALUE,
+  OUTPUT_MESSAGES,
+  INPUT_DOCUMENTS,
+  OUTPUT_DOCUMENTS,
+  OUTPUT_VALUE,
+  METRICS,
+  ML_APP,
+  TAGS,
+  PARENT_ID_KEY,
+  SESSION_ID,
+  NAME
+} = require('./constants/tags')
+const { UNSERIALIZABLE_VALUE_TEXT } = require('./constants/text')
+
+const {
+  ERROR_MESSAGE,
+  ERROR_TYPE,
+  ERROR_STACK
+} = require('../constants')
+
+const LLMObsTagger = require('./tagger')
+
+const tracerVersion = require('../../../../package.json').version
+const logger = require('../log')
+
+class LLMObsSpanProcessor {
+  constructor (config) {
+    this._config = config
+  }
+
+  setWriter (writer) {
+    this._writer = writer
+  }
+
+  // TODO: instead of relying on the tagger's weakmap registry, can we use some namespaced storage correlation?
+  process ({ span }) {
+    if (!this._config.llmobs.enabled) return
+    // if the span is not in our private tagger map, it is not an llmobs span
+    if (!LLMObsTagger.tagMap.has(span)) return
+
+    try {
+      const formattedEvent = this.format(span)
+      this._writer.append(formattedEvent)
+    } catch (e) {
+      // this should be a rare case
+      // we protect against unserializable properties in the format function, and in
+      // safeguards in the tagger
+      logger.warn(`
+        Failed to append span to LLM Observability writer, likely due to an unserializable property.
+        Span won't be sent to LLM Observability: ${e.message}
+      `)
+    }
+  }
+
+  format (span) {
+    const spanTags = span.context()._tags
+    const mlObsTags = LLMObsTagger.tagMap.get(span)
+
+    const spanKind = mlObsTags[SPAN_KIND]
+
+    const meta = { 'span.kind': spanKind, input: {}, output: {} }
+    const input = {}
+    const output = {}
+
+    if (['llm', 'embedding'].includes(spanKind)) {
+      meta.model_name = mlObsTags[MODEL_NAME] || 'custom'
+      meta.model_provider = (mlObsTags[MODEL_PROVIDER] || 'custom').toLowerCase()
+    }
+    if (mlObsTags[METADATA]) {
+      this._addObject(mlObsTags[METADATA], meta.metadata = {})
+    }
+    if (spanKind === 'llm' && mlObsTags[INPUT_MESSAGES]) {
+      input.messages = mlObsTags[INPUT_MESSAGES]
+    }
+    if (mlObsTags[INPUT_VALUE]) {
+      input.value = mlObsTags[INPUT_VALUE]
+    }
+    if (spanKind === 'llm' && mlObsTags[OUTPUT_MESSAGES]) {
+      output.messages = mlObsTags[OUTPUT_MESSAGES]
+    }
+    if (spanKind === 'embedding' && mlObsTags[INPUT_DOCUMENTS]) {
+      input.documents = mlObsTags[INPUT_DOCUMENTS]
+    }
+    if (mlObsTags[OUTPUT_VALUE]) {
+      output.value = mlObsTags[OUTPUT_VALUE]
+    }
+    if (spanKind === 'retrieval' && mlObsTags[OUTPUT_DOCUMENTS]) {
+      output.documents = mlObsTags[OUTPUT_DOCUMENTS]
+    }
+
+    const error = spanTags.error || spanTags[ERROR_TYPE]
+    if (error) {
+      meta[ERROR_MESSAGE] = spanTags[ERROR_MESSAGE] || error.message || error.code
+      meta[ERROR_TYPE] = spanTags[ERROR_TYPE] || error.name
+      meta[ERROR_STACK] = spanTags[ERROR_STACK] || error.stack
+    }
+
+    if (input) meta.input = input
+    if (output) meta.output = output
+
+    const metrics = mlObsTags[METRICS] || {}
+
+    const mlApp = mlObsTags[ML_APP]
+    const sessionId = mlObsTags[SESSION_ID]
+    const parentId = mlObsTags[PARENT_ID_KEY]
+
+    const name = mlObsTags[NAME] || span._name
+
+    const llmObsSpanEvent = {
+      trace_id: span.context().toTraceId(true),
+      span_id: span.context().toSpanId(),
+      parent_id: parentId,
+      name,
+      tags: this._processTags(span, mlApp, sessionId, error),
+      start_ns: Math.round(span._startTime * 1e6),
+      duration: Math.round(span._duration * 1e6),
+      status: error ? 'error' : 'ok',
+      meta,
+      metrics,
+      _dd: {
+        span_id: span.context().toSpanId(),
+        trace_id: span.context().toTraceId(true)
+      }
+    }
+
+    if (sessionId) llmObsSpanEvent.session_id = sessionId
+
+    return llmObsSpanEvent
+  }
+
+  // For now, this only applies to metadata, as we let users annotate this field with any object
+  // However, we want to protect against circular references or BigInts (unserializable)
+  // This function can be reused for other fields if needed
+  // Messages, Documents, and Metrics are safeguarded in `llmobs/tagger.js`
+  _addObject (obj, carrier) {
+    const seenObjects = new WeakSet()
+    seenObjects.add(obj) // capture root object
+
+    const isCircular = value => {
+      if (typeof value !== 'object') return false
+      if (seenObjects.has(value)) return true
+      seenObjects.add(value)
+      return false
+    }
+
+    const add = (obj, carrier) => {
+      for (const key in obj) {
+        const value = obj[key]
+        if (!Object.prototype.hasOwnProperty.call(obj, key)) continue
+        if (typeof value === 'bigint' || isCircular(value)) {
+          // mark as unserializable instead of dropping
+          logger.warn(`Unserializable property found in metadata: ${key}`)
+          carrier[key] = UNSERIALIZABLE_VALUE_TEXT
+          continue
+        }
+        if (typeof value === 'object') {
+          add(value, carrier[key] = {})
+        } else {
+          carrier[key] = value
+        }
+      }
+    }
+
+    add(obj, carrier)
+  }
+
+  _processTags (span, mlApp, sessionId, error) {
+    let tags = {
+      version: this._config.version,
+      env: this._config.env,
+      service: this._config.service,
+      source: 'integration',
+      ml_app: mlApp,
+      'dd-trace.version': tracerVersion,
+      error: Number(!!error) || 0,
+      language: 'javascript'
+    }
+    const errType = span.context()._tags[ERROR_TYPE] || error?.name
+    if (errType) tags.error_type = errType
+    if (sessionId) tags.session_id = sessionId
+    const existingTags = LLMObsTagger.tagMap.get(span)?.[TAGS] || {}
+    if (existingTags) tags = { ...tags, ...existingTags }
+    return Object.entries(tags).map(([key, value]) => `${key}:${value ?? ''}`)
+  }
+}
+
+module.exports = LLMObsSpanProcessor

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

@@ -0,0 +1,7 @@
+'use strict'
+
+// TODO: remove this and use namespaced storage once available
+const { AsyncLocalStorage } = require('async_hooks')
+const storage = new AsyncLocalStorage()
+
+module.exports = { storage }