chore: Changed token_count to only use tokenCountCallback (#2070)

newrelic · Mar 11, 2024 · 66f94b0 · 66f94b0
1 parent 64b4ca2
commit 66f94b0
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 15 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/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',