feat: node cancel

apps/application/workflow/nodes/image_generate_node/image_generate_node.py

@@ -7,6 +7,7 @@
 @desc:
 """
 
+import asyncio
 import base64
 from functools import reduce
 
@@ -17,7 +18,7 @@
 from langchain_core.messages import HumanMessage, AIMessage
 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.struct.content import NodeInfo, Position
 from application.workflow.message.struct.text_content import TextContent
@@ -62,6 +63,14 @@ class ImageGenerateNode(INode):
     supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
     type = "image-generate-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 execute(self):
         node_params = self.get_parameters()
         workflow_params = self.get_workflow_parameters()
@@ -112,7 +121,7 @@ def execute(self):
         self.write_context("dialogue_type", dialogue_type)
 
         self._check_cancelled()
-        image_urls = tti_model.generate_image(question, negative_prompt)
+        image_urls = self.async_runner.run_async(lambda: self._await_generate(tti_model, question, negative_prompt))
 
         file_urls = []
         for image_url in image_urls:
@@ -141,6 +150,9 @@ def execute(self):
                 TextContent(str(uuid.uuid7()), answer, Status.SUCCESS, node_info, position=Position(self.get_node_id()))
             )
 
+    async def _await_generate(self, tti_model, question, negative_prompt):
+        return await asyncio.to_thread(tti_model.generate_image, question, negative_prompt)
+
     def _get_history_message(self, history_chat_record, dialogue_number):
         start_index = len(history_chat_record) - dialogue_number
         return reduce(

apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py

@@ -1,4 +1,5 @@
 # coding=utf-8
+import asyncio
 import base64
 import uuid_utils.compat as uuid
 import requests
@@ -10,7 +11,7 @@
 from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
 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.struct.content import NodeInfo, Position
 from application.workflow.message.struct.text_content import TextContent
@@ -58,6 +59,14 @@ class ImageToVideoNode(INode):
     supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
     type = "image-to-video-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 execute(self):
         maxkb_logger.info(f"[ImageToVideoNode] execute START, node_id={self.get_node_id()}")
         workflow_params = self.get_workflow_parameters()
@@ -136,7 +145,9 @@ def execute(self):
         last_frame_url = self._get_file_base64(last_frame_url)
 
         self._check_cancelled()
-        video_urls = ttv_model.generate_video(question, negative_prompt, first_frame_url, last_frame_url)
+        video_urls = self.async_runner.run_async(
+            lambda: self._await_generate(ttv_model, question, negative_prompt, first_frame_url, last_frame_url)
+        )
         maxkb_logger.info(
             f"[ImageToVideoNode] generate_video result: {video_urls is not None}, node_id={self.get_node_id()}"
         )
@@ -164,6 +175,11 @@ def execute(self):
                 TextContent(str(uuid.uuid7()), video_label, Status.SUCCESS, node_info, Position(self.get_node_id()))
             )
 
+    async def _await_generate(self, ttv_model, question, negative_prompt, first_frame_url, last_frame_url):
+        return await asyncio.to_thread(
+            ttv_model.generate_video, question, negative_prompt, first_frame_url, last_frame_url
+        )
+
     def _get_file_base64(self, image_url):
         try:
             if isinstance(image_url, list):

apps/application/workflow/nodes/intent_node/intent_node.py

@@ -9,7 +9,7 @@
 from langchain_core.messages import HumanMessage, SystemMessage
 from rest_framework import serializers
 
-from application.workflow.common import WorkflowType
+from application.workflow.common import AsyncRunner, CancelledException, WorkflowType
 from application.workflow.i_node import INode
 from application.workflow.status import Status
 from common.exception.app_exception import AppApiException
@@ -57,6 +57,14 @@ class IntentNode(INode):
     supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
     type = "intent-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 execute(self):
         workflow_params = self.get_workflow_parameters()
         node_params = self.get_parameters()
@@ -120,7 +128,7 @@ def execute(self):
 
         try:
             self._check_cancelled()
-            r = chat_model.invoke(message_list)
+            r = self.async_runner.run_async(lambda: self._await_invoke(chat_model, message_list))
             classification_result = r.content.strip()
             matched_branch = self._parse_classification_result(classification_result, branch)
 
@@ -135,6 +143,8 @@ def execute(self):
 
             self.complete(Status.SUCCESS, [self.branch_anchor(matched_branch["id"])])
 
+        except CancelledException:
+            raise  # 取消必须放行给 INode.run 置 CANCELLED,不能被 except Exception 吞成 other_branch
         except Exception as e:
             other_branch = self._find_other_branch(branch)
             if other_branch:
@@ -145,6 +155,10 @@ def execute(self):
             else:
                 raise Exception(f"error: {str(e)}")
 
+    async def _await_invoke(self, chat_model, message_list):
+        """走原生 async 调用;协程被取消时 langchain 会断开底层连接,真正停止模型请求。"""
+        return await chat_model.ainvoke(message_list)
+
     def _get_history_message(self, history_chat_record, dialogue_number):
         start_index = len(history_chat_record) - dialogue_number
         history_message = reduce(

apps/application/workflow/nodes/speech_to_text_node/speech_to_text_node.py

@@ -5,6 +5,7 @@
 @desc:
 """
 
+import asyncio
 import os
 import tempfile
 from concurrent.futures import ThreadPoolExecutor
@@ -13,7 +14,7 @@
 from django.utils.translation import gettext_lazy as _
 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.struct.content import NodeInfo, Position
 from application.workflow.message.struct.text_content import TextContent
@@ -73,6 +74,14 @@ class SpeechToTextNode(INode):
     supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
     type = "speech-to-text-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 execute(self):
         node_params = self.get_parameters()
         workflow_params = self.get_workflow_parameters()
