dd-trace 5.24.0 → 5.26.0

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,
+      '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 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 }

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

@@ -0,0 +1,322 @@
+'use strict'
+
+const log = require('../log')
+const {
+  MODEL_NAME,
+  MODEL_PROVIDER,
+  SESSION_ID,
+  ML_APP,
+  SPAN_KIND,
+  INPUT_VALUE,
+  OUTPUT_DOCUMENTS,
+  INPUT_DOCUMENTS,
+  OUTPUT_VALUE,
+  METADATA,
+  METRICS,
+  PARENT_ID_KEY,
+  INPUT_MESSAGES,
+  OUTPUT_MESSAGES,
+  TAGS,
+  NAME,
+  PROPAGATED_PARENT_ID_KEY,
+  ROOT_PARENT_ID,
+  INPUT_TOKENS_METRIC_KEY,
+  OUTPUT_TOKENS_METRIC_KEY,
+  TOTAL_TOKENS_METRIC_KEY
+} = require('./constants/tags')
+
+// global registry of LLMObs spans
+// maps LLMObs spans to their annotations
+const registry = new WeakMap()
+
+class LLMObsTagger {
+  constructor (config, softFail = false) {
+    this._config = config
+
+    this.softFail = softFail
+  }
+
+  static get tagMap () {
+    return registry
+  }
+
+  registerLLMObsSpan (span, {
+    modelName,
+    modelProvider,
+    sessionId,
+    mlApp,
+    parent,
+    kind,
+    name
+  } = {}) {
+    if (!this._config.llmobs.enabled) return
+    if (!kind) return // do not register it in the map if it doesn't have an llmobs span kind
+
+    this._register(span)
+
+    if (name) this._setTag(span, NAME, name)
+
+    this._setTag(span, SPAN_KIND, kind)
+    if (modelName) this._setTag(span, MODEL_NAME, modelName)
+    if (modelProvider) this._setTag(span, MODEL_PROVIDER, modelProvider)
+
+    sessionId = sessionId || parent?.context()._tags[SESSION_ID]
+    if (sessionId) this._setTag(span, SESSION_ID, sessionId)
+
+    if (!mlApp) mlApp = parent?.context()._tags[ML_APP] || this._config.llmobs.mlApp
+    this._setTag(span, ML_APP, mlApp)
+
+    const parentId =
+      parent?.context().toSpanId() ||
+      span.context()._trace.tags[PROPAGATED_PARENT_ID_KEY] ||
+      ROOT_PARENT_ID
+    this._setTag(span, PARENT_ID_KEY, parentId)
+  }
+
+  // TODO: similarly for the following `tag` methods,
+  // how can we transition from a span weakmap to core API functionality
+  tagLLMIO (span, inputData, outputData) {
+    this._tagMessages(span, inputData, INPUT_MESSAGES)
+    this._tagMessages(span, outputData, OUTPUT_MESSAGES)
+  }
+
+  tagEmbeddingIO (span, inputData, outputData) {
+    this._tagDocuments(span, inputData, INPUT_DOCUMENTS)
+    this._tagText(span, outputData, OUTPUT_VALUE)
+  }
+
+  tagRetrievalIO (span, inputData, outputData) {
+    this._tagText(span, inputData, INPUT_VALUE)
+    this._tagDocuments(span, outputData, OUTPUT_DOCUMENTS)
+  }
+
+  tagTextIO (span, inputData, outputData) {
+    this._tagText(span, inputData, INPUT_VALUE)
+    this._tagText(span, outputData, OUTPUT_VALUE)
+  }
+
+  tagMetadata (span, metadata) {
+    this._setTag(span, METADATA, metadata)
+  }
+
+  tagMetrics (span, metrics) {
+    const filterdMetrics = {}
+    for (const [key, value] of Object.entries(metrics)) {
+      let processedKey = key
+
+      // processing these specifically for our metrics ingestion
+      switch (key) {
+        case 'inputTokens':
+          processedKey = INPUT_TOKENS_METRIC_KEY
+          break
+        case 'outputTokens':
+          processedKey = OUTPUT_TOKENS_METRIC_KEY
+          break
+        case 'totalTokens':
+          processedKey = TOTAL_TOKENS_METRIC_KEY
+          break
+      }
+
+      if (typeof value === 'number') {
+        filterdMetrics[processedKey] = value
+      } else {
+        this._handleFailure(`Value for metric '${key}' must be a number, instead got ${value}`)
+      }
+    }
+
+    this._setTag(span, METRICS, filterdMetrics)
+  }
+
+  tagSpanTags (span, tags) {
+    // new tags will be merged with existing tags
+    const currentTags = registry.get(span)?.[TAGS]
+    if (currentTags) {
+      Object.assign(tags, currentTags)
+    }
+    this._setTag(span, TAGS, tags)
+  }
+
+  _tagText (span, data, key) {
+    if (data) {
+      if (typeof data === 'string') {
+        this._setTag(span, key, data)
+      } else {
+        try {
+          this._setTag(span, key, JSON.stringify(data))
+        } catch {
+          const type = key === INPUT_VALUE ? 'input' : 'output'
+          this._handleFailure(`Failed to parse ${type} value, must be JSON serializable.`)
+        }
+      }
+    }
+  }
+
+  _tagDocuments (span, data, key) {
+    if (data) {
+      if (!Array.isArray(data)) {
+        data = [data]
+      }
+
+      const documents = data.map(document => {
+        if (typeof document === 'string') {
+          return { text: document }
+        }
+
+        if (document == null || typeof document !== 'object') {
+          this._handleFailure('Documents must be a string, object, or list of objects.')
+          return undefined
+        }
+
+        const { text, name, id, score } = document
+        let validDocument = true
+
+        if (typeof text !== 'string') {
+          this._handleFailure('Document text must be a string.')
+          validDocument = false
+        }
+
+        const documentObj = { text }
+
+        validDocument = this._tagConditionalString(name, 'Document name', documentObj, 'name') && validDocument
+        validDocument = this._tagConditionalString(id, 'Document ID', documentObj, 'id') && validDocument
+        validDocument = this._tagConditionalNumber(score, 'Document score', documentObj, 'score') && validDocument
+
+        return validDocument ? documentObj : undefined
+      }).filter(doc => !!doc)
+
+      if (documents.length) {
+        this._setTag(span, key, documents)
+      }
+    }
+  }
+
+  _tagMessages (span, data, key) {
+    if (data) {
+      if (!Array.isArray(data)) {
+        data = [data]
+      }
+
+      const messages = data.map(message => {
+        if (typeof message === 'string') {
+          return { content: message }
+        }
+
+        if (message == null || typeof message !== 'object') {
+          this._handleFailure('Messages must be a string, object, or list of objects')
+          return undefined
+        }
+
+        let validMessage = true
+
+        const { content = '', role } = message
+        let toolCalls = message.toolCalls
+        const messageObj = { content }
+
+        if (typeof content !== 'string') {
+          this._handleFailure('Message content must be a string.')
+          validMessage = false
+        }
+
+        validMessage = this._tagConditionalString(role, 'Message role', messageObj, 'role') && validMessage
+
+        if (toolCalls) {
+          if (!Array.isArray(toolCalls)) {
+            toolCalls = [toolCalls]
+          }
+
+          const filteredToolCalls = toolCalls.map(toolCall => {
+            if (typeof toolCall !== 'object') {
+              this._handleFailure('Tool call must be an object.')
+              return undefined
+            }
+
+            let validTool = true
+
+            const { name, arguments: args, toolId, type } = toolCall
+            const toolCallObj = {}
+
+            validTool = this._tagConditionalString(name, 'Tool name', toolCallObj, 'name') && validTool
+            validTool = this._tagConditionalObject(args, 'Tool arguments', toolCallObj, 'arguments') && validTool
+            validTool = this._tagConditionalString(toolId, 'Tool ID', toolCallObj, 'tool_id') && validTool
+            validTool = this._tagConditionalString(type, 'Tool type', toolCallObj, 'type') && validTool
+
+            return validTool ? toolCallObj : undefined
+          }).filter(toolCall => !!toolCall)
+
+          if (filteredToolCalls.length) {
+            messageObj.tool_calls = filteredToolCalls
+          }
+        }
+
+        return validMessage ? messageObj : undefined
+      }).filter(msg => !!msg)
+
+      if (messages.length) {
+        this._setTag(span, key, messages)
+      }
+    }
+  }
+
+  _tagConditionalString (data, type, carrier, key) {
+    if (!data) return true
+    if (typeof data !== 'string') {
+      this._handleFailure(`"${type}" must be a string.`)
+      return false
+    }
+    carrier[key] = data
+    return true
+  }
+
+  _tagConditionalNumber (data, type, carrier, key) {
+    if (!data) return true
+    if (typeof data !== 'number') {
+      this._handleFailure(`"${type}" must be a number.`)
+      return false
+    }
+    carrier[key] = data
+    return true
+  }
+
+  _tagConditionalObject (data, type, carrier, key) {
+    if (!data) return true
+    if (typeof data !== 'object') {
+      this._handleFailure(`"${type}" must be an object.`)
+      return false
+    }
+    carrier[key] = data
+    return true
+  }
+
+  // any public-facing LLMObs APIs using this tagger should not soft fail
+  // auto-instrumentation should soft fail
+  _handleFailure (msg) {
+    if (this.softFail) {
+      log.warn(msg)
+    } else {
+      throw new Error(msg)
+    }
+  }
+
+  _register (span) {
+    if (!this._config.llmobs.enabled) return
+    if (registry.has(span)) {
+      this._handleFailure(`LLMObs Span "${span._name}" already registered.`)
+      return
+    }
+
+    registry.set(span, {})
+  }
+
+  _setTag (span, key, value) {
+    if (!this._config.llmobs.enabled) return
+    if (!registry.has(span)) {
+      this._handleFailure('Span must be an LLMObs generated span.')
+      return
+    }
+
+    const tagsCarrier = registry.get(span)
+    Object.assign(tagsCarrier, { [key]: value })
+  }
+}
+
+module.exports = LLMObsTagger

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

