feat: Updated token calculation for openai LLM Events (newrelic#3430)

Author: svetlanabrennan

lib/llm-events/openai/chat-completion-message.js

@@ -5,6 +5,8 @@
 
 'use strict'
 const LlmEvent = require('./event')
+const { tokenUsageAttributesExist } = require('./utils')
+const { setTokenFromCallback } = require('../utils')
 
 module.exports = class LlmChatCompletionMessage extends LlmEvent {
   constructor({
@@ -38,17 +40,42 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
     if (agent.config.ai_monitoring.record_content.enabled === true) {
       this.content = content
     }
+    this.setTokenCount(agent, request, response)
+  }
 
-    // Calculate token count if the callback is available.
+  setTokenCount(agent, request, response) {
     const tokenCB = agent.llm?.tokenCountCallback
-    if (typeof tokenCB !== 'function') {
+
+    if (tokenCB) {
+      const messages = request?.input || request?.messages
+
+      const promptContent = typeof messages === 'string'
+        ? messages
+        : messages?.map((msg) => msg.content).join(' ')
+
+      const completionContent = response?.output
+        ? response.output.map((resContent) => resContent.content[0].text).join(' ')
+        : response?.choices?.map((resContent) => resContent.message.content).join(' ')
+
+      if (promptContent && completionContent) {
+        setTokenFromCallback(
+          {
+            context: this,
+            tokenCB,
+            reqModel: request.model,
+            resModel: this['response.model'],
+            promptContent,
+            completionContent
+          }
+        )
+      }
       return
     }
 
-    if (this.is_response) {
-      this.token_count = tokenCB(this['response.model'], content)
-    } else {
-      this.token_count = tokenCB(request.model || request.engine, content)
+    // If no token count callback is available, we need to check the response object
+    // for usage information and set token_count to 0 if all usage attributes are present.
+    if (tokenUsageAttributesExist(response)) {
+      this.token_count = 0
     }
   }
 }

lib/llm-events/openai/chat-completion-summary.js

@@ -5,6 +5,8 @@
 
 'use strict'
 const LlmEvent = require('./event')
+const { setUsageTokens } = require('./utils')
+const { setTokenUsageFromCallback } = require('../utils')
 
 module.exports = class LlmChatCompletionSummary extends LlmEvent {
   constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -24,5 +26,39 @@ module.exports = class LlmChatCompletionSummary extends LlmEvent {
       this['response.number_of_messages'] = request?.messages?.length + response?.choices?.length
       this['response.choices.finish_reason'] = response?.choices?.[0]?.finish_reason
     }
+
+    this.setTokens(agent, request, response)
+  }
+
+  setTokens(agent, request, response) {
+    const tokenCB = agent.llm?.tokenCountCallback
+
+    // Prefer callback for prompt and completion tokens; if unavailable, fall back to response data.
+    if (tokenCB) {
+      const messages = request?.input || request?.messages
+
+      const promptContent = typeof messages === 'string'
+        ? messages
+        : messages?.map((msg) => msg.content).join(' ')
+
+      const completionContent = response?.output
+        ? response.output.map((resContent) => resContent.content[0].text).join(' ')
+        : response?.choices?.map((resContent) => resContent.message.content).join(' ')
+
+      if (promptContent && completionContent) {
+        setTokenUsageFromCallback(
+          {
+            context: this,
+            tokenCB,
+            reqModel: request.model,
+            resModel: this['response.model'],
+            promptContent,
+            completionContent
+          }
+        )
+      }
+      return
+    }
+    setUsageTokens(response, this)
   }
 }

lib/llm-events/openai/embedding.js

@@ -5,6 +5,7 @@
 
 'use strict'
 const LlmEvent = require('./event')
+const { validCallbackTokenCount, calculateCallbackTokens } = require('../utils')
 
 module.exports = class LlmEmbedding extends LlmEvent {
   constructor({ agent, segment, request = {}, response = {}, withError = false, transaction }) {
@@ -14,9 +15,38 @@ module.exports = class LlmEmbedding extends LlmEvent {
     if (agent.config.ai_monitoring.record_content.enabled === true) {
       this.input = request.input?.toString()
     }
-    this.token_count = agent.llm?.tokenCountCallback?.(
-      this['request.model'],
-      request.input?.toString()
-    )
+
+    this.setTotalTokens(agent, request, response)
+  }
+
+  setTotalTokens(agent, request, response) {
+    const tokenCB = agent.llm?.tokenCountCallback
+
+    // For embedding events, only total token count is relevant.
+    // Prefer callback for total tokens; if unavailable, fall back to response data.
+    if (tokenCB) {
+      const content = request.input?.toString()
+
+      if (content === undefined) {
+        return
+      }
+
+      const totalTokenCount = calculateCallbackTokens(tokenCB, this['request.model'], content)
+      const hasValidCallbackCounts = validCallbackTokenCount(totalTokenCount)
+
+      if (hasValidCallbackCounts) {
+        this['response.usage.total_tokens'] = Number(totalTokenCount)
+      }
+      return
+    }
+
+    const totalTokens = this.getTotalTokens(response)
+    if (totalTokens) {
+      this['response.usage.total_tokens'] = Number(totalTokens)
+    }
+  }
+
+  getTotalTokens(response) {
+    return response?.usage?.total_tokens || response?.usage?.totalTokens
   }
 }

lib/llm-events/openai/utils.js

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2024 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+const { setTokensInResponse } = require('../utils')
+
+function setUsageTokens(response, context) {
+  // input and output token counts must available in order to add all usage attributes to response
+  // if total tokens is not available, we can manually add it up (from input and output token count)
+  if (tokenUsageAttributesExist(response) === false) {
+    return
+  }
+
+  const promptTokens = Number(response?.usage?.prompt_tokens || response?.usage?.input_tokens)
+  const completionTokens = Number(response?.usage?.completion_tokens || response?.usage?.output_tokens)
+  const totalTokens = Number(response?.usage?.total_tokens || response?.usage?.totalTokens)
+
+  setTokensInResponse(context, { promptTokens, completionTokens, totalTokens })
+}
+
+function tokenUsageAttributesExist(response) {
+  const tokensA = response?.usage?.prompt_tokens && response?.usage?.completion_tokens
+  const tokensB = response?.usage?.input_tokens && response?.usage?.output_tokens
+
+  return tokensA !== undefined || tokensB !== undefined
+}
+
+module.exports = {
+  tokenUsageAttributesExist,
+  setUsageTokens
+}

lib/llm-events/utils.js

@@ -0,0 +1,110 @@
+/*
+ * Copyright 2025 New Relic Corporation. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+'use strict'
+
+/**
+ * Determines if the provided token count is valid.
+ * A valid token count is greater than 0 and not null.
+ * @param {number} tokenCount The token count obtained from the token callback
+ * @returns {boolean} Whether the token count is valid
+ */
+function validCallbackTokenCount(tokenCount) {
+  return tokenCount !== null && tokenCount > 0
+}
+
+/**
+ * Calculates the total token count from the prompt tokens and completion tokens
+ * set in the context.
+ * @param {LlmEvent} context The context object containing token counts
+ * @returns {number} The total token count
+ */
+function getTotalTokenCount(context) {
+  return Number(context['response.usage.prompt_tokens']) + Number(context['response.usage.completion_tokens'])
+}
+
+/**
+ * Sets the provided tokens counts on the LLM event.
+ * @param {LlmChatCompletionMessage} context The context object to set token usage counts on.
+ * @param {object} tokens The object contains the token prompt, completion and total counts.
+ */
+function setTokensInResponse(context, tokens) {
+  context['response.usage.prompt_tokens'] = tokens.promptTokens
+  context['response.usage.completion_tokens'] = tokens.completionTokens
+  context['response.usage.total_tokens'] = tokens.totalTokens || getTotalTokenCount(context)
+}
+
+/**
+ * Calculates prompt and completion token counts using the provided callback and models.
+ * If both counts are valid, sets context.token_count to 0.
+ *
+ * @param {object} options - The params object.
+ * @param {LlmChatCompletionMessage} options.context - The context object to set token count on.
+ * @param {Function} options.tokenCB - The token counting callback function.
+ * @param {string} options.reqModel - The model used for the prompt.
+ * @param {string} options.resModel - The model used for the completion.
+ * @param {string} options.promptContent - The prompt content to count tokens for.
+ * @param {string} options.completionContent - The completion content to count tokens for.
+ * @returns {void}
+ */
+function setTokenFromCallback({ context, tokenCB, reqModel, resModel, promptContent, completionContent }) {
+  const promptToken = calculateCallbackTokens(tokenCB, reqModel, promptContent)
+  const completionToken = calculateCallbackTokens(tokenCB, resModel, completionContent)
+
+  const hasValidCallbackCounts =
+    validCallbackTokenCount(promptToken) && validCallbackTokenCount(completionToken)
+
+  if (hasValidCallbackCounts) {
+    context.token_count = 0
+  }
+}
+
+/**
+ * Calculates prompt and completion token counts using the provided callback and models.
+ * If both counts are valid, sets token prompt, completion and total counts on the context.
+ *
+ * @param {object} options - The params object.
+ * @param {LlmEvent} options.context - The context object (llm summary or llm embedding) to set token count on.
+ * @param {Function} options.tokenCB - The token counting callback function.
+ * @param {string} options.reqModel - The model used for the prompt.
+ * @param {string} options.resModel - The model used for the completion.
+ * @param {string} options.promptContent - The prompt content to count tokens for.
+ * @param {string} options.completionContent - The completion content to count tokens for.
+ * @returns {void}
+ */
+function setTokenUsageFromCallback({ context, tokenCB, reqModel, resModel, promptContent, completionContent }) {
+  const promptTokens = calculateCallbackTokens(tokenCB, reqModel, promptContent)
+  const completionTokens = calculateCallbackTokens(tokenCB, resModel, completionContent)
+
+  const hasValidCallbackCounts =
+    validCallbackTokenCount(promptTokens) && validCallbackTokenCount(completionTokens)
+
+  if (hasValidCallbackCounts) {
+    setTokensInResponse(context, { promptTokens, completionTokens, totalTokens: promptTokens + completionTokens })
+  }
+}
+
+/**
+ * Calculate the token counts using the provided callback.
+ * @param {Function} tokenCB - The token count callback function.
+ * @param {string} model - The model.
+ * @param {string} content - The content to calculate tokens for, such as prompt or completion response.
+ * @returns {number|undefined} - The calculated token count or undefined if callback is not a function.
+ */
+function calculateCallbackTokens(tokenCB, model, content) {
+  if (typeof tokenCB === 'function') {
+    return tokenCB(model, content)
+  }
+  return undefined
+}
+
+module.exports = {
+  validCallbackTokenCount,
+  getTotalTokenCount,
+  setTokensInResponse,
+  setTokenFromCallback,
+  setTokenUsageFromCallback,
+  calculateCallbackTokens
+}