address feedback

jsumners-nr · Mar 11, 2024 · 19e0a2b · 19e0a2b
1 parent 3f8076b
commit 19e0a2b
Showing 1 changed file with 18 additions and 1 deletion.
diff --git a/test/versioned/openai/common.js b/test/versioned/openai/common.js
@@ -39,7 +39,15 @@ common.afterHook = function afterHook(t) {
   t.context.agent && helper.unloadAgent(t.context.agent)
 }
 
-function assertChatCompletionMessages({ tx, chatMsgs, id, model, reqContent, resContent }) {
+function assertChatCompletionMessages({
+  tx,
+  chatMsgs,
+  id,
+  model,
+  reqContent,
+  resContent,
+  tokenUsage
+}) {
   const baseMsg = {
     'appName': 'New Relic for Node.js tests',
     'request_id': '49dbbffbd3c3f4612aa48def69059aad',
@@ -60,16 +68,25 @@ function assertChatCompletionMessages({ tx, chatMsgs, id, model, reqContent, res
       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')