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