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 20df7753c4e..2aa54796399 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
@@ -1,5 +1,6 @@
# coding=utf-8
import base64
+import contextlib
import uuid_utils.compat as uuid
from functools import reduce
@@ -8,7 +9,7 @@
from langchain_core.messages import HumanMessage, SystemMessage, 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.reasoning_content import ReasoningContent
@@ -52,6 +53,14 @@ class ImageUnderstandNode(INode):
supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
type = "image-understand-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()
@@ -137,45 +146,75 @@ def execute(self):
node_info = NodeInfo(self.get_node_id(), self.get_node_name(), Status.RUNNING)
- r = chat_model.stream(message_list)
- self._stream_response(
- r, chat_model, message_list, question, reasoning_content_id, text_content_id, node_info, is_result
+ self.async_runner.run_async(
+ lambda: self._stream_response(
+ chat_model.astream(message_list),
+ chat_model,
+ message_list,
+ question,
+ reasoning_content_id,
+ text_content_id,
+ node_info,
+ is_result,
+ )
)
- def _stream_response(
+ async def _stream_response(
self, response, chat_model, message_list, question, reasoning_content_id, text_content_id, node_info, is_result
):
- 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
-
- if is_result:
- if isinstance(chunk.content, list):
- for chunk_item in chunk.content:
- text = chunk_item.get("text", "")
- if text:
+ 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:
+ 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
+
+ if is_result:
+ if isinstance(chunk.content, list):
+ for chunk_item in chunk.content:
+ text = chunk_item.get("text", "")
+ if text:
+ self.write(
+ TextContent(
+ text_content_id, text, 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()),
+ )
+ )
+ else:
+ if content_chunk:
self.write(
TextContent(
- text_content_id, text, Status.RUNNING, node_info, Position(self.get_node_id())
+ 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):
@@ -188,52 +227,39 @@ def _stream_response(
Position(self.get_node_id()),
)
)
- else:
- if content_chunk:
- 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,
- 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 is_result:
- if reasoning_end.get("content"):
- self.write(
- TextContent(
- text_content_id,
- reasoning_end.get("content"),
- 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 is_result:
+ 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()),
+ 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)
+
+ finally:
+ 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)
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 8f9d3b3b8ac..e7bd08f3917 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
@@ -1,4 +1,5 @@
# coding=utf-8
+import contextlib
import uuid_utils.compat as uuid
from functools import reduce
@@ -7,7 +8,7 @@
from langchain_core.messages import HumanMessage, SystemMessage, 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.reasoning_content import ReasoningContent
@@ -50,6 +51,14 @@ class VideoUnderstandNode(INode):
supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL]
type = "video-understand-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()
@@ -135,45 +144,75 @@ def execute(self):
node_info = NodeInfo(self.get_node_id(), self.get_node_name(), Status.RUNNING)
- r = chat_model.stream(message_list)
- self._stream_response(
- r, chat_model, message_list, question, reasoning_content_id, text_content_id, node_info, is_result
+ self.async_runner.run_async(
+ lambda: self._stream_response(
+ chat_model.astream(message_list),
+ chat_model,
+ message_list,
+ question,
+ reasoning_content_id,
+ text_content_id,
+ node_info,
+ is_result,
+ )
)
- def _stream_response(
+ async def _stream_response(
self, response, chat_model, message_list, question, reasoning_content_id, text_content_id, node_info, is_result
):
- 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
-
- if is_result:
- if isinstance(chunk.content, list):
- for chunk_item in chunk.content:
- text = chunk_item.get("text", "")
- if text:
+ 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:
+ 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
+
+ if is_result:
+ if isinstance(chunk.content, list):
+ for chunk_item in chunk.content:
+ text = chunk_item.get("text", "")
+ if text:
+ self.write(
+ TextContent(
+ text_content_id, text, 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()),
+ )
+ )
+ else:
+ if content_chunk:
self.write(
TextContent(
- text_content_id, text, Status.RUNNING, node_info, Position(self.get_node_id())
+ 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):
@@ -186,52 +225,38 @@ def _stream_response(
Position(self.get_node_id()),
)
)
- else:
- if content_chunk:
- 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,
- 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 is_result:
- if reasoning_end.get("content"):
- self.write(
- TextContent(
- text_content_id,
- reasoning_end.get("content"),
- 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 is_result:
+ 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()),
+ 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)
+ finally:
+ 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)