chore: Changed token_count to only use tokenCountCallback

jsumners-nr · Mar 8, 2024 · 3f8076b · 3f8076b
1 parent 64b4ca2
commit 3f8076b
Show file tree

Hide file tree

Showing 7 changed files with 15 additions and 33 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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',