diff --git a/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py b/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py index 25068d9ccba..c5946a77eca 100644 --- a/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py +++ b/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py @@ -17,7 +17,14 @@ import uuid_utils.compat as uuid from django.db.models import QuerySet from django.utils.translation import gettext_lazy as _ -from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage, AIMessageChunk +from langchain_core.messages import ( + AIMessage, + HumanMessage, + SystemMessage, + ToolMessage, + AIMessageChunk, + get_buffer_string, +) from rest_framework import serializers from application.workflow.common import AsyncRunner, WorkflowType @@ -35,7 +42,9 @@ ) from application.workflow.status import Status from application.workflow.tools import Reasoning +from common.config.tokenizer_manage_config import TokenizerManage from common.utils.common import guess_image_format +from common.utils.count_tokens import count_local_tokens from common.utils.messages_util import to_ai_message_list, to_human_message_list from common.utils.shared_resource_auth import filter_authorized_ids from common.utils.tool_code import ToolExecutor @@ -388,6 +397,8 @@ def _generate_prompt_question(self, question_str, model, vision, image_list, vid async def _stream_response( self, response, chat_model, message_list, question, reasoning_content_id, text_content_id ): + answer = "" + reasoning_content = "" try: node_info = NodeInfo(self.get_node_id(), self.get_node_name(), Status.RUNNING) model_setting = self.get_context("model_setting") or {} @@ -395,8 +406,6 @@ async def _stream_response( model_setting.get("reasoning_content_start", ""), model_setting.get("reasoning_content_end", ""), ) - answer = "" - reasoning_content = "" response_reasoning_content = False async for chunk in response: @@ -462,9 +471,9 @@ async def _stream_response( ) ) - self._write_final_context(chat_model, message_list, question, answer, reasoning_content) - # 无论正常结束、抛异常还是被取消中断,都关闭 LLM 流式响应底层连接 + # 无论正常结束、抛异常还是被取消中断,都落盘已累积内容并关闭 LLM 流式响应底层连接 finally: + self._write_final_context(chat_model, message_list, question, answer, reasoning_content) with contextlib.suppress(Exception): await response.aclose() @@ -490,8 +499,12 @@ def _invoke_response(self, response, chat_model, message_list, question, is_resu self.write(TextContent(text_content_id, content, Status.SUCCESS, node_info, Position(self.get_node_id()))) def _write_final_context(self, chat_model, message_list, question, answer, reasoning_content): - message_tokens = chat_model.get_num_tokens_from_messages(message_list) - answer_tokens = chat_model.get_num_tokens(answer) + # 中断时模型已不可用,只能走本地 tokenizer 计算 + message_tokens, answer_tokens = ( + count_local_tokens(message_list, answer) + if self.async_runner.cancel_event.is_set() + else (chat_model.get_num_tokens_from_messages(message_list), chat_model.get_num_tokens(answer)) + ) self.data["message_tokens"] = message_tokens self.data["answer_tokens"] = answer_tokens self.write_context("answer", answer) diff --git a/apps/application/workflow/nodes/image_understand_node/image_understand_node.py b/apps/application/workflow/nodes/image_understand_node/image_understand_node.py index 2aa54796399..2378fb83995 100644 --- a/apps/application/workflow/nodes/image_understand_node/image_understand_node.py +++ b/apps/application/workflow/nodes/image_understand_node/image_understand_node.py @@ -17,6 +17,7 @@ from application.workflow.status import Status from application.workflow.tools import Reasoning from common.utils.common import guess_image_format +from common.utils.count_tokens import count_local_tokens from common.exception.app_exception import AppApiException from knowledge.models import File from models_provider.models import Model @@ -162,14 +163,14 @@ def execute(self): async def _stream_response( self, response, chat_model, message_list, question, reasoning_content_id, text_content_id, node_info, is_result ): + answer = "" + reasoning_content = "" try: model_setting = self.get_context("model_setting") or {} reasoning = Reasoning( model_setting.get("reasoning_content_start", ""), model_setting.get("reasoning_content_end", ""), ) - answer = "" - reasoning_content = "" response_reasoning_content = False async for chunk in response: @@ -255,15 +256,18 @@ async def _stream_response( ) ) - self._write_final_context(chat_model, message_list, question, answer, reasoning_content) - + # 无论正常结束、抛异常还是被取消中断,都落盘已累积内容并关闭 LLM 流式响应底层连接 finally: + self._write_final_context(chat_model, message_list, question, answer, reasoning_content) with contextlib.suppress(Exception): await response.aclose() def _write_final_context(self, chat_model, message_list, question, answer, reasoning_content): - message_tokens = chat_model.get_num_tokens_from_messages(message_list) - answer_tokens = chat_model.get_num_tokens(answer) + message_tokens, answer_tokens = ( + count_local_tokens(message_list, answer) + if self.async_runner.cancel_event.is_set() + else (chat_model.get_num_tokens_from_messages(message_list), chat_model.get_num_tokens(answer)) + ) self.write_context("message_tokens", message_tokens) self.write_context("answer_tokens", answer_tokens) self.write_context("answer", answer) diff --git a/apps/application/workflow/nodes/question_node/question_node.py b/apps/application/workflow/nodes/question_node/question_node.py index 36af7444af7..7db242b690f 100644 --- a/apps/application/workflow/nodes/question_node/question_node.py +++ b/apps/application/workflow/nodes/question_node/question_node.py @@ -19,6 +19,7 @@ from application.workflow.message.struct.content import NodeInfo, Position from application.workflow.message.struct.text_content import TextContent from application.workflow.status import Status +from common.utils.count_tokens import count_local_tokens from common.exception.app_exception import AppApiException from models_provider.models import Model from models_provider.tools import get_model_instance_by_model_workspace_id, get_model_credential @@ -158,8 +159,8 @@ def execute(self): ) async def _stream_response(self, response, chat_model, message_list, is_result): + answer = "" try: - answer = "" async for chunk in response: self._check_cancelled() answer += chunk.content @@ -170,16 +171,22 @@ async def _stream_response(self, response, chat_model, message_list, is_result): self.get_node_id(), chunk.content, Status.SUCCESS, node_info, Position(self.get_node_id()) ) ) - - message_tokens = chat_model.get_num_tokens_from_messages(message_list) - answer_tokens = chat_model.get_num_tokens(answer) - self.write_context("message_tokens", message_tokens) - self.write_context("answer_tokens", answer_tokens) - self.write_context("answer", answer) + # 无论正常结束、抛异常还是被取消中断,都落盘已累积内容并关闭 LLM 流式响应底层连接 finally: + self._write_final_context(chat_model, message_list, answer) with contextlib.suppress(Exception): await response.aclose() + def _write_final_context(self, chat_model, message_list, answer): + message_tokens, answer_tokens = ( + count_local_tokens(message_list, answer) + if self.async_runner.cancel_event.is_set() + else (chat_model.get_num_tokens_from_messages(message_list), chat_model.get_num_tokens(answer)) + ) + self.write_context("message_tokens", message_tokens) + self.write_context("answer_tokens", answer_tokens) + self.write_context("answer", answer) + def get_details(self, index: int = 0, position: dict = None, old_details: dict = None, **kwargs): details = super().get_details(index, position, old_details, **kwargs) details.update( diff --git a/apps/application/workflow/nodes/video_understand_node/video_understand_node.py b/apps/application/workflow/nodes/video_understand_node/video_understand_node.py index e7bd08f3917..956af5333c9 100644 --- a/apps/application/workflow/nodes/video_understand_node/video_understand_node.py +++ b/apps/application/workflow/nodes/video_understand_node/video_understand_node.py @@ -15,6 +15,7 @@ from application.workflow.message.struct.text_content import TextContent from application.workflow.status import Status from application.workflow.tools import Reasoning +from common.utils.count_tokens import count_local_tokens from common.exception.app_exception import AppApiException from knowledge.models import File from models_provider.models import Model @@ -160,14 +161,14 @@ def execute(self): async def _stream_response( self, response, chat_model, message_list, question, reasoning_content_id, text_content_id, node_info, is_result ): + answer = "" + reasoning_content = "" try: model_setting = self.get_context("model_setting") or {} reasoning = Reasoning( model_setting.get("reasoning_content_start", ""), model_setting.get("reasoning_content_end", ""), ) - answer = "" - reasoning_content = "" response_reasoning_content = False async for chunk in response: @@ -253,14 +254,18 @@ async def _stream_response( ) ) - self._write_final_context(chat_model, message_list, question, answer, reasoning_content) + # 无论正常结束、抛异常还是被取消中断,都落盘已累积内容并关闭 LLM 流式响应底层连接 finally: + self._write_final_context(chat_model, message_list, question, answer, reasoning_content) with contextlib.suppress(Exception): await response.aclose() def _write_final_context(self, chat_model, message_list, question, answer, reasoning_content): - message_tokens = chat_model.get_num_tokens_from_messages(message_list) - answer_tokens = chat_model.get_num_tokens(answer) + message_tokens, answer_tokens = ( + count_local_tokens(message_list, answer) + if self.async_runner.cancel_event.is_set() + else (chat_model.get_num_tokens_from_messages(message_list), chat_model.get_num_tokens(answer)) + ) self.write_context("message_tokens", message_tokens) self.write_context("answer_tokens", answer_tokens) self.write_context("answer", answer) diff --git a/apps/common/utils/count_tokens.py b/apps/common/utils/count_tokens.py new file mode 100644 index 00000000000..56c82fa201b --- /dev/null +++ b/apps/common/utils/count_tokens.py @@ -0,0 +1,11 @@ +from langchain_core.messages import get_buffer_string + +from common.config.tokenizer_manage_config import TokenizerManage + + +def count_local_tokens(message_list, answer): + """模型不可用时,用本地 tokenizer 估算输入/输出 token 数。""" + tokenizer = TokenizerManage.get_tokenizer() + message_tokens = sum([len(tokenizer.encode(get_buffer_string([m]))) for m in message_list]) + answer_tokens = len(tokenizer.encode(answer)) + return message_tokens, answer_tokens