From bc6a5a8869cdded238e234f02a069df8d2ffa7d1 Mon Sep 17 00:00:00 2001 From: shaohuzhang1 <80892890+shaohuzhang1@users.noreply.github.com> Date: Tue, 22 Jul 2025 18:43:10 +0800 Subject: [PATCH] fix: Non streaming sessions cannot count token consumption #3635 (#3709) --- apps/setting/models_provider/impl/base_chat_open_ai.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/apps/setting/models_provider/impl/base_chat_open_ai.py b/apps/setting/models_provider/impl/base_chat_open_ai.py index ccafc0e0b..626a751f7 100644 --- a/apps/setting/models_provider/impl/base_chat_open_ai.py +++ b/apps/setting/models_provider/impl/base_chat_open_ai.py @@ -99,7 +99,7 @@ class BaseChatOpenAI(ChatOpenAI): except Exception as e: tokenizer = TokenizerManage.get_tokenizer() return sum([len(tokenizer.encode(get_buffer_string([m]))) for m in messages]) - return self.usage_metadata.get('input_tokens', 0) + return self.usage_metadata.get('input_tokens', self.usage_metadata.get('prompt_tokens', 0)) def get_num_tokens(self, text: str) -> int: if self.usage_metadata is None or self.usage_metadata == {}: @@ -108,7 +108,8 @@ class BaseChatOpenAI(ChatOpenAI): except Exception as e: tokenizer = TokenizerManage.get_tokenizer() return len(tokenizer.encode(text)) - return self.get_last_generation_info().get('output_tokens', 0) + return self.get_last_generation_info().get('output_tokens', + self.get_last_generation_info().get('completion_tokens', 0)) def _stream(self, *args: Any, **kwargs: Any) -> Iterator[ChatGenerationChunk]: kwargs['stream_usage'] = True