dd-trace 5.31.0 → 5.33.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/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/opentelemetry/context_manager.js

@@ -1,7 +1,7 @@
 'use strict'
 
 const { storage } = require('../../../datadog-core')
-const { trace, ROOT_CONTEXT } = require('@opentelemetry/api')
+const { trace, ROOT_CONTEXT, propagation } = require('@opentelemetry/api')
 const DataDogSpanContext = require('../opentracing/span_context')
 
 const SpanContext = require('./span_context')
@@ -18,17 +18,40 @@ class ContextManager {
     const context = (activeSpan && activeSpan.context()) || store || ROOT_CONTEXT
 
     if (!(context instanceof DataDogSpanContext)) {
+      const span = trace.getSpan(context)
+      // span instanceof NonRecordingSpan
+      if (span && span._spanContext && span._spanContext._ddContext && span._spanContext._ddContext._baggageItems) {
+        const baggages = span._spanContext._ddContext._baggageItems
+        const entries = {}
+        for (const [key, value] of Object.entries(baggages)) {
+          entries[key] = { value }
+        }
+        const otelBaggages = propagation.createBaggage(entries)
+        return propagation.setBaggage(context, otelBaggages)
+      }
       return context
     }
 
+    const baggages = JSON.parse(activeSpan.getAllBaggageItems())
+    const entries = {}
+    for (const [key, value] of Object.entries(baggages)) {
+      entries[key] = { value }
+    }
+    const otelBaggages = propagation.createBaggage(entries)
+
     if (!context._otelSpanContext) {
       const newSpanContext = new SpanContext(context)
       context._otelSpanContext = newSpanContext
     }
     if (store && trace.getSpanContext(store) === context._otelSpanContext) {
-      return store
+      return otelBaggages
+        ? propagation.setBaggage(store, otelBaggages)
+        : store
     }
-    return trace.setSpanContext(store || ROOT_CONTEXT, context._otelSpanContext)
+    const wrappedContext = trace.setSpanContext(store || ROOT_CONTEXT, context._otelSpanContext)
+    return otelBaggages
+      ? propagation.setBaggage(wrappedContext, otelBaggages)
+      : wrappedContext
   }
 
   with (context, fn, thisArg, ...args) {
@@ -38,9 +61,26 @@ class ContextManager {
       const cb = thisArg == null ? fn : fn.bind(thisArg)
       return this._store.run(context, cb, ...args)
     }
+    const baggages = propagation.getBaggage(context)
+    let baggageItems = []
+    if (baggages) {
+      baggageItems = baggages.getAllEntries()
+    }
     if (span && span._ddSpan) {
+      // does otel always override datadog?
+      span._ddSpan.removeAllBaggageItems()
+      for (const baggage of baggageItems) {
+        span._ddSpan.setBaggageItem(baggage[0], baggage[1].value)
+      }
       return ddScope.activate(span._ddSpan, run)
     }
+    // span instanceof NonRecordingSpan
+    if (span && span._spanContext && span._spanContext._ddContext && span._spanContext._ddContext._baggageItems) {
+      span._spanContext._ddContext._baggageItems = {}
+      for (const baggage of baggageItems) {
+        span._spanContext._ddContext._baggageItems[baggage[0]] = baggage[1].value
+      }
+    }
     return run()
   }