Knowledge workflow

apps/application/flow/common.py

@@ -6,7 +6,7 @@
     @date：2024/12/11 17:57
     @desc:
 """
-
+from enum import Enum
 from typing import List, Dict
 
 from django.db.models import QuerySet
@@ -90,6 +90,16 @@ def __init__(self, edge, node):
         self.node = node
 
 
+class WorkflowMode(Enum):
+    APPLICATION = "application"
+
+    APPLICATION_LOOP = "application-loop"
+
+    KNOWLEDGE = "knowledge"
+
+    KNOWLEDGE_LOOP = "knowledge-loop"
+
+
 class Workflow:
     """
     节点列表
@@ -112,7 +122,10 @@ class Workflow:
     """
     next_node_map: Dict[str, List[EdgeNode]]
 
-    def __init__(self, nodes: List[Node], edges: List[Edge]):
+    workflow_mode: WorkflowMode
+
+    def __init__(self, nodes: List[Node], edges: List[Edge],
+                 workflow_mode: WorkflowMode = WorkflowMode.APPLICATION.value):
         self.nodes = nodes
         self.edges = edges
         self.node_map = {node.id: node for node in nodes}
@@ -125,6 +138,7 @@ def __init__(self, nodes: List[Node], edges: List[Edge]):
         self.next_node_map = {key: [EdgeNode(edge, self.node_map.get(edge.targetNodeId)) for edge in edges] for
                               key, edges in
                               group_by(edges, key=lambda edge: edge.sourceNodeId).items()}
+        self.workflow_mode = workflow_mode
 
     def get_node(self, node_id):
         """
@@ -167,13 +181,13 @@ def get_next_nodes(self, node_id) -> List[Node]:
         return [en.node for en in self.next_node_map.get(node_id, [])]
 
     @staticmethod
-    def new_instance(flow_obj: Dict):
+    def new_instance(flow_obj: Dict, workflow_mode: WorkflowMode = WorkflowMode.APPLICATION):
         nodes = flow_obj.get('nodes')
         edges = flow_obj.get('edges')
         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)
+        return Workflow(nodes, edges, workflow_mode)
 
     def get_start_node(self):
         return self.get_node('start-node')
@@ -190,10 +204,9 @@ def is_valid(self):
         self.is_valid_base_node()
         self.is_valid_work_flow()
 
-    @staticmethod
-    def is_valid_node_params(node: Node):
+    def is_valid_node_params(self, node: Node):
         from application.flow.step_node import get_node
-        get_node(node.type)(node, None, None)
+        get_node(node.type, self.workflow_mode)(node, None, None)
 
     def is_valid_node(self, node: Node):
         self.is_valid_node_params(node)

apps/application/flow/i_step_node.py

@@ -21,6 +21,7 @@
 from application.models import ApplicationChatUserStats
 from application.models import ChatRecord, ChatUserType
 from common.field.common import InstanceField
+from knowledge.models.knowledge_action import KnowledgeAction, State
 
 chat_cache = cache
 
@@ -78,7 +79,8 @@ def handler(self, workflow):
                                      message_tokens=message_tokens,
                                      answer_tokens=answer_tokens,
                                      answer_text_list=answer_text_list,
-                                     run_time=time.time() - workflow.context['start_time'],
+                                     run_time=time.time() - workflow.context.get('start_time') if workflow.context.get(
+                                         'start_time') is not None else 0,
                                      index=0)
 
         self.chat_info.append_chat_record(chat_record)
@@ -97,6 +99,16 @@ def handler(self, workflow):
         self.chat_info = None
 
 
+class KnowledgeWorkflowPostHandler(WorkFlowPostHandler):
+    def __init__(self, chat_info, knowledge_action_id):
+        super().__init__(chat_info)
+        self.knowledge_action_id = knowledge_action_id
+
+    def handler(self, workflow):
+        QuerySet(KnowledgeAction).filter(id=self.knowledge_action_id).update(
+            state=State.SUCCESS)
+
+
 class NodeResult:
     def __init__(self, node_variable: Dict, workflow_variable: Dict,
                  _write_context=write_context, _is_interrupt=is_interrupt):
