dd-trace 5.30.0 → 5.32.0

package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/chat_model.js

@@ -0,0 +1,111 @@
+'use strict'
+
+const LangChainLLMObsHandler = require('.')
+const LLMObsTagger = require('../../../tagger')
+const { spanHasError } = require('../../../util')
+
+const LLM = 'llm'
+
+class LangChainLLMObsChatModelHandler extends LangChainLLMObsHandler {
+  setMetaTags ({ span, inputs, results, options, integrationName }) {
+    if (integrationName === 'openai' && options?.response_format) {
+      // langchain-openai will call a beta client if "response_format" is passed in on the options object
+      // we do not trace these calls, so this should be an llm span
+      this._tagger.changeKind(span, LLM)
+    }
+    const spanKind = LLMObsTagger.getSpanKind(span)
+    const isWorkflow = spanKind === 'workflow'
+
+    const inputMessages = []
+    if (!Array.isArray(inputs)) inputs = [inputs]
+
+    for (const messageSet of inputs) {
+      for (const message of messageSet) {
+        const content = message.content || ''
+        const role = this.getRole(message)
+        inputMessages.push({ content, role })
+      }
+    }
+
+    if (spanHasError(span)) {
+      if (isWorkflow) {
+        this._tagger.tagTextIO(span, inputMessages, [{ content: '' }])
+      } else {
+        this._tagger.tagLLMIO(span, inputMessages, [{ content: '' }])
+      }
+      return
+    }
+
+    const outputMessages = []
+    let inputTokens = 0
+    let outputTokens = 0
+    let totalTokens = 0
+    let tokensSetTopLevel = false
+    const tokensPerRunId = {}
+
+    if (!isWorkflow) {
+      const tokens = this.checkTokenUsageChatOrLLMResult(results)
+      inputTokens = tokens.inputTokens
+      outputTokens = tokens.outputTokens
+      totalTokens = tokens.totalTokens
+      tokensSetTopLevel = totalTokens > 0
+    }
+
+    for (const messageSet of results.generations) {
+      for (const chatCompletion of messageSet) {
+        const chatCompletionMessage = chatCompletion.message
+        const role = this.getRole(chatCompletionMessage)
+        const content = chatCompletionMessage.text || ''
+        const toolCalls = this.extractToolCalls(chatCompletionMessage)
+        outputMessages.push({ content, role, toolCalls })
+
+        if (!isWorkflow && !tokensSetTopLevel) {
+          const { tokens, runId } = this.checkTokenUsageFromAIMessage(chatCompletionMessage)
+          if (!tokensPerRunId[runId]) {
+            tokensPerRunId[runId] = tokens
+          } else {
+            tokensPerRunId[runId].inputTokens += tokens.inputTokens
+            tokensPerRunId[runId].outputTokens += tokens.outputTokens
+            tokensPerRunId[runId].totalTokens += tokens.totalTokens
+          }
+        }
+      }
+    }
+
+    if (!isWorkflow && !tokensSetTopLevel) {
+      inputTokens = Object.values(tokensPerRunId).reduce((acc, val) => acc + val.inputTokens, 0)
+      outputTokens = Object.values(tokensPerRunId).reduce((acc, val) => acc + val.outputTokens, 0)
+      totalTokens = Object.values(tokensPerRunId).reduce((acc, val) => acc + val.totalTokens, 0)
+    }
+
+    if (isWorkflow) {
+      this._tagger.tagTextIO(span, inputMessages, outputMessages)
+    } else {
+      this._tagger.tagLLMIO(span, inputMessages, outputMessages)
+      this._tagger.tagMetrics(span, {
+        inputTokens,
+        outputTokens,
+        totalTokens
+      })
+    }
+  }
+
+  extractToolCalls (message) {
+    let toolCalls = message.tool_calls
+    if (!toolCalls) return []
+
+    const toolCallsInfo = []
+    if (!Array.isArray(toolCalls)) toolCalls = [toolCalls]
+    for (const toolCall of toolCalls) {
+      toolCallsInfo.push({
+        name: toolCall.name || '',
+        arguments: toolCall.args || {},
+        tool_id: toolCall.id || ''
+      })
+    }
+
+    return toolCallsInfo
+  }
+}
+
+module.exports = LangChainLLMObsChatModelHandler

package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/embedding.js

