diff --git a/lib/llm-events/openai/chat-completion-message.js b/lib/llm-events/openai/chat-completion-message.js index ebfbe57a7e..61f101a464 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 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/unit/llm-events/openai/chat-completion-message.test.js b/test/unit/llm-events/openai/chat-completion-message.test.js index a42a725287..0599bb5f2f 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 e5c9a722a2..07921fa50a 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 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',