diff --git a/apps/application/workflow/common.py b/apps/application/workflow/common.py
index 802b2a4ee39..1fc5da8a578 100644
--- a/apps/application/workflow/common.py
+++ b/apps/application/workflow/common.py
@@ -7,6 +7,9 @@
@desc:
"""
+import asyncio
+import contextlib
+import threading
from enum import Enum
from typing import List, Dict
@@ -246,3 +249,52 @@ def new_instance(flow_obj: Dict, workflow_type: WorkflowType = WorkflowType.APPL
nodes = [Node(node.get("id"), node.get("type"), **node) for node in nodes]
edges = [Edge(edge.get("id"), edge.get("type"), **edge) for edge in edges]
return Workflow(nodes, edges)
+
+
+class CancelledException(Exception):
+ """工作流取消异常"""
+
+ pass
+
+
+class AsyncRunner:
+ """跨线程取消感知地运行异步协程;cancel() 置位事件使阻塞中的流式中断并抛 CancelledException。"""
+
+ def __init__(self):
+ self.cancel_event = threading.Event()
+
+ def cancel(self):
+ """请求取消(任意线程可调);只置旗标,真正中断发生在事件循环内。"""
+ self.cancel_event.set()
+
+ def run_async(self, async_factory):
+ """在独立事件循环里运行 async_factory() 的协程;被取消则抛 CancelledException。"""
+ return asyncio.run(self.await_or_cancel(async_factory()))
+
+ async def _wait_cancel(self):
+ # 轮询取消旗标,置位即结束,充当竞速对手
+ while not self.cancel_event.is_set():
+ await asyncio.sleep(0)
+
+ async def await_or_cancel(self, _async):
+ """竞速等待:流式完成返回其结果;取消先到则抛 CancelledException。"""
+
+ # 真正要跑的流式协程,包成 task 才能与 waiter 一起放进 asyncio.wait
+ stream = asyncio.create_task(_async)
+ # 把"是否已取消"伪装成一个可等待的 task,充当 stream 的竞速对手。
+ waiter = asyncio.create_task(self._wait_cancel())
+ # 同时等两个 task,谁先完成(返回)谁进入 done;FIRST_COMPLETED = 取先到者。
+ done, _ = await asyncio.wait({stream, waiter}, return_when=asyncio.FIRST_COMPLETED)
+ # stream_done 为真 -> 流式自己先跑完;为假 -> 是 waiter 先完成,即取消先到。
+ stream_done = stream in done
+ # 无论谁赢,都把两个 task cancel 并等它们退干净,防止"输掉的那一方"
+ # 还悬在事件循环里。await 吞掉 CancelledError(取消是本意,不算错误)。
+ for task in (stream, waiter):
+ task.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await task
+
+ # 流式正常结束:把它的返回值交给调用方。
+ if stream_done:
+ return stream.result()
+ raise CancelledException()
diff --git a/apps/application/workflow/i_node.py b/apps/application/workflow/i_node.py
index 0b3a17c0c94..e29ff7f9157 100644
--- a/apps/application/workflow/i_node.py
+++ b/apps/application/workflow/i_node.py
@@ -14,19 +14,13 @@
from rest_framework import serializers
-from application.workflow.common import Node
+from application.workflow.common import CancelledException, Node
from application.workflow.message.struct.content import Content, NodeInfo, Position
from application.workflow.message.struct.progress_content import ProgressContent
from application.workflow.status import Status
from common.utils.logger import maxkb_logger
-class CancelledException(Exception):
- """工作流取消异常"""
-
- pass
-
-
class Signal(str, Enum):
BREAK = "BREAK"
CONTINUE = "CONTINUE"
diff --git a/apps/application/workflow/nodes/ai_chat_node/agent.py b/apps/application/workflow/nodes/ai_chat_node/agent.py
index 9e8a561bd83..3d47543e500 100644
--- a/apps/application/workflow/nodes/ai_chat_node/agent.py
+++ b/apps/application/workflow/nodes/ai_chat_node/agent.py
@@ -24,7 +24,7 @@
from langgraph.checkpoint.memory import MemorySaver
from application.workflow.backend.sandbox_shell import SandboxShellBackend
-from application.workflow.i_node import CancelledException
+from application.workflow.common import CancelledException
from application.workflow.nodes.ai_chat_node.tools.skill import init_skills
from maxkb.const import CONFIG
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 3c9e0f9a281..25068d9ccba 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
@@ -8,6 +8,7 @@
"""
import base64
+import contextlib
import json
import re
from functools import reduce
@@ -19,7 +20,7 @@
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage, AIMessageChunk
from rest_framework import serializers
-from application.workflow.common import WorkflowType
+from application.workflow.common import AsyncRunner, WorkflowType
from application.workflow.i_node import INode
from application.workflow.message.aggregator import AggregationManager
from application.workflow.message.struct.content import NodeInfo, Position, Content
@@ -218,6 +219,14 @@ class AIChatNode(INode):
supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
type = "ai-chat-node"
+ def __init__(self, node, workflow_manage, get_node_parameters):
+ super().__init__(node, workflow_manage, get_node_parameters)
+ self.async_runner = AsyncRunner()
+
+ def cancel(self):
+ self.async_runner.cancel()
+ super().cancel()
+
def write(self, message: Content):
super().write(message)
if not self.data.get("messages"):
@@ -345,9 +354,15 @@ def execute(self):
message_list_with_system = [SystemMessage(system)] + message_list
if stream:
- r = chat_model.stream(message_list_with_system)
- self._stream_response(
- r, chat_model, message_list_with_system, question.content, reasoning_content_id, text_content_id
+ self.async_runner.run_async(
+ lambda: self._stream_response(
+ chat_model.astream(message_list_with_system),
+ chat_model,
+ message_list_with_system,
+ question.content,
+ reasoning_content_id,
+ text_content_id,
+ )
)
else:
r = chat_model.invoke(message_list_with_system)
@@ -370,40 +385,71 @@ def _generate_prompt_question(self, question_str, model, vision, image_list, vid
videos = _process_videos(video, model)
return HumanMessage(content=[*videos, *images, {"type": "text", "text": question_str}])
- def _stream_response(self, response, chat_model, message_list, question, reasoning_content_id, text_content_id):
- 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", ""),
- model_setting.get("reasoning_content_end", ""),
- )
- answer = ""
- reasoning_content = ""
- response_reasoning_content = False
-
- for chunk in response:
- self._check_cancelled()
- reasoning_chunk = reasoning.get_reasoning_content(chunk)
- content_chunk = reasoning_chunk.get("content")
- if "reasoning_content" in chunk.additional_kwargs:
- response_reasoning_content = True
- reasoning_content_chunk = chunk.additional_kwargs.get("reasoning_content", "")
- else:
- reasoning_content_chunk = reasoning_chunk.get("reasoning_content")
- answer += content_chunk
- if reasoning_content_chunk is None:
- reasoning_content_chunk = ""
- reasoning_content += reasoning_content_chunk
- reasoning_end = False
- if content_chunk:
- if not reasoning_end:
+ async def _stream_response(
+ self, response, chat_model, message_list, question, reasoning_content_id, text_content_id
+ ):
+ 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", ""),
+ model_setting.get("reasoning_content_end", ""),
+ )
+ answer = ""
+ reasoning_content = ""
+ response_reasoning_content = False
+
+ async for chunk in response:
+ self._check_cancelled()
+ reasoning_chunk = reasoning.get_reasoning_content(chunk)
+ content_chunk = reasoning_chunk.get("content")
+ if "reasoning_content" in chunk.additional_kwargs:
+ response_reasoning_content = True
+ reasoning_content_chunk = chunk.additional_kwargs.get("reasoning_content", "")
+ else:
+ reasoning_content_chunk = reasoning_chunk.get("reasoning_content")
+ answer += content_chunk
+ if reasoning_content_chunk is None:
+ reasoning_content_chunk = ""
+ reasoning_content += reasoning_content_chunk
+ reasoning_end = False
+ if content_chunk:
+ if not reasoning_end:
+ self.write(
+ ReasoningContent(
+ reasoning_content_id, "", Status.SUCCESS, node_info, Position(self.get_node_id())
+ )
+ )
+ self.write(
+ TextContent(
+ text_content_id, content_chunk, Status.RUNNING, node_info, Position(self.get_node_id())
+ )
+ )
+ if reasoning_content_chunk and model_setting.get("reasoning_content_enable", False):
self.write(
ReasoningContent(
- reasoning_content_id, "", Status.SUCCESS, node_info, Position(self.get_node_id())
+ reasoning_content_id,
+ reasoning_content_chunk,
+ Status.RUNNING,
+ node_info,
+ Position(self.get_node_id()),
)
)
+
+ reasoning_end = reasoning.get_end_reasoning_content()
+ answer += reasoning_end.get("content")
+ reasoning_content_chunk = ""
+ if not response_reasoning_content:
+ reasoning_content_chunk = reasoning_end.get("reasoning_content")
+ if reasoning_end.get("content"):
self.write(
- TextContent(text_content_id, content_chunk, Status.RUNNING, node_info, Position(self.get_node_id()))
+ TextContent(
+ text_content_id,
+ reasoning_end.get("content"),
+ Status.RUNNING,
+ node_info,
+ Position(self.get_node_id()),
+ )
)
if reasoning_content_chunk and model_setting.get("reasoning_content_enable", False):
self.write(
@@ -416,33 +462,11 @@ def _stream_response(self, response, chat_model, message_list, question, reasoni
)
)
- reasoning_end = reasoning.get_end_reasoning_content()
- answer += reasoning_end.get("content")
- reasoning_content_chunk = ""
- if not response_reasoning_content:
- reasoning_content_chunk = reasoning_end.get("reasoning_content")
- if reasoning_end.get("content"):
- self.write(
- TextContent(
- text_content_id,
- reasoning_end.get("content"),
- Status.RUNNING,
- node_info,
- Position(self.get_node_id()),
- )
- )
- if reasoning_content_chunk and model_setting.get("reasoning_content_enable", False):
- self.write(
- ReasoningContent(
- reasoning_content_id,
- reasoning_content_chunk,
- Status.RUNNING,
- node_info,
- Position(self.get_node_id()),
- )
- )
-
- self._write_final_context(chat_model, message_list, question, answer, reasoning_content)
+ self._write_final_context(chat_model, message_list, question, answer, reasoning_content)
+ # 无论正常结束、抛异常还是被取消中断,都关闭 LLM 流式响应底层连接
+ finally:
+ with contextlib.suppress(Exception):
+ await response.aclose()
def _invoke_response(self, response, chat_model, message_list, question, is_result=False, text_content_id=None):
model_setting = self.get_context("model_setting") or {}
diff --git a/apps/application/workflow/nodes/data_source_web_node/data_source_web_node.py b/apps/application/workflow/nodes/data_source_web_node/data_source_web_node.py
index 06d75fd22aa..a5cfb5aae36 100644
--- a/apps/application/workflow/nodes/data_source_web_node/data_source_web_node.py
+++ b/apps/application/workflow/nodes/data_source_web_node/data_source_web_node.py
@@ -12,8 +12,8 @@
from django.utils.translation import gettext_lazy as _
from rest_framework import serializers
-from application.workflow.common import WorkflowType
-from application.workflow.i_node import CancelledException, INode
+from application.workflow.common import CancelledException, WorkflowType
+from application.workflow.i_node import INode
from common.utils.fork import ChildLink, Fork, ForkManage
from common.utils.logger import maxkb_logger
diff --git a/apps/application/workflow/workflow_manage.py b/apps/application/workflow/workflow_manage.py
index c0e6b7e9df1..1ea3e5bc3f7 100644
--- a/apps/application/workflow/workflow_manage.py
+++ b/apps/application/workflow/workflow_manage.py
@@ -273,6 +273,10 @@ def from_context(cls, get_context, workflow, parameters, workflow_type, call_bac
return None
def cancel(self):
- self.signal = Signal.CANCELLED
- for node in self.nodes:
+ with self._lock:
+ if self.signal == Signal.CANCELLED:
+ return
+ self.signal = Signal.CANCELLED
+ nodes = list(self.nodes)
+ for node in nodes:
node.cancel()