@@ -0,0 +1,42 @@
+'use strict'
+
+const LangChainLLMObsHandler = require('.')
+const LLMObsTagger = require('../../../tagger')
+const { spanHasError } = require('../../../util')
+
+class LangChainLLMObsEmbeddingHandler extends LangChainLLMObsHandler {
+  setMetaTags ({ span, inputs, results }) {
+    const isWorkflow = LLMObsTagger.getSpanKind(span) === 'workflow'
+    let embeddingInput, embeddingOutput
+
+    if (isWorkflow) {
+      embeddingInput = this.formatIO(inputs)
+    } else {
+      const input = Array.isArray(inputs) ? inputs : [inputs]
+      embeddingInput = input.map(doc => ({ text: doc }))
+    }
+
+    if (spanHasError(span) || !results) {
+      embeddingOutput = ''
+    } else {
+      let embeddingDimensions, embeddingsCount
+      if (typeof results[0] === 'number') {
+        embeddingsCount = 1
+        embeddingDimensions = results.length
+      } else {
+        embeddingsCount = results.length
+        embeddingDimensions = results[0].length
+      }
+
+      embeddingOutput = `[${embeddingsCount} embedding(s) returned with size ${embeddingDimensions}]`
+    }
+
+    if (isWorkflow) {
+      this._tagger.tagTextIO(span, embeddingInput, embeddingOutput)
+    } else {
+      this._tagger.tagEmbeddingIO(span, embeddingInput, embeddingOutput)
+    }
+  }
+}
+
+module.exports = LangChainLLMObsEmbeddingHandler

package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/index.js

@@ -0,0 +1,102 @@
+'use strict'
+
+const ROLE_MAPPINGS = {
+  human: 'user',
+  ai: 'assistant',
+  system: 'system'
+}
+
+class LangChainLLMObsHandler {
+  constructor (tagger) {
+    this._tagger = tagger
+  }
+
+  setMetaTags () {}
+
+  formatIO (messages) {
+    if (messages.constructor.name === 'Object') { // plain JSON
+      const formatted = {}
+      for (const [key, value] of Object.entries(messages)) {
+        formatted[key] = this.formatIO(value)
+      }
+
+      return formatted
+    } else if (Array.isArray(messages)) {
+      return messages.map(message => this.formatIO(message))
+    } else { // either a BaseMesage type or a string
+      return this.getContentFromMessage(messages)
+    }
+  }
+
+  getContentFromMessage (message) {
+    if (typeof message === 'string') {
+      return message
+    } else {
+      try {
+        const messageContent = {}
+        messageContent.content = message.content || ''
+
+        const role = this.getRole(message)
+        if (role) messageContent.role = role
+
+        return messageContent
+      } catch {
+        return JSON.stringify(message)
+      }
+    }
+  }
+
+  checkTokenUsageChatOrLLMResult (results) {
+    const llmOutput = results.llmOutput
+    const tokens = {
+      inputTokens: 0,
+      outputTokens: 0,
+      totalTokens: 0
+    }
+    if (!llmOutput) return tokens
+    const tokenUsage = llmOutput.tokenUsage || llmOutput.usageMetadata || llmOutput.usage || {}
+    if (!tokenUsage) return tokens
+
+    tokens.inputTokens = tokenUsage.promptTokens || tokenUsage.inputTokens || 0
+    tokens.outputTokens = tokenUsage.completionTokens || tokenUsage.outputTokens || 0
+    tokens.totalTokens = tokenUsage.totalTokens || tokens.inputTokens + tokens.outputTokens
+
+    return tokens
+  }
+
+  checkTokenUsageFromAIMessage (message) {
+    let usage = message.usage_metadata || message.additional_kwargs?.usage
+    const runId = message.run_id || message.id || ''
+    const runIdBase = runId ? runId.split('-').slice(0, -1).join('-') : ''
+
+    const responseMetadata = message.response_metadata || {}
+    usage = usage || responseMetadata.usage || responseMetadata.tokenUsage || {}
+
+    const inputTokens = usage.promptTokens || usage.inputTokens || usage.prompt_tokens || usage.input_tokens || 0
+    const outputTokens =
+      usage.completionTokens || usage.outputTokens || usage.completion_tokens || usage.output_tokens || 0
+    const totalTokens = usage.totalTokens || inputTokens + outputTokens
+
+    return {
+      tokens: {
+        inputTokens,
+        outputTokens,
+        totalTokens
+      },
+      runId: runIdBase
+    }
+  }
+
+  getRole (message) {
+    if (message.role) return ROLE_MAPPINGS[message.role] || message.role
+
+    const type = (
+      (typeof message.getType === 'function' && message.getType()) ||
+      (typeof message._getType === 'function' && message._getType())
+    )
+
+    return ROLE_MAPPINGS[type] || type
+  }
+}
+
+module.exports = LangChainLLMObsHandler