@@ -0,0 +1,176 @@
+'use strict'
+
+const { SPAN_KINDS } = require('./constants/tags')
+
+function encodeUnicode (str) {
+  if (!str) return str
+  return str.split('').map(char => {
+    const code = char.charCodeAt(0)
+    if (code > 127) {
+      return `\\u${code.toString(16).padStart(4, '0')}`
+    }
+    return char
+  }).join('')
+}
+
+function validateKind (kind) {
+  if (!SPAN_KINDS.includes(kind)) {
+    throw new Error(`
+      Invalid span kind specified: "${kind}"
+      Must be one of: ${SPAN_KINDS.join(', ')}
+    `)
+  }
+
+  return kind
+}
+
+// extracts the argument names from a function string
+function parseArgumentNames (str) {
+  const result = []
+  let current = ''
+  let closerCount = 0
+  let recording = true
+  let inSingleLineComment = false
+  let inMultiLineComment = false
+
+  for (let i = 0; i < str.length; i++) {
+    const char = str[i]
+    const nextChar = str[i + 1]
+
+    // Handle single-line comments
+    if (!inMultiLineComment && char === '/' && nextChar === '/') {
+      inSingleLineComment = true
+      i++ // Skip the next character
+      continue
+    }
+
+    // Handle multi-line comments
+    if (!inSingleLineComment && char === '/' && nextChar === '*') {
+      inMultiLineComment = true
+      i++ // Skip the next character
+      continue
+    }
+
+    // End of single-line comment
+    if (inSingleLineComment && char === '\n') {
+      inSingleLineComment = false
+      continue
+    }
+
+    // End of multi-line comment
+    if (inMultiLineComment && char === '*' && nextChar === '/') {
+      inMultiLineComment = false
+      i++ // Skip the next character
+      continue
+    }
+
+    // Skip characters inside comments
+    if (inSingleLineComment || inMultiLineComment) {
+      continue
+    }
+
+    if (['{', '[', '('].includes(char)) {
+      closerCount++
+    } else if (['}', ']', ')'].includes(char)) {
+      closerCount--
+    } else if (char === '=' && nextChar !== '>' && closerCount === 0) {
+      recording = false
+      // record the variable name early, and stop counting characters until we reach the next comma
+      result.push(current.trim())
+      current = ''
+      continue
+    } else if (char === ',' && closerCount === 0) {
+      if (recording) {
+        result.push(current.trim())
+        current = ''
+      }
+
+      recording = true
+      continue
+    }
+
+    if (recording) {
+      current += char
+    }
+  }
+
+  if (current && recording) {
+    result.push(current.trim())
+  }
+
+  return result
+}
+
+// finds the bounds of the arguments in a function string
+function findArgumentsBounds (str) {
+  let start = -1
+  let end = -1
+  let closerCount = 0
+
+  for (let i = 0; i < str.length; i++) {
+    const char = str[i]
+
+    if (char === '(') {
+      if (closerCount === 0) {
+        start = i
+      }
+
+      closerCount++
+    } else if (char === ')') {
+      closerCount--
+
+      if (closerCount === 0) {
+        end = i
+        break
+      }
+    }
+  }
+
+  return { start, end }
+}
+
+const memo = new WeakMap()
+function getFunctionArguments (fn, args = []) {
+  if (!fn) return
+  if (!args.length) return
+  if (args.length === 1) return args[0]
+
+  try {
+    let names
+    if (memo.has(fn)) {
+      names = memo.get(fn)
+    } else {
+      const fnString = fn.toString()
+      const { start, end } = findArgumentsBounds(fnString)
+      names = parseArgumentNames(fnString.slice(start + 1, end))
+      memo.set(fn, names)
+    }
+
+    const argsObject = {}
+
+    for (const argIdx in args) {
+      const name = names[argIdx]
+      const arg = args[argIdx]
+
+      const spread = name?.startsWith('...')
+
+      // this can only be the last argument
+      if (spread) {
+        argsObject[name.slice(3)] = args.slice(argIdx)
+        break
+      }
+
+      argsObject[name] = arg
+    }
+
+    return argsObject
+  } catch {
+    return args
+  }
+}
+
+module.exports = {
+  encodeUnicode,
+  validateKind,
+  getFunctionArguments
+}