span_processor.js

'use strict'

const util = require('node:util')

const tracerVersion = require('../../../../package.json').version
const logger = require('../log')
const {
  ERROR_MESSAGE,
  ERROR_TYPE,
  ERROR_STACK,
} = require('../constants')
const PublicSpan = require('../opentracing/public/span')
const {
  SPAN_KIND,
  MODEL_NAME,
  MODEL_PROVIDER,
  METADATA,
  INPUT_MESSAGES,
  INPUT_VALUE,
  INTEGRATION,
  OUTPUT_MESSAGES,
  INPUT_DOCUMENTS,
  OUTPUT_DOCUMENTS,
  OUTPUT_VALUE,
  METRICS,
  ML_APP,
  TAGS,
  PARENT_ID_KEY,
  SESSION_ID,
  NAME,
  INPUT_PROMPT,
  ROUTING_API_KEY,
  ROUTING_SITE,
} = require('./constants/tags')
const { UNSERIALIZABLE_VALUE_TEXT } = require('./constants/text')
const telemetry = require('./telemetry')
const LLMObsTagger = require('./tagger')

/**
 * Resolves the tagMap key for a raw span from the finish channel.
 * Plugin-instrumented spans are registered with raw Span instances,
 * while SDK spans are registered with PublicSpan wrappers.
 */
function resolveTagMapKey (span) {
  if (LLMObsTagger.tagMap.has(span)) return span
  const wrapped = new PublicSpan(span)
  if (LLMObsTagger.tagMap.has(wrapped)) return wrapped
}

class LLMObservabilitySpan {
  constructor () {
    this.input = []
    this.output = []

    this._tags = {}
  }

  getTag (key) {
    return this._tags[key]
  }
}

class LLMObsSpanProcessor {
  /** @type {import('../config')} */
  #config

  /** @type {((span: LLMObservabilitySpan) => LLMObservabilitySpan | null) | null} */
  #userSpanProcessor

  /** @type {import('./writers/spans')} */
  #writer

  constructor (config) {
    this.#config = config
  }

  setUserSpanProcessor (userSpanProcessor) {
    this.#userSpanProcessor = userSpanProcessor
  }

  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
    // The finish channel always publishes raw Span instances, but the tagMap
    // may be keyed by either a raw Span (plugin-instrumented) or a PublicSpan
    // wrapper (SDK-instrumented). Resolve to whichever key is present.
    const tagMapKey = resolveTagMapKey(span)
    // if the span is not in our private tagger map, it is not an llmobs span
    if (!tagMapKey) return

    try {
      const formattedEvent = this.format(span, tagMapKey)
      telemetry.incrementLLMObsSpanFinishedCount(tagMapKey)
      if (formattedEvent == null) return

      const mlObsTags = LLMObsTagger.tagMap.get(tagMapKey)
      const routing = {
        apiKey: mlObsTags[ROUTING_API_KEY],
        site: mlObsTags[ROUTING_SITE],
      }

      this.#writer.append(formattedEvent, routing)
    } 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, tagMapKey) {
    const llmObsSpan = new LLMObservabilitySpan()
    let inputType, outputType

    const spanTags = span.context()._tags
    const mlObsTags = LLMObsTagger.tagMap.get(tagMapKey)

    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]) {
      llmObsSpan.input = mlObsTags[INPUT_MESSAGES]
      inputType = 'messages'
    } else if (spanKind === 'embedding' && mlObsTags[INPUT_DOCUMENTS]) {
      input.documents = mlObsTags[INPUT_DOCUMENTS]
    } else if (mlObsTags[INPUT_VALUE]) {
      llmObsSpan.input = [{ role: '', content: mlObsTags[INPUT_VALUE] }]
      inputType = 'value'
    }

    if (spanKind === 'llm' && mlObsTags[OUTPUT_MESSAGES]) {
      llmObsSpan.output = mlObsTags[OUTPUT_MESSAGES]
      outputType = 'messages'
    } else if (spanKind === 'retrieval' && mlObsTags[OUTPUT_DOCUMENTS]) {
      output.documents = mlObsTags[OUTPUT_DOCUMENTS]
    } else if (mlObsTags[OUTPUT_VALUE]) {
      llmObsSpan.output = [{ role: '', content: mlObsTags[OUTPUT_VALUE] }]
      outputType = 'value'
    }

    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
    }

    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 tags = this.#getTags(span, tagMapKey, mlApp, sessionId, error)
    llmObsSpan._tags = tags

    const processedSpan = this.#runProcessor(llmObsSpan)
    if (processedSpan === undefined) return null

    if (processedSpan.input) {
      if (inputType === 'messages') {
        input.messages = processedSpan.input
      } else if (inputType === 'value') {
        input.value = processedSpan.input[0].content
      }
    }

    if (processedSpan.output) {
      if (outputType === 'messages') {
        output.messages = processedSpan.output
      } else if (outputType === 'value') {
        output.value = processedSpan.output[0].content
      }
    }

    if (input) meta.input = input
    if (output) meta.output = output

    const prompt = mlObsTags[INPUT_PROMPT]
    if (prompt && spanKind === 'llm') {
      // by this point, we should have logged a warning if the span kind was not llm
      meta.input.prompt = prompt
    }

    const llmObsSpanEvent = {
      trace_id: span.context().toTraceId(true),
      span_id: span.context().toSpanId(),
      parent_id: parentId,
      name,
      tags: this.#objectTagsToStringArrayTags(tags),
      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) {
    // Capture root object by default
    const seenObjects = new WeakSet([obj])

    const isCircular = value => {
      if (value == null || 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.hasOwn(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 (value !== null && typeof value === 'object') {
          carrier[key] = Array.isArray(value) ? [] : {}
          add(value, carrier[key])
        } else {
          carrier[key] = value
        }
      }
    }

    add(obj, carrier)
  }

  #getTags (span, tagMapKey, mlApp, sessionId, error) {
    let tags = {
      ...this.#config.parsedDdTags,
      version: this.#config.version,
      env: this.#config.env,
      service: this.#config.service,
      source: 'integration',
      ml_app: mlApp,
      'ddtrace.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 integration = LLMObsTagger.tagMap.get(tagMapKey)?.[INTEGRATION]
    if (integration) tags.integration = integration

    const existingTags = LLMObsTagger.tagMap.get(tagMapKey)?.[TAGS] || {}
    if (existingTags) tags = { ...tags, ...existingTags }

    return tags
  }

  #objectTagsToStringArrayTags (tags) {
    return Object.entries(tags).map(([key, value]) => `${key}:${value ?? ''}`)
  }

  /**
   * Runs the user span processor, emitting telemetry and adding some guardrails against invalid return types
   * @param {LLMObservabilitySpan} span
   * @returns {LLMObservabilitySpan | undefined}
   */
  #runProcessor (span) {
    const processor = this.#userSpanProcessor
    if (!processor) return span

    let error = false

    try {
      const processedLLMObsSpan = processor(span)
      if (processedLLMObsSpan === null) return

      if (!(processedLLMObsSpan instanceof LLMObservabilitySpan)) {
        error = true
        logger.warn('User span processor must return an instance of an LLMObservabilitySpan or null, dropping span.')
        return
      }

      return processedLLMObsSpan
    } catch (e) {
      logger.error(`[LLMObs] Error in LLMObs span processor (${util.inspect(processor)}): ${util.inspect(e)}`)
      error = true
    } finally {
      telemetry.recordLLMObsUserProcessorCalled(error)
    }
  }
}

module.exports = LLMObsSpanProcessor