package/packages/dd-trace/src/llmobs/plugins/langchain/handlers/llm.js

@@ -0,0 +1,32 @@
+'use strict'
+
+const LangChainLLMObsHandler = require('.')
+const LLMObsTagger = require('../../../tagger')
+const { spanHasError } = require('../../../util')
+
+class LangChainLLMObsLlmHandler extends LangChainLLMObsHandler {
+  setMetaTags ({ span, inputs, results }) {
+    const isWorkflow = LLMObsTagger.getSpanKind(span) === 'workflow'
+    const prompts = Array.isArray(inputs) ? inputs : [inputs]
+
+    let outputs
+    if (spanHasError(span)) {
+      outputs = [{ content: '' }]
+    } else {
+      outputs = results.generations.map(completion => ({ content: completion[0].text }))
+
+      if (!isWorkflow) {
+        const tokens = this.checkTokenUsageChatOrLLMResult(results)
+        this._tagger.tagMetrics(span, tokens)
+      }
+    }
+
+    if (isWorkflow) {
+      this._tagger.tagTextIO(span, prompts, outputs)
+    } else {
+      this._tagger.tagLLMIO(span, prompts, outputs)
+    }
+  }
+}
+
+module.exports = LangChainLLMObsLlmHandler

package/packages/dd-trace/src/llmobs/plugins/langchain/index.js

@@ -0,0 +1,131 @@
+'use strict'
+
+const log = require('../../../log')
+const LLMObsPlugin = require('../base')
+
+const pluginManager = require('../../../../../..')._pluginManager
+
+const ANTHROPIC_PROVIDER_NAME = 'anthropic'
+const BEDROCK_PROVIDER_NAME = 'amazon_bedrock'
+const OPENAI_PROVIDER_NAME = 'openai'
+
+const SUPPORTED_INTEGRATIONS = ['openai']
+const LLM_SPAN_TYPES = ['llm', 'chat_model', 'embedding']
+const LLM = 'llm'
+const WORKFLOW = 'workflow'
+const EMBEDDING = 'embedding'
+
+const ChainHandler = require('./handlers/chain')
+const ChatModelHandler = require('./handlers/chat_model')
+const LlmHandler = require('./handlers/llm')
+const EmbeddingHandler = require('./handlers/embedding')
+
+class LangChainLLMObsPlugin extends LLMObsPlugin {
+  static get prefix () {
+    return 'tracing:apm:langchain:invoke'
+  }
+
+  constructor () {
+    super(...arguments)
+
+    this._handlers = {
+      chain: new ChainHandler(this._tagger),
+      chat_model: new ChatModelHandler(this._tagger),
+      llm: new LlmHandler(this._tagger),
+      embedding: new EmbeddingHandler(this._tagger)
+    }
+  }
+
+  getLLMObsSpanRegisterOptions (ctx) {
+    const span = ctx.currentStore?.span
+    const tags = span?.context()._tags || {}
+
+    const modelProvider = tags['langchain.request.provider'] // could be undefined
+    const modelName = tags['langchain.request.model'] // could be undefined
+    const kind = this.getKind(ctx.type, modelProvider)
+    const name = tags['resource.name']
+
+    return {
+      modelProvider,
+      modelName,
+      kind,
+      name
+    }
+  }
+
+  setLLMObsTags (ctx) {
+    const span = ctx.currentStore?.span
+    const type = ctx.type // langchain operation type (oneof chain,chat_model,llm,embedding)
+
+    if (!Object.keys(this._handlers).includes(type)) {
+      log.warn(`Unsupported LangChain operation type: ${type}`)
+      return
+    }
+
+    const provider = span?.context()._tags['langchain.request.provider']
+    const integrationName = this.getIntegrationName(type, provider)
+    this.setMetadata(span, provider)
+
+    const inputs = ctx.args?.[0]
+    const options = ctx.args?.[1]
+    const results = ctx.result
+
+    this._handlers[type].setMetaTags({ span, inputs, results, options, integrationName })
+  }
+
+  setMetadata (span, provider) {
+    if (!provider) return
+
+    const metadata = {}
+
+    // these fields won't be set for non model-based operations
+    const temperature =
+      span?.context()._tags[`langchain.request.${provider}.parameters.temperature`] ||
+      span?.context()._tags[`langchain.request.${provider}.parameters.model_kwargs.temperature`]
+
+    const maxTokens =
+      span?.context()._tags[`langchain.request.${provider}.parameters.max_tokens`] ||
+      span?.context()._tags[`langchain.request.${provider}.parameters.maxTokens`] ||
+      span?.context()._tags[`langchain.request.${provider}.parameters.model_kwargs.max_tokens`]
+
+    if (temperature) {
+      metadata.temperature = parseFloat(temperature)
+    }
+
+    if (maxTokens) {
+      metadata.maxTokens = parseInt(maxTokens)
+    }
+
+    this._tagger.tagMetadata(span, metadata)
+  }
+
+  getKind (type, provider) {
+    if (LLM_SPAN_TYPES.includes(type)) {
+      const llmobsIntegration = this.getIntegrationName(type, provider)
+
+      if (!this.isLLMIntegrationEnabled(llmobsIntegration)) {
+        return type === 'embedding' ? EMBEDDING : LLM
+      }
+    }
+
+    return WORKFLOW
+  }
+
+  getIntegrationName (type, provider = 'custom') {
+    if (provider.startsWith(BEDROCK_PROVIDER_NAME)) {
+      return 'bedrock'
+    } else if (provider.startsWith(OPENAI_PROVIDER_NAME)) {
+      return 'openai'
+    } else if (type === 'chat_model' && provider.startsWith(ANTHROPIC_PROVIDER_NAME)) {
+      return 'anthropic'
+    }
+
+    return provider
+  }
+
+  isLLMIntegrationEnabled (integration) {
+    return SUPPORTED_INTEGRATIONS.includes(integration) && pluginManager?._pluginsByName[integration]?.llmobs?._enabled
+  }
+}
+
+module.exports = LangChainLLMObsPlugin

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

