feat: count tokens and runtime when cancel

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,15 +397,15 @@ 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 {}
             reasoning = Reasoning(
                 model_setting.get("reasoning_content_start", "<think>"),
                 model_setting.get("reasoning_content_end", "</think>"),
             )
-            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)

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", "<think>"),
                 model_setting.get("reasoning_content_end", "</think>"),
             )
-            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)

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(

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", "<think>"),
                 model_setting.get("reasoning_content_end", "</think>"),
             )
-            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)

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