From f7147ab312b5e78ae91d7946486e68371c86d132 Mon Sep 17 00:00:00 2001 From: KernelDeimos Date: Thu, 6 Feb 2025 14:32:21 -0500 Subject: [PATCH] dev: generalize how OpenAIUtil handles usage The current implementation would break for XAIService because it gets cumulative usage reports as the stream goes on (whereas OpenAI and DeepSeek only provide this at the end of the stream). --- src/backend/src/modules/puterai/lib/OpenAIUtil.js | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/backend/src/modules/puterai/lib/OpenAIUtil.js b/src/backend/src/modules/puterai/lib/OpenAIUtil.js index 965ac09d..b124822a 100644 --- a/src/backend/src/modules/puterai/lib/OpenAIUtil.js +++ b/src/backend/src/modules/puterai/lib/OpenAIUtil.js @@ -66,14 +66,12 @@ module.exports = class OpenAIUtil { let mode = 'text'; const tool_call_blocks = []; + let last_usage = null; for await ( const chunk of completion ) { - if ( chunk.usage ) { - usage_promise.resolve({ - input_tokens: chunk.usage.prompt_tokens, - output_tokens: chunk.usage.completion_tokens, - }); - continue; + if ( process.env.DEBUG ) { + console.log(`AI CHUNK`, chunk); } + if ( chunk.usage ) last_usage = chunk.usage; if ( chunk.choices.length < 1 ) continue; const choice = chunk.choices[0]; @@ -108,6 +106,10 @@ module.exports = class OpenAIUtil { } } } + usage_promise.resolve({ + input_tokens: last_usage.prompt_tokens, + output_tokens: last_usage.completion_tokens, + }); if ( mode === 'text' ) textblock.end(); if ( mode === 'tool' ) toolblock.end();