@@ -7,7 +7,7 @@ class OpenAiLLMObsPlugin extends LLMObsPlugin {
     return 'tracing:apm:openai:request'
   }
 
-  getLLMObsSPanRegisterOptions (ctx) {
+  getLLMObsSpanRegisterOptions (ctx) {
     const resource = ctx.methodName
     const methodName = gateResource(normalizeOpenAIResourceName(resource))
     if (!methodName) return // we will not trace all openai methods for llmobs

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

@@ -1,12 +1,12 @@
 'use strict'
 
-const { SPAN_KIND, OUTPUT_VALUE } = require('./constants/tags')
+const { SPAN_KIND, OUTPUT_VALUE, INPUT_VALUE } = require('./constants/tags')
 
 const {
   getFunctionArguments,
   validateKind
 } = require('./util')
-const { isTrue } = require('../util')
+const { isTrue, isError } = require('../util')
 
 const { storage } = require('./storage')
 
@@ -134,29 +134,63 @@ class LLMObs extends NoopLLMObs {
 
     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) })
+      const fnArgs = arguments
+
+      const lastArgId = fnArgs.length - 1
+      const cb = fnArgs[lastArgId]
+      const hasCallback = typeof cb === 'function'
+
+      if (hasCallback) {
+        const scopeBoundCb = llmobs._bind(cb)
+        fnArgs[lastArgId] = function () {
+          // 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 = arguments[0]
+          const maybeResult = arguments[1]
+
+          llmobs._autoAnnotate(
+            span,
+            kind,
+            getFunctionArguments(fn, fnArgs),
+            isError(maybeError) || maybeError == null ? maybeResult : maybeError
+          )
+
+          return scopeBoundCb.apply(this, arguments)
         }
+      }
 
-        return fn.apply(this, arguments)
-      })
+      try {
+        const result = llmobs._activate(span, { kind, options: 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
+            }
+          )
+        }
 
-      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
-        })
-      }
+        // 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)
+        }
 
-      if (result && !['llm', 'retrieval'].includes(kind) && !LLMObsTagger.tagMap.get(span)?.[OUTPUT_VALUE]) {
-        llmobs.annotate(span, { outputData: result })
+        return result
+      } catch (e) {
+        llmobs._autoAnnotate(span, kind, getFunctionArguments(fn, fnArgs))
+        throw e
       }
-
-      return result
     }
 
     return this._tracer.wrap(name, spanOptions, wrapped)
@@ -333,20 +367,34 @@ class LLMObs extends NoopLLMObs {
     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)
+  }
+
   _active () {
     const store = storage.getStore()
     return store?.span
   }
 
