fix: Loop nodes cannot be used in the knowledge base workflow

Author: shaohuzhang1

apps/application/flow/i_step_node.py

@@ -185,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,

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

@@ -8,13 +8,12 @@
 """
 import traceback
 from concurrent.futures import ThreadPoolExecutor
-from typing import List
 
 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, INode
+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

apps/application/flow/loop_workflow_manage.py

@@ -116,7 +116,7 @@ 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/loop_node/i_loop_node.py

@@ -54,5 +54,5 @@ def _run(self):
                 array[1:])
         return self.execute(**{**self.node_params_serializer.data, "array": array}, **self.flow_params_serializer.data)
 
-    def execute(self, loop_type, array, number, loop_body, stream, **kwargs) -> NodeResult:
+    def execute(self, loop_type, array, number, loop_body, **kwargs) -> NodeResult:
         pass

apps/application/flow/step_node/loop_node/impl/base_loop_node.py

@@ -11,7 +11,7 @@
 
 from django.utils.translation import gettext as _
 
-from application.flow.common import Answer
+from application.flow.common import Answer, WorkflowMode
 from application.flow.i_step_node import NodeResult, WorkFlowPostHandler, INode
 from application.flow.step_node.loop_node.i_loop_node import ILoopNode
 from application.flow.tools import Reasoning
@@ -197,6 +197,7 @@ def loop(workflow_manage_new_instance, node: INode, generate_loop):
         insert_or_replace(loop_node_data, index, instance.get_runtime_details())
         insert_or_replace(loop_answer_data, index,
                           get_answer_list(instance, child_node_node_dict, node.runtime_node_id))
+        instance._cleanup()
         if break_outer:
             break
     node.context['is_interrupt_exec'] = is_interrupt_exec
@@ -206,7 +207,7 @@ def loop(workflow_manage_new_instance, node: INode, generate_loop):
     node.context["item"] = current_index
 
 
-def get_write_context(loop_type, array, number, loop_body, stream):
+def get_write_context(loop_type, array, number, loop_body):
     def inner_write_context(node_variable: Dict, workflow_variable: Dict, node: INode, workflow):
         if loop_type == 'ARRAY':
             return loop(node_variable['workflow_manage_new_instance'], node, generate_loop_array(array))
@@ -248,27 +249,31 @@ def get_answer_list(self) -> List[Answer] | None:
     def get_loop_context(self):
         return self.context
 
-    def execute(self, loop_type, array, number, loop_body, stream, **kwargs) -> NodeResult:
+    def execute(self, loop_type, array, number, loop_body, **kwargs) -> NodeResult:
         from application.flow.loop_workflow_manage import LoopWorkflowManage, Workflow
+        from application.flow.knowledge_loop_workflow_manage import KnowledgeLoopWorkflowManage
         def workflow_manage_new_instance(loop_data, global_data, start_node_id=None,
                                          start_node_data=None, chat_record=None, child_node=None):
-            workflow_manage = LoopWorkflowManage(Workflow.new_instance(loop_body), self.workflow_manage.params,
-                                                 LoopWorkFlowPostHandler(
-                                                     self.workflow_manage.work_flow_post_handler.chat_info),
-                                                 self.workflow_manage,
-                                                 loop_data,
-                                                 self.get_loop_context,
-                                                 base_to_response=LoopToResponse(),
-                                                 start_node_id=start_node_id,
-                                                 start_node_data=start_node_data,
-                                                 chat_record=chat_record,
-                                                 child_node=child_node
-                                                 )
+            workflow_mode = WorkflowMode.KNOWLEDGE_LOOP if WorkflowMode.KNOWLEDGE == self.workflow_manage.flow.workflow_mode else WorkflowMode.APPLICATION_LOOP
+            c = KnowledgeLoopWorkflowManage if workflow_mode == WorkflowMode.KNOWLEDGE_LOOP else LoopWorkflowManage
+            workflow_manage = c(Workflow.new_instance(loop_body, workflow_mode),
+                                self.workflow_manage.params,
+                                LoopWorkFlowPostHandler(
+                                    self.workflow_manage.work_flow_post_handler.chat_info),
+                                self.workflow_manage,
+                                loop_data,
+                                self.get_loop_context,
+                                base_to_response=LoopToResponse(),
+                                start_node_id=start_node_id,
+                                start_node_data=start_node_data,
+                                chat_record=chat_record,
+                                child_node=child_node
+                                )
 
             return workflow_manage
 
         return NodeResult({'workflow_manage_new_instance': workflow_manage_new_instance}, {},
-                          _write_context=get_write_context(loop_type, array, number, loop_body, stream),
+                          _write_context=get_write_context(loop_type, array, number, loop_body),
                           _is_interrupt=_is_interrupt_exec)
 
     def get_loop_context_data(self):

apps/application/flow/step_node/loop_start_node/impl/base_start_node.py

@@ -10,6 +10,7 @@
 
 from rest_framework import serializers
 
+from application.flow.common import WorkflowMode
 from application.flow.i_step_node import NodeResult
 from application.flow.step_node.loop_start_node.i_loop_start_node import ILoopStarNode
 
@@ -31,7 +32,8 @@ def execute(self, **kwargs) -> NodeResult:
             'index': loop_params.get("index"),
             'item': loop_params.get("item")
         }
-        self.workflow_manage.chat_context = self.workflow_manage.get_chat_info().get_chat_variable()
+        if WorkflowMode.APPLICATION_LOOP == self.workflow_manage.flow.workflow_mode:
+            self.workflow_manage.chat_context = self.workflow_manage.get_chat_info().get_chat_variable()
         return NodeResult(node_variable, {})
 
     def get_details(self, index: int, **kwargs):

apps/application/flow/workflow_manage.py

@@ -316,7 +316,7 @@ def is_run(self, timeout=0.5):
         except Exception as e:
             return True
 
-    def await_result(self):
+    def await_result(self, is_cleanup=True):
         try:
             while self.is_run():
                 while True:
@@ -344,7 +344,8 @@ def await_result(self):
                                                                  '',
                                                                  [],
                                                                  '', True, message_tokens, answer_tokens, {})
-            self._cleanup()
+            if is_cleanup:
+                self._cleanup()
 
     def run_chain_async(self, current_node, node_result_future, language='zh'):
         future = executor.submit(self.run_chain_manage, current_node, node_result_future, language)

ui/src/components/workflow-dropdown-menu/index.vue

@@ -20,7 +20,9 @@ defineProps({
 })
 const kw: any = {
   [WorkflowMode.Application]: ApplicationDropdownMenu,
+  [WorkflowMode.ApplicationLoop]: ApplicationDropdownMenu,
   [WorkflowMode.Knowledge]: KnowledgeDropdownMenu,
+  [WorkflowMode.KnowledgeLoop]: KnowledgeDropdownMenu,
 }
 </script>
 <style lang="scss">

ui/src/workflow/common/NodeContainer.vue

@@ -46,10 +46,7 @@
                   <p class="mt-8 lighter">
                     <span>{{ $t('views.workflow.condition.front') }}</span>
                     <el-select v-model="condition" size="small" style="width: 60px; margin: 0 8px">
-                      <el-option
-                        :label="$t('views.workflow.condition.AND')"
-                        value="AND"
-                      />
+                      <el-option :label="$t('views.workflow.condition.AND')" value="AND" />
                       <el-option :label="$t('views.workflow.condition.OR')" value="OR" />
                     </el-select>
                     <span>{{ $t('views.workflow.condition.text') }}</span>