From 0a95e60c8ebd43ae64391504f74df801d6afeb47 Mon Sep 17 00:00:00 2001
From: James Sumners <jsumners@newrelic.com>
Date: Fri, 8 Mar 2024 09:15:45 -0500
Subject: [PATCH] chore: Changed token_count to only use tokenCountCallback

---
 .../openai/chat-completion-message.js         | 13 +++++++------
 lib/llm-events/openai/embedding.js            |  7 ++++---
 test/versioned/openai/chat-completions.tap.js |  3 +--
 test/versioned/openai/common.js               | 19 +------------------
 test/versioned/openai/embeddings.tap.js       |  2 +-
 5 files changed, 14 insertions(+), 30 deletions(-)

diff --git a/lib/llm-events/openai/chat-completion-message.js b/lib/llm-events/openai/chat-completion-message.js
index ebfbe57a7e..c9bc3edd14 100644
--- a/lib/llm-events/openai/chat-completion-message.js
+++ b/lib/llm-events/openai/chat-completion-message.js
@@ -19,14 +19,15 @@ module.exports = class LlmChatCompletionMessage extends LlmEvent {
       this.content = message?.content
     }
 
+    const tokenCB = agent.llm?.tokenCountCallback
+    if (tokenCB == null) {
+      return
+    }
+
     if (this.is_response) {
-      this.token_count =
-        response?.usage?.completion_tokens ||
-        agent.llm?.tokenCountCallback?.(this['response.model'], message?.content)
+      this.token_count = tokenCB(this['response.model'], message?.content)
     } else {
-      this.token_count =
-        response?.usage?.prompt_tokens ||
-        agent.llm?.tokenCountCallback?.(request.model || request.engine, message?.content)
+      this.token_count = tokenCB(request.model || request.engine, message?.content)
     }
   }
 }
diff --git a/lib/llm-events/openai/embedding.js b/lib/llm-events/openai/embedding.js
index 74a60e855e..1f2ee78bfb 100644
--- a/lib/llm-events/openai/embedding.js
+++ b/lib/llm-events/openai/embedding.js
@@ -14,8 +14,9 @@ module.exports = class LlmEmbedding extends LlmEvent {
     if (agent.config.ai_monitoring.record_content.enabled === true) {
       this.input = request.input?.toString()
     }
-    this.token_count =
-      response?.usage?.prompt_tokens ||
-      agent.llm?.tokenCountCallback?.(this['request.model'], request.input?.toString())
+    this.token_count = agent.llm?.tokenCountCallback?.(
+      this['request.model'],
+      request.input?.toString()
+    )
   }
 }
diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js
index efdbc63790..67334d7fd9 100644
--- a/test/versioned/openai/chat-completions.tap.js
+++ b/test/versioned/openai/chat-completions.tap.js
@@ -96,8 +96,7 @@ tap.test('OpenAI instrumentation - chat completions', (t) => {
         model,
         id: 'chatcmpl-87sb95K4EF2nuJRcTs43Tm9ntTeat',
         resContent: '1 plus 2 is 3.',
-        reqContent: content,
-        tokenUsage: true
+        reqContent: content
       })
 
       const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0]
diff --git a/test/versioned/openai/common.js b/test/versioned/openai/common.js
index 1f54535c92..9be90bc573 100644
--- a/test/versioned/openai/common.js
+++ b/test/versioned/openai/common.js
@@ -39,15 +39,7 @@ common.afterHook = function afterHook(t) {
   t.context.agent && helper.unloadAgent(t.context.agent)
 }
 
-function assertChatCompletionMessages({
-  tx,
-  chatMsgs,
-  id,
-  model,
-  reqContent,
-  resContent,
-  tokenUsage
-}) {
+function assertChatCompletionMessages({ tx, chatMsgs, id, model, reqContent, resContent }) {
   const baseMsg = {
     'appName': 'New Relic for Node.js tests',
     'request_id': '49dbbffbd3c3f4612aa48def69059aad',
@@ -68,25 +60,16 @@ function assertChatCompletionMessages({
       expectedChatMsg.sequence = 0
       expectedChatMsg.id = `${id}-0`
       expectedChatMsg.content = reqContent
-      if (tokenUsage) {
-        expectedChatMsg.token_count = 53
-      }
     } else if (msg[1].sequence === 1) {
       expectedChatMsg.sequence = 1
       expectedChatMsg.id = `${id}-1`
       expectedChatMsg.content = 'What does 1 plus 1 equal?'
-      if (tokenUsage) {
-        expectedChatMsg.token_count = 53
-      }
     } else {
       expectedChatMsg.sequence = 2
       expectedChatMsg.role = 'assistant'
       expectedChatMsg.id = `${id}-2`
       expectedChatMsg.content = resContent
       expectedChatMsg.is_response = true
-      if (tokenUsage) {
-        expectedChatMsg.token_count = 11
-      }
     }
 
     this.equal(msg[0].type, 'LlmChatCompletionMessage')
diff --git a/test/versioned/openai/embeddings.tap.js b/test/versioned/openai/embeddings.tap.js
index 4d4d68e08d..6ac1cd96a8 100644
--- a/test/versioned/openai/embeddings.tap.js
+++ b/test/versioned/openai/embeddings.tap.js
@@ -100,7 +100,7 @@ tap.test('OpenAI instrumentation - embedding', (t) => {
         'response.organization': 'new-relic-nkmd8b',
         'response.usage.total_tokens': 6,
         'response.usage.prompt_tokens': 6,
-        'token_count': 6,
+        'token_count': undefined,
         'response.headers.llmVersion': '2020-10-01',
         'response.headers.ratelimitLimitRequests': '200',
         'response.headers.ratelimitLimitTokens': '150000',