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

Datadog APM tracing client for JavaScript

'use strict'

class LangChainLLMObsHandler {
  constructor (tagger) {
    /** @type {import('../../../tagger')} */
    this._tagger = tagger
  }

  getName ({ span }) {
    return span?.context()._tags?.['resource.name']
  }

  setMetaTags () {}

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

module.exports = LangChainLLMObsHandler