-  _activate (span, { kind, options } = {}, fn) {
+  _activate (span, options, fn) {
     const parent = this._active()
     if (this.enabled) storage.enterWith({ span })
 
-    this._tagger.registerLLMObsSpan(span, {
-      ...options,
-      parent,
-      kind
-    })
+    if (options) {
+      this._tagger.registerLLMObsSpan(span, {
+        ...options,
+        parent
+      })
+    }
 
     try {
       return fn()
@@ -355,6 +403,22 @@ class LLMObs extends NoopLLMObs {
     }
   }
 
+  // bind function to active LLMObs span
+  _bind (fn) {
+    if (typeof fn !== 'function') return fn
+
+    const llmobs = this
+    const activeSpan = llmobs._active()
+
+    const bound = function () {
+      return llmobs._activate(activeSpan, null, () => {
+        return fn.apply(this, arguments)
+      })
+    }
+
+    return bound
+  }
+
   _extractOptions (options) {
     const {
       modelName,

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

@@ -40,6 +40,10 @@ class LLMObsTagger {
     return registry
   }
 
+  static getSpanKind (span) {
+    return registry.get(span)?.[SPAN_KIND]
+  }
+
   registerLLMObsSpan (span, {
     modelName,
     modelProvider,
@@ -60,10 +64,10 @@ class LLMObsTagger {
     if (modelName) this._setTag(span, MODEL_NAME, modelName)
     if (modelProvider) this._setTag(span, MODEL_PROVIDER, modelProvider)
 
-    sessionId = sessionId || parent?.context()._tags[SESSION_ID]
+    sessionId = sessionId || registry.get(parent)?.[SESSION_ID]
     if (sessionId) this._setTag(span, SESSION_ID, sessionId)
 
-    if (!mlApp) mlApp = parent?.context()._tags[ML_APP] || this._config.llmobs.mlApp
+    if (!mlApp) mlApp = registry.get(parent)?.[ML_APP] || this._config.llmobs.mlApp
     this._setTag(span, ML_APP, mlApp)
 
     const parentId =
@@ -136,6 +140,10 @@ class LLMObsTagger {
     this._setTag(span, TAGS, tags)
   }
 
+  changeKind (span, newKind) {
+    this._setTag(span, SPAN_KIND, newKind)
+  }
+
   _tagText (span, data, key) {
     if (data) {
       if (typeof data === 'string') {
@@ -310,7 +318,7 @@ class LLMObsTagger {
   _setTag (span, key, value) {
     if (!this._config.llmobs.enabled) return
     if (!registry.has(span)) {
-      this._handleFailure('Span must be an LLMObs generated span.')
+      this._handleFailure(`Span "${span._name}" must be an LLMObs generated span.`)
       return
     }
 

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

@@ -169,8 +169,14 @@ function getFunctionArguments (fn, args = []) {
   }
 }
 
+function spanHasError (span) {
+  const tags = span.context()._tags
+  return !!(tags.error || tags['error.type'])
+}
+
 module.exports = {
   encodeUnicode,
   validateKind,
-  getFunctionArguments
+  getFunctionArguments,
+  spanHasError
 }

package/packages/dd-trace/src/llmobs/writers/spans/agentProxy.js

@@ -10,10 +10,10 @@ const LLMObsBaseSpanWriter = require('./base')
 class LLMObsAgentProxySpanWriter extends LLMObsBaseSpanWriter {
   constructor (config) {
     super({
-      intake: config.hostname || 'localhost',
-      protocol: 'http:',
+      intake: config.url?.hostname || config.hostname || 'localhost',
+      protocol: config.url?.protocol || 'http:',
       endpoint: EVP_PROXY_AGENT_ENDPOINT,
-      port: config.port
+      port: config.url?.port || config.port
     })
 
     this._headers[EVP_SUBDOMAIN_HEADER_NAME] = EVP_SUBDOMAIN_HEADER_VALUE

package/packages/dd-trace/src/log/index.js

@@ -63,15 +63,14 @@ const log = {
 
       Error.captureStackTrace(logRecord, this.trace)
 
-      const fn = logRecord.stack.split('\n')[1].replace(/^\s+at ([^\s]+) .+/, '$1')
-      const params = args.map(a => {
-        return a && a.hasOwnProperty('toString') && typeof a.toString === 'function'
-          ? a.toString()
-          : inspect(a, { depth: 3, breakLength: Infinity, compact: true })
-      }).join(', ')
-      const formatted = logRecord.stack.replace('Error: ', `Trace: ${fn}(${params})`)
-
-      traceChannel.publish(Log.parse(formatted))
+      const stack = logRecord.stack.split('\n')
+      const fn = stack[1].replace(/^\s+at ([^\s]+) .+/, '$1')
+      const options = { depth: 2, breakLength: Infinity, compact: true, maxArrayLength: Infinity }
+      const params = args.map(a => inspect(a, options)).join(', ')
+
+      stack[0] = `Trace: ${fn}(${params})`
+
+      traceChannel.publish(Log.parse(stack.join('\n')))
     }
     return this
   },