@@ -153,6 +165,14 @@ class FlowParamsSerializer(serializers.Serializer):
     debug = serializers.BooleanField(required=True, label="是否debug")
 
 
+class KnowledgeFlowParamsSerializer(serializers.Serializer):
+    knowledge_id = serializers.UUIDField(required=True, label="知识库id")
+    workspace_id = serializers.CharField(required=True, label="工作空间id")
+    knowledge_action_id = serializers.UUIDField(required=True, label="知识库任务执行器id")
+    data_source = serializers.DictField(required=True, label="数据源")
+    knowledge_base = serializers.DictField(required=False, label="知识库设置")
+
+
 class INode:
     view_type = 'many_view'
 
@@ -165,7 +185,8 @@ def get_answer_list(self) -> List[Answer] | None:
             return None
         reasoning_content_enable = self.context.get('model_setting', {}).get('reasoning_content_enable', False)
         return [
-            Answer(self.answer_text, self.view_type, self.runtime_node_id, self.workflow_params['chat_record_id'], {},
+            Answer(self.answer_text, self.view_type, self.runtime_node_id, self.workflow_params.get('chat_record_id'),
+                   {},
                    self.runtime_node_id, self.context.get('reasoning_content', '') if reasoning_content_enable else '')]
 
     def __init__(self, node, workflow_params, workflow_manage, up_node_id_list=None,
@@ -222,13 +243,14 @@ def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
         pass
 
     def get_flow_params_serializer_class(self) -> Type[serializers.Serializer]:
-        return FlowParamsSerializer
+        return self.workflow_manage.get_params_serializer_class()
 
     def get_write_error_context(self, e):
         self.status = 500
         self.answer_text = str(e)
         self.err_message = str(e)
-        self.context['run_time'] = time.time() - self.context['start_time']
+        current_time = time.time()
+        self.context['run_time'] = current_time - (self.context.get('start_time') or current_time)
 
         def write_error_context(answer, status=200):
             pass

apps/application/flow/knowledge_loop_workflow_manage.py

@@ -0,0 +1,15 @@
+# coding=utf-8
+"""
+    @project: maxkb
+    @Author：虎
+    @file： workflow_manage.py
+    @date：2024/1/9 17:40
+    @desc:
+"""
+from application.flow.i_step_node import KnowledgeFlowParamsSerializer
+from application.flow.loop_workflow_manage import LoopWorkflowManage
+
+
+class KnowledgeLoopWorkflowManage(LoopWorkflowManage):
+    def get_params_serializer_class(self):
+        return KnowledgeFlowParamsSerializer

apps/application/flow/knowledge_workflow_manage.py

@@ -0,0 +1,102 @@
+# coding=utf-8
+"""
+    @project: MaxKB
+    @Author：虎虎
+    @file： Knowledge_workflow_manage.py
+    @date：2025/11/13 19:02
+    @desc:
+"""
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+
+from django.db.models import QuerySet
+from django.utils.translation import get_language
+
+from application.flow.common import Workflow
+from application.flow.i_step_node import WorkFlowPostHandler, KnowledgeFlowParamsSerializer
+from application.flow.workflow_manage import WorkflowManage
+from common.handle.base_to_response import BaseToResponse
+from common.handle.impl.response.system_to_response import SystemToResponse
+from knowledge.models.knowledge_action import KnowledgeAction, State
+
+executor = ThreadPoolExecutor(max_workers=200)
+
+
+class KnowledgeWorkflowManage(WorkflowManage):
+
+    def __init__(self, flow: Workflow,
+                 params,
+                 work_flow_post_handler: WorkFlowPostHandler,
+                 base_to_response: BaseToResponse = SystemToResponse(),
+                 start_node_id=None,
+                 start_node_data=None, chat_record=None, child_node=None):
+        super().__init__(flow, params, work_flow_post_handler, base_to_response, None, None, None,
+                         None,
+                         None, None, start_node_id, start_node_data, chat_record, child_node)
+
+    def get_params_serializer_class(self):
+        return KnowledgeFlowParamsSerializer
+
+    def get_start_node(self):
+        start_node_list = [node for node in self.flow.nodes if
+                           self.params.get('data_source', {}).get('node_id') == node.id]
+        return start_node_list[0]
+
+    def run(self):
+        executor.submit(self._run)
+
+    def _run(self):
+        QuerySet(KnowledgeAction).filter(id=self.params.get('knowledge_action_id')).update(
+            state=State.STARTED)
+        language = get_language()
+        self.run_chain_async(self.start_node, None, language)
+        while self.is_run():
+            pass
+        self.work_flow_post_handler.handler(self)
+
+    @staticmethod
+    def get_node_details(current_node, node, index):
+        if current_node == node:
+            return {
+                'name': node.node.properties.get('stepName'),
+                "index": index,
+                'run_time': 0,
+                'type': node.type,
+                'status': 202,
+                'err_message': ""
+            }
+
+        return node.get_details(index)
+
+    def run_chain(self, current_node, node_result_future=None):
+        QuerySet(KnowledgeAction).filter(id=self.params.get('knowledge_action_id')).update(
+            details=self.get_runtime_details(lambda node, index: self.get_node_details(current_node, node, index)))
+        if node_result_future is None:
+            node_result_future = self.run_node_future(current_node)
+        try:
+            result = self.hand_node_result(current_node, node_result_future)
+            return result
+        except Exception as e:
+            traceback.print_exc()
+        return None
+
+    def hand_node_result(self, current_node, node_result_future):
+        try:
+            current_result = node_result_future.result()
+            result = current_result.write_context(current_node, self)
+            if result is not None:
+                # 阻塞获取结果
+                list(result)
+            return current_result
+        except Exception as e:
+            traceback.print_exc()
+            self.status = 500
+            current_node.get_write_error_context(e)
+            self.answer += str(e)
+            QuerySet(KnowledgeAction).filter(id=self.params.get('knowledge_action_id')).update(
+                details=self.get_runtime_details(),
+                state=State.FAILURE)
+        finally:
+            current_node.node_chunk.end()
+            QuerySet(KnowledgeAction).filter(id=self.params.get('knowledge_action_id')).update(
+                details=self.get_runtime_details())

apps/application/flow/loop_workflow_manage.py

@@ -105,18 +105,18 @@ def get_node_cls_by_id(self, node_id, up_node_id_list=None,
                            get_node_params=lambda node: node.properties.get('node_data')):
         for node in self.flow.nodes:
             if node.id == node_id:
-                node_instance = get_node(node.type)(node,
-                                                    self.params, self, up_node_id_list,
-                                                    get_node_params,
-                                                    salt=self.get_index())
+                node_instance = get_node(node.type, self.flow.workflow_mode)(node,
+                                                                             self.params, self, up_node_id_list,
+                                                                             get_node_params,
+                                                                             salt=self.get_index())
                 return node_instance
         return None
 
     def stream(self):
         close_old_connections()
         language = get_language()
         self.run_chain_async(self.start_node, None, language)
-        return self.await_result()
+        return self.await_result(is_cleanup=False)
 
     def get_index(self):
         return self.loop_params.get('index')

apps/application/flow/step_node/__init__.py

@@ -9,13 +9,16 @@
 from .ai_chat_step_node import *
 from .application_node import BaseApplicationNode
 from .condition_node import *
+from .data_source_local_node.impl.base_data_source_local_node import BaseDataSourceLocalNode
+from .data_source_web_node.impl.base_data_source_web_node import BaseDataSourceWebNode
 from .direct_reply_node import *
 from .document_extract_node import *
 from .form_node import *
 from .image_generate_step_node import *
 from .image_to_video_step_node import BaseImageToVideoNode
 from .image_understand_step_node import *
 from .intent_node import *
+from .knowledge_write_node.impl.base_knowledge_write_node import BaseKnowledgeWriteNode
 from .loop_break_node import BaseLoopBreakNode
 from .loop_continue_node import BaseLoopContinueNode
 from .loop_node import *
@@ -36,6 +39,7 @@
 from .variable_assign_node import BaseVariableAssignNode
 from .variable_splitting_node import BaseVariableSplittingNode
 from .video_understand_step_node import BaseVideoUnderstandNode
+from .document_split_node import BaseDocumentSplitNode
 
 node_list = [BaseStartStepNode, BaseChatNode, BaseSearchKnowledgeNode, BaseSearchDocumentNode, BaseQuestionNode,
              BaseConditionNode, BaseReplyNode,
@@ -46,11 +50,11 @@
              BaseVideoUnderstandNode,
              BaseIntentNode, BaseLoopNode, BaseLoopStartStepNode,
              BaseLoopContinueNode,
-             BaseLoopBreakNode, BaseVariableSplittingNode, BaseParameterExtractionNode, BaseVariableAggregationNode]
+             BaseLoopBreakNode, BaseVariableSplittingNode, BaseParameterExtractionNode, BaseVariableAggregationNode,
+             BaseDataSourceLocalNode, BaseDataSourceWebNode, BaseKnowledgeWriteNode, BaseDocumentSplitNode]
 
+node_map = {n.type: {w: n for w in n.support} for n in node_list}
 
-def get_node(node_type):
-    find_list = [node for node in node_list if node.type == node_type]
-    if len(find_list) > 0:
-        return find_list[0]
-    return None
+
+def get_node(node_type, workflow_model):
+    return node_map.get(node_type).get(workflow_model)

apps/application/flow/step_node/ai_chat_step_node/i_chat_node.py

@@ -11,6 +11,7 @@
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
 
+from application.flow.common import WorkflowMode
 from application.flow.i_step_node import INode, NodeResult
 
 
@@ -34,22 +35,31 @@ class ChatNodeSerializer(serializers.Serializer):
     mcp_enable = serializers.BooleanField(required=False, label=_("Whether to enable MCP"))
     mcp_servers = serializers.JSONField(required=False, label=_("MCP Server"))
     mcp_tool_id = serializers.CharField(required=False, allow_blank=True, allow_null=True, label=_("MCP Tool ID"))
-    mcp_tool_ids = serializers.ListField(child=serializers.UUIDField(), required=False, allow_empty=True, label=_("MCP Tool IDs"), )
+    mcp_tool_ids = serializers.ListField(child=serializers.UUIDField(), required=False, allow_empty=True,
+                                         label=_("MCP Tool IDs"), )
     mcp_source = serializers.CharField(required=False, allow_blank=True, allow_null=True, label=_("MCP Source"))
 
     tool_enable = serializers.BooleanField(required=False, default=False, label=_("Whether to enable tools"))
     tool_ids = serializers.ListField(child=serializers.UUIDField(), required=False, allow_empty=True,
                                      label=_("Tool IDs"), )
     mcp_output_enable = serializers.BooleanField(required=False, default=True, label=_("Whether to enable MCP output"))
 
+
 class IChatNode(INode):
     type = 'ai-chat-node'
+    support = [WorkflowMode.APPLICATION, WorkflowMode.APPLICATION_LOOP, WorkflowMode.KNOWLEDGE_LOOP,
+               WorkflowMode.KNOWLEDGE]
 
     def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
         return ChatNodeSerializer
 
     def _run(self):
-        return self.execute(**self.node_params_serializer.data, **self.flow_params_serializer.data)
+        if [WorkflowMode.KNOWLEDGE, WorkflowMode.KNOWLEDGE_LOOP].__contains__(
+                self.workflow_manage.flow.workflow_mode):
+            return self.execute(**self.node_params_serializer.data, **self.flow_params_serializer.data,
+                                **{'history_chat_record': [], 'stream': True, 'chat_id': None, 'chat_record_id': None})
+        else:
+            return self.execute(**self.node_params_serializer.data, **self.flow_params_serializer.data)
 
     def execute(self, model_id, system, prompt, dialogue_number, history_chat_record, stream, chat_id,
                 chat_record_id,

apps/application/flow/step_node/application_node/i_application_node.py

@@ -3,6 +3,7 @@
 
 from rest_framework import serializers
 
+from application.flow.common import WorkflowMode
 from application.flow.i_step_node import INode, NodeResult
 
 from django.utils.translation import gettext_lazy as _
@@ -25,6 +26,7 @@ class ApplicationNodeSerializer(serializers.Serializer):
 
 class IApplicationNode(INode):
     type = 'application-node'
+    support = [WorkflowMode.APPLICATION, WorkflowMode.APPLICATION_LOOP]
 
     def get_node_params_serializer_class(self) -> Type[serializers.Serializer]:
         return ApplicationNodeSerializer