@@ -115,7 +124,7 @@ def execute(self):
         self.write_context("audio_list", audio_list)
 
         self._check_cancelled()
-        result = _process_audio_items(audio_list, stt_model)
+        result = self.async_runner.run_async(lambda: self._await_process(audio_list, stt_model))
         content = []
         result_content = []
         for item in result:
@@ -132,6 +141,9 @@ def execute(self):
             node_info = NodeInfo(self.get_node_id(), self.get_node_name(), Status.SUCCESS)
             self.write(TextContent(self.get_node_id(), answer, Status.SUCCESS, node_info, Position(self.get_node_id())))
 
+    async def _await_process(self, audio_list, stt_model):
+        return await asyncio.to_thread(_process_audio_items, audio_list, stt_model)
+
     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/text_to_speech_node/text_to_speech_node.py

@@ -5,6 +5,7 @@
 @desc:
 """
 
+import asyncio
 import io
 import mimetypes
 
@@ -14,7 +15,7 @@
 from pydub import AudioSegment
 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.struct.content import NodeInfo, Position
 from application.workflow.message.struct.text_content import TextContent
@@ -69,6 +70,14 @@ class TextToSpeechNode(INode):
     supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
     type = "text-to-speech-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 execute(self):
         node_params = self.get_parameters()
         workflow_params = self.get_workflow_parameters()
@@ -109,19 +118,11 @@ def execute(self):
         max_length = 1024
         content_chunks = [content[i : i + max_length] for i in range(0, len(content), max_length)]
 
-        audio_segments = []
-        temp_files = []
-
-        for chunk in content_chunks:
-            self._check_cancelled()
-            self.write_context("content", chunk)
-            workspace_id = workflow_params.get("workspace_id")
-            model = get_model_instance_by_model_workspace_id(tts_model_id, workspace_id, **(model_params_setting or {}))
-            audio_byte = model.text_to_speech(chunk)
-            temp_file = io.BytesIO(audio_byte)
-            audio_segment = AudioSegment.from_file(temp_file)
-            audio_segments.append(audio_segment)
-            temp_files.append(temp_file)
+        workspace_id = workflow_params.get("workspace_id")
+        audio_segments = self.async_runner.run_async(
+            lambda: self._process_tts_chunks(content_chunks, tts_model_id, workspace_id, model_params_setting)
+        )
+        self.write_context("content", content_chunks[-1] if content_chunks else "")
 
         combined_audio = AudioSegment.empty()
         for segment in audio_segments:
@@ -137,8 +138,6 @@ def execute(self):
         file_id = file_url.split("/")[-1]
         audio_list = [{"file_id": file_id, "file_name": file_name, "url": file_url}]
 
-        for temp_file in temp_files:
-            temp_file.close()
         output_buffer.close()
 
         audio_label = f'<audio src="{file_url}" controls style="width: 300px; height: 43px"></audio>'
@@ -151,6 +150,22 @@ def execute(self):
                 TextContent(self.get_node_id(), audio_label, Status.SUCCESS, node_info, Position(self.get_node_id()))
             )
 
+    async def _process_tts_chunks(self, content_chunks, tts_model_id, workspace_id, model_params_setting):
+        def process():
+            audio_segments = []
+            for chunk in content_chunks:
+                self._check_cancelled()
+                model = get_model_instance_by_model_workspace_id(
+                    tts_model_id, workspace_id, **(model_params_setting or {})
+                )
+                audio_byte = model.text_to_speech(chunk)
+                temp_file = io.BytesIO(audio_byte)
+                audio_segment = AudioSegment.from_file(temp_file)
+                audio_segments.append(audio_segment)
+            return audio_segments
+
+        return await asyncio.to_thread(process)
+
     def _upload_file(self, file, workflow_params, workflow_type):
         if workflow_type == WorkflowType.KNOWLEDGE:
             return self._upload_knowledge_file(file, workflow_params)

apps/application/workflow/nodes/text_to_video_node/text_to_video_node.py

@@ -1,4 +1,5 @@
 # coding=utf-8
+import asyncio
 import uuid_utils.compat as uuid
 import requests
 from functools import reduce
@@ -9,7 +10,7 @@
 from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
 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.struct.content import NodeInfo, Position
 from application.workflow.message.struct.text_content import TextContent
@@ -55,6 +56,14 @@ class TextToVideoNode(INode):
     supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
     type = "text-to-video-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 execute(self):
         maxkb_logger.info(f"[TextToVideoNode] execute START, node_id={self.get_node_id()}")
         workflow_params = self.get_workflow_parameters()
@@ -117,7 +126,7 @@ def execute(self):
         self.write_context("negative_prompt", self.workflow_manage.generate_prompt(negative_prompt))
 
         self._check_cancelled()
-        video_urls = ttv_model.generate_video(question, negative_prompt)
+        video_urls = self.async_runner.run_async(lambda: self._await_generate(ttv_model, question, negative_prompt))
         maxkb_logger.info(
             f"[TextToVideoNode] generate_video result: {video_urls is not None}, node_id={self.get_node_id()}"
         )
@@ -145,6 +154,9 @@ def execute(self):
                 TextContent(str(uuid.uuid7()), video_label, Status.SUCCESS, node_info, Position(self.get_node_id()))
             )
 
+    async def _await_generate(self, ttv_model, question, negative_prompt):
+        return await asyncio.to_thread(ttv_model.generate_video, question, negative_prompt)
+
     def _upload_file(self, file, workflow_type, workflow_params):
         if workflow_type == WorkflowType.KNOWLEDGE:
             return self._upload_knowledge_file(file, workflow_params)