From 01645d7bee9c649373031a849d5040da1b0c6cb4 Mon Sep 17 00:00:00 2001 From: hexiaonan-800 Date: Sat, 10 Oct 2026 15:23:30 +0800 Subject: [PATCH] feat: node cancel --- .../image_generate_node.py | 16 ++++++- .../image_to_video_node.py | 20 +++++++- .../workflow/nodes/intent_node/intent_node.py | 18 ++++++- .../speech_to_text_node.py | 16 ++++++- .../text_to_speech_node.py | 47 ++++++++++++------- .../text_to_video_node/text_to_video_node.py | 16 ++++++- 6 files changed, 107 insertions(+), 26 deletions(-) diff --git a/apps/application/workflow/nodes/image_generate_node/image_generate_node.py b/apps/application/workflow/nodes/image_generate_node/image_generate_node.py index c2776a1c50d..a18de5d2b2c 100644 --- a/apps/application/workflow/nodes/image_generate_node/image_generate_node.py +++ b/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( diff --git a/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py b/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py index 624a71ab0e9..33ba9d0bd26 100644 --- a/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py +++ b/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): diff --git a/apps/application/workflow/nodes/intent_node/intent_node.py b/apps/application/workflow/nodes/intent_node/intent_node.py index bb0e884c2f6..7012ddc9831 100644 --- a/apps/application/workflow/nodes/intent_node/intent_node.py +++ b/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( diff --git a/apps/application/workflow/nodes/speech_to_text_node/speech_to_text_node.py b/apps/application/workflow/nodes/speech_to_text_node/speech_to_text_node.py index 37b194d5535..0cc637fc7ee 100644 --- a/apps/application/workflow/nodes/speech_to_text_node/speech_to_text_node.py +++ b/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( diff --git a/apps/application/workflow/nodes/text_to_speech_node/text_to_speech_node.py b/apps/application/workflow/nodes/text_to_speech_node/text_to_speech_node.py index 79e2ad9d8cc..36b844223e4 100644 --- a/apps/application/workflow/nodes/text_to_speech_node/text_to_speech_node.py +++ b/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'' @@ -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) diff --git a/apps/application/workflow/nodes/text_to_video_node/text_to_video_node.py b/apps/application/workflow/nodes/text_to_video_node/text_to_video_node.py index c4ad6b989db..07c8b723b05 100644 --- a/apps/application/workflow/nodes/text_to_video_node/text_to_video_node.py +++ b/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)