From 66f94b05aa7018af6fa1db1d2d87de9013f23a5b Mon Sep 17 00:00:00 2001
From: James Sumners <jsumners@newrelic.com>
Date: Mon, 11 Mar 2024 10:58:53 -0400
Subject: [PATCH] chore: Changed token_count to only use tokenCountCallback
 (#2070)

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

diff --git a/lib/llm-events/openai/chat-completion-message.js b/lib/llm-events/openai/chat-completion-message.js
index ebfbe57a7..61f101a46 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 (typeof tokenCB !== 'function') {
+      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 74a60e855..1f2ee78bf 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/unit/llm-events/openai/chat-completion-message.test.js b/test/unit/llm-events/openai/chat-completion-message.test.js
index a42a72528..0599bb5f2 100644
--- a/test/unit/llm-events/openai/chat-completion-message.test.js
+++ b/test/unit/llm-events/openai/chat-completion-message.test.js
@@ -62,7 +62,6 @@ tap.test('LlmChatCompletionMessage', (t) => {
         expected.content = chatRes.choices[0].message.content
         expected.role = chatRes.choices[0].message.role
         expected.is_response = true
-        expected.token_count = 20
         t.same(chatMessageEvent, expected)
         t.end()
       })
diff --git a/test/unit/llm-events/openai/common.js b/test/unit/llm-events/openai/common.js
index e5c9a722a..07921fa50 100644
--- a/test/unit/llm-events/openai/common.js
+++ b/test/unit/llm-events/openai/common.js
@@ -75,7 +75,7 @@ function getExpectedResult(tx, event, type, completionId) {
       expected = { ...expected, ...resKeys }
       expected.input = 'This is my test input'
       expected.error = false
-      expected.token_count = 10
+      expected.token_count = undefined
       break
     case 'summary':
       expected = {
@@ -96,7 +96,6 @@ function getExpectedResult(tx, event, type, completionId) {
         role: 'inquisitive-kid',
         sequence: 0,
         completion_id: completionId,
-        token_count: 10,
         is_response: false
       }
   }
diff --git a/test/versioned/openai/chat-completions.tap.js b/test/versioned/openai/chat-completions.tap.js
index efdbc6379..67334d7fd 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/embeddings.tap.js b/test/versioned/openai/embeddings.tap.js
index 4d4d68e08..6ac1cd96a 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',