diff --git a/apps/application/workflow/loop_workflow_manage.py b/apps/application/workflow/loop_workflow_manage.py new file mode 100644 index 00000000000..4f78a73d9e4 --- /dev/null +++ b/apps/application/workflow/loop_workflow_manage.py @@ -0,0 +1,41 @@ +# coding=utf-8 +""" + @project: MaxKB + @Author:虎虎虎 + @file: loop_workflow_manage.py + @date:2026/7/2 10:00 + @desc: +""" +from typing import Dict, Optional, Callable + +from application.workflow.common import Workflow, WorkflowType, Node +from application.workflow.i_node import INode +from application.workflow.workflow_manage import WorkflowManage, CallBack + + +class LoopWorkFlowManage(WorkflowManage): + + def __init__(self, + workflow: Workflow, + parameters: Dict, + workflow_type: WorkflowType, + call_back: CallBack, + get_start_node: Callable[[Workflow, WorkflowManage], INode], + parent_workflow_manage: WorkflowManage, + loop_context: Dict = None): + self.parent_workflow_manage = parent_workflow_manage + self.loop_context = loop_context or {} + super().__init__(workflow, parameters, workflow_type, call_back, get_start_node) + + def get_parameters(self): + return {**self.parameters, **self.loop_context} + + def get_parent_context(self, node_id, key): + return self.parent_workflow_manage.get_context(node_id, key) + + def generate_prompt(self, prompt): + prompt = self.workflow.reset_prompt(prompt) + context = {**self.context, **self.parent_workflow_manage.context} + from langchain_core.prompts import PromptTemplate + prompt_template = PromptTemplate.from_template(prompt, template_format='jinja2') + return prompt_template.format(context=context) diff --git a/apps/application/workflow/nodes/loop_node/__init__.py b/apps/application/workflow/nodes/loop_node/__init__.py new file mode 100644 index 00000000000..826ece19728 --- /dev/null +++ b/apps/application/workflow/nodes/loop_node/__init__.py @@ -0,0 +1,9 @@ +# coding=utf-8 +""" + @project: MaxKB + @Author:虎虎虎 + @file: __init__.py + @date:2026/7/2 10:00 + @desc: +""" +from .loop_node import LoopNode diff --git a/apps/application/workflow/nodes/loop_node/loop_node.py b/apps/application/workflow/nodes/loop_node/loop_node.py new file mode 100644 index 00000000000..d2b7accea0c --- /dev/null +++ b/apps/application/workflow/nodes/loop_node/loop_node.py @@ -0,0 +1,151 @@ +# coding=utf-8 +""" + @project: MaxKB + @Author:虎虎虎 + @file: loop_node.py + @date:2026/7/2 10:00 + @desc: +""" +import time + +from django.utils.translation import gettext_lazy as _ +from rest_framework import serializers + +from application.workflow.common import WorkflowType, new_instance +from application.workflow.i_node import INode + +from application.workflow.message.struct.content import NodeInfo +from application.workflow.message.struct.text_content import TextContent +from application.workflow.status import Status + +from common.exception.app_exception import AppApiException + +MAX_LOOP_COUNT = 500 + + +class LoopNodeSerializer(serializers.Serializer): + loop_type = serializers.CharField(required=True, label=_("loop_type")) + array = serializers.ListField(required=False, allow_null=True, label=_("array")) + number = serializers.IntegerField(required=False, allow_null=True, label=_("number")) + loop_body = serializers.DictField(required=True, label="循环体") + + def is_valid(self, *, raise_exception=False): + super().is_valid(raise_exception=True) + loop_type = self.data.get('loop_type') + if loop_type == 'ARRAY': + array = self.data.get('array') + if array is None or len(array) == 0: + message = _('{field}, this field is required.', field='array') + raise AppApiException(500, message) + elif loop_type == 'NUMBER': + number = self.data.get('number') + if number is None: + message = _('{field}, this field is required.', field='number') + raise AppApiException(500, message) + + +def _generate_loop_number(number): + return iter([(i, i) for i in range(number)]) + + +def _generate_loop_array(array): + return iter([(item, i) for i, item in enumerate(array)]) + + +def _generate_while_loop(number): + return iter([(i, i) for i in range(number)]) + + +class LoopNode(INode): + serializer_class = LoopNodeSerializer + supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL] + type = 'loop-node' + + def _run(self): + self.execute() + + def execute(self): + node_params = self.get_parameters() + workflow_params = self.get_workflow_parameters() + + loop_type = node_params.get('loop_type') + array = node_params.get('array') + number = node_params.get('number') + loop_body = node_params.get('loop_body') + + if loop_type == 'ARRAY' and isinstance(array, list) and len(array) >= 2: + array = self.workflow_manage.get_reference_field(array[0], array[1:]) + + self.write_context('params', {'loop_type': loop_type, 'array': array, 'number': number}) + + if loop_type == 'ARRAY': + iterator = _generate_loop_array(array) + elif loop_type == 'LOOP': + iterator = _generate_while_loop(number or MAX_LOOP_COUNT) + else: + iterator = _generate_loop_number(number) + + self._loop_node_data = [] + self._loop_answer_data = [] + self._answer_text = '' + self._workflow_params = workflow_params + self._loop_body = loop_body + self._iterator = iterator + + self._run_next() + + def _run_next(self): + try: + item, index = next(self._iterator) + except StopIteration: + self.write_context('answer', self._answer_text) + self.write_context('run_time', time.time() - self.data.get('start_time', time.time())) + self.complete(Status.SUCCESS) + return + loop_context = {'index': index, 'item': item} + workflow = new_instance(self._loop_body, self.get_workflow_type()) + + chunk_list = [] + + def on_next(wf_manage, content): + chunk_list.append(content) + if hasattr(content, 'content'): + self._answer_text += content.content + self.write(content) + + def on_complete(wf_manage, error): + self._loop_node_data.append(wf_manage.context) + self._loop_answer_data.append(chunk_list) + self.write_context('loop_node_data', self._loop_node_data) + self.write_context('loop_answer_data', self._loop_answer_data) + self.write_context('index', index) + self.write_context('item', item) + + if error: + self.write_context('error_message', str(error)) + return + + self._run_next() + + from application.workflow.workflow_manage import CallBack + from application.workflow.loop_workflow_manage import LoopWorkFlowManage + from application.workflow.nodes import get_node_class + call_back = CallBack(on_next, on_complete) + + loop_start_class = get_node_class('loop-start-node', self.get_workflow_type()) + + def get_start_node_fn(wf, wf_manage): + start_node = wf.get_node('loop-start-node') + return loop_start_class(start_node, wf_manage, lambda n: n.properties.get('node_data', {})) + + loop_manage = LoopWorkFlowManage( + workflow=workflow, + parameters=self._workflow_params, + workflow_type=self.get_workflow_type(), + call_back=call_back, + get_start_node=get_start_node_fn, + parent_workflow_manage=self.workflow_manage, + loop_context=loop_context, + ) + loop_manage.start_node.workflow_manage = loop_manage + loop_manage.run() diff --git a/apps/application/workflow/nodes/loop_start_node/__init__.py b/apps/application/workflow/nodes/loop_start_node/__init__.py new file mode 100644 index 00000000000..7d1dbb5f046 --- /dev/null +++ b/apps/application/workflow/nodes/loop_start_node/__init__.py @@ -0,0 +1,9 @@ +# coding=utf-8 +""" + @project: MaxKB + @Author:虎虎虎 + @file: __init__.py + @date:2026/7/2 10:00 + @desc: +""" +from .loop_start_node import LoopStartNode diff --git a/apps/application/workflow/nodes/loop_start_node/loop_start_node.py b/apps/application/workflow/nodes/loop_start_node/loop_start_node.py new file mode 100644 index 00000000000..aa977562f72 --- /dev/null +++ b/apps/application/workflow/nodes/loop_start_node/loop_start_node.py @@ -0,0 +1,30 @@ +# coding=utf-8 +""" + @project: MaxKB + @Author:虎虎虎 + @file: loop_start_node.py + @date:2026/7/2 10:00 + @desc: +""" +from rest_framework import serializers + +from application.workflow.common import WorkflowType +from application.workflow.i_node import INode + + +class LoopStartNodeSerializer(serializers.Serializer): + pass + + +class LoopStartNode(INode): + serializer_class = LoopStartNodeSerializer + supported_workflow_type_list = [WorkflowType.APPLICATION, WorkflowType.KNOWLEDGE, WorkflowType.TOOL] + type = 'loop-start-node' + + def execute(self): + loop_context = getattr(self.workflow_manage, 'loop_context', {}) + index = loop_context.get('index', 0) + item = loop_context.get('item', None) + + self.write_context('index', index) + self.write_context('item', item) diff --git a/apps/application/workflow/nodes/reply_node/reply_node.py b/apps/application/workflow/nodes/reply_node/reply_node.py index cdd79e7acef..3c27efa263b 100644 --- a/apps/application/workflow/nodes/reply_node/reply_node.py +++ b/apps/application/workflow/nodes/reply_node/reply_node.py @@ -7,7 +7,7 @@ @desc: """ from typing import List - +import uuid_utils.compat as uuid from django.utils.translation import gettext_lazy as _ from rest_framework import serializers @@ -34,7 +34,7 @@ class ReplyNode(INode): def execute(self): node_params = self.get_parameters() workflow_params = self.get_workflow_parameters() - chat_record_id = workflow_params.get('chat_record_id') + chunk_id = uuid.uuid7() reply_type = node_params.get('reply_type') fields = node_params.get('fields') @@ -50,7 +50,7 @@ def execute(self): if is_result: node_info = NodeInfo(self.get_node_id(), self.get_node_name(), Status.SUCCESS) - self.write(TextContent(str(chat_record_id), result, Status.SUCCESS, node_info)) + self.write(TextContent(str(chunk_id), result, Status.SUCCESS, node_info)) def _generate_reply_content(self, prompt): if prompt is None: diff --git a/apps/application/workflow/workflow_manage.py b/apps/application/workflow/workflow_manage.py index 70f4461a1d2..b83fd11a82c 100644 --- a/apps/application/workflow/workflow_manage.py +++ b/apps/application/workflow/workflow_manage.py @@ -16,7 +16,7 @@ from application.workflow.common import Workflow, WorkflowType, Node, get_node_parameters from application.workflow.i_node import INode from application.workflow.message.struct.content import Content -from application.workflow.nodes import get_node_class + from application.workflow.status import Status @@ -77,6 +77,7 @@ def next_nodes(self, nodes: Optional[List[Node]]): if nodes is None or len(nodes) == 0: return with self._lock: + from application.workflow.nodes import get_node_class instances = [get_node_class(n.type, self.workflow_type)(n, self, get_node_parameters) for n in nodes] self.nodes.extend(instances) @@ -89,7 +90,8 @@ def assertion_end(self): return if not self.is_end(): return - self.end() + self.done = True # 锁内抢占,保证只有一个线程能往下走 + self.end() # 回调放到锁外,避免回调里再触碰本 manage 造成重入/死锁 def is_end(self): """ @@ -148,10 +150,7 @@ def end(self): 工作流输出结束的时候调用 @return: None """ - if self.done: - return self.call_back.on_complete(self, None) - self.done = True def get_parameters(self): """