fix: Optimize knowledge base; AI-driven generation of tools and workflows

apps/application/tests/test_workflow_loop_variables.py

@@ -0,0 +1,146 @@
+import threading
+
+from django.test import SimpleTestCase
+
+from application.workflow.common import WorkflowType, new_instance
+from application.workflow.nodes import get_node_class
+from application.workflow.nodes.loop_node.loop_node import LoopNodeSerializer
+from application.workflow.workflow_manage import CallBack, WorkflowManage
+from common.exception.app_exception import AppApiException
+
+
+def _node(node_id, node_type, node_data=None, **properties):
+    return {
+        "id": node_id,
+        "type": node_type,
+        "x": 0,
+        "y": 0,
+        "properties": {"stepName": node_id, "node_data": node_data or {}, **properties},
+    }
+
+
+def _edge(source, target):
+    return {
+        "id": f"{source}-{target}",
+        "type": "app-edge",
+        "sourceNodeId": source,
+        "targetNodeId": target,
+        "sourceAnchorId": f"{source}_right",
+        "targetAnchorId": f"{target}_left",
+    }
+
+
+def _loop_start(*variables):
+    return _node(
+        "loop-start-node",
+        "loop-start-node",
+        stepName="循环开始",
+        config={"fields": [{"label": "下标", "value": "index"}, {"label": "循环元素", "value": "item"}]},
+        loop_input_field_list=[{"field": field, "label": field} for field in variables],
+    )
+
+
+def _assign(node_id, *variables):
+    return _node(
+        node_id,
+        "variable-assign-node",
+        {
+            "variable_list": [
+                {"name": ".".join(fields), "fields": fields, "source": "custom", "type": "string", "value": value}
+                for fields, value in variables
+            ]
+        },
+    )
+
+
+def _loop(node_id, number, body_nodes, body_edges):
+    return _node(
+        node_id,
+        "loop-node",
+        {"loop_type": "NUMBER", "number": number, "loop_body": {"nodes": body_nodes, "edges": body_edges}},
+    )
+
+
+def _run(loop_node):
+    """以循环节点作为起始节点跑完一个根工作流,返回根 WorkflowManage"""
+    workflow = new_instance({"nodes": [loop_node], "edges": []})
+    finished = threading.Event()
+    errors = []
+
+    def on_complete(_manage, error):
+        errors.append(error)
+        finished.set()
+
+    def get_start_node(wf, manage):
+        node = wf.get_node(loop_node["id"])
+        node_class = get_node_class(node.type, WorkflowType.APPLICATION)
+        return node_class(node, manage, lambda n: n.properties.get("node_data", {}))
+
+    def on_next(wf_manage, _content):
+        # 知识库工作流每收到一块输出就取一次详情,第一轮循环尚未结束时也必须可用
+        wf_manage.get_details()
+
+    manage = WorkflowManage(workflow, {}, WorkflowType.APPLICATION, CallBack(on_next, on_complete), get_start_node)
+    manage.run()
+    assert finished.wait(10), "workflow did not finish"
+    assert errors == [None], errors
+    return manage
+
+
+class LoopVariableTests(SimpleTestCase):
+    def test_loop_variable_is_shared_across_iterations_and_readable_after_the_loop(self):
+        loop = _loop(
+            "loop1",
+            3,
+            [
+                _loop_start("total"),
+                _assign(
+                    "assign",
+                    (["loop", "total"], "{{loop.total}}{{循环开始.index}}"),
+                    (["global", "last"], "{{循环变量.total}}"),
+                ),
+            ],
+            [_edge("loop-start-node", "assign")],
+        )
+
+        manage = _run(loop)
+
+        # 循环变量存放在循环节点的上下文里,下游用 [循环节点ID, 变量名] 引用
+        self.assertEqual(manage.get_reference_field("loop1", ["total"]), "012")
+        # 循环体内给全局变量赋值,写到根工作流
+        self.assertEqual(manage.get_reference_field("global", ["last"]), "012")
+        self.assertNotIn("loop", manage.context["loop1"]["last_context"])
+        self.assertNotIn("global", manage.context["loop1"]["last_context"])
+
+    def test_nested_loop_writes_reach_the_owning_loop_and_the_root_workflow(self):
+        inner = _loop(
+            "inner",
+            2,
+            [
+                _loop_start("cnt"),
+                _assign(
+                    "assign",
+                    (["loop", "acc"], "{{loop.acc}}x"),
+                    (["loop", "cnt"], "{{loop.cnt}}y"),
+                    (["global", "g"], "{{loop.acc}}|{{loop.cnt}}"),
+                ),
+            ],
+            [_edge("loop-start-node", "assign")],
+        )
+        outer = _loop("outer", 2, [_loop_start("acc"), inner], [_edge("loop-start-node", "inner")])
+
+        manage = _run(outer)
+
+        # 内层循环修改外层循环声明的变量
+        self.assertEqual(manage.get_reference_field("outer", ["acc"]), "xxxx")
+        # 内层循环自己的变量存放在外层循环体的上下文里,且每次进入内层循环都从空开始
+        self.assertEqual(manage.context["outer"]["last_context"]["inner"]["cnt"], "yy")
+        # 两层嵌套里给全局变量赋值,依然写到根工作流
+        self.assertEqual(manage.context["global"], {"g": "xxxx|yy"})
+        self.assertNotIn("global", manage.context["outer"]["last_context"])
+
+    def test_loop_variable_cannot_use_a_name_reserved_by_the_loop_node(self):
+        data = {"loop_type": "NUMBER", "number": 1, "loop_body": {"nodes": [_loop_start("last_context")], "edges": []}}
+
+        with self.assertRaises(AppApiException):
+            LoopNodeSerializer(data=data).is_valid(raise_exception=True)

apps/application/workflow/loop_workflow_manage.py

@@ -7,15 +7,33 @@
 @desc:
 """
 
-from typing import Dict, Callable
+from typing import Dict, Callable, List
 
 from application.workflow.common import Workflow, WorkflowType
 from application.workflow.i_node import INode
 from application.workflow.workflow_manage import WorkflowManage, CallBack
-from common.utils.prompt_template import render_prompt
+
+
+def get_loop_variable_fields(loop_start_properties) -> List[str]:
+    """
+    从循环开始节点的属性里取出声明的循环变量名
+    @param loop_start_properties: loop-start-node 的 properties
+    @return: 循环变量名列表
+    """
+    field_list = (loop_start_properties or {}).get("loop_input_field_list") or []
+    fields = [item.get("field") or item.get("variable") for item in field_list if isinstance(item, dict)]
+    return [field for field in fields if field]
 
 
 class LoopWorkFlowManage(WorkflowManage):
+    """
+    循环体工作流。
+
+    循环变量不存放在循环体自己的上下文里,而是存放在父工作流中“循环节点”的上下文下
+    (parent.context[loop_node_id][变量名]):各轮共享、循环结束后下游可直接用 [循环节点ID, 变量名] 引用。
+    循环体内的 ["loop", 变量名] 只是它的别名。
+    """
+
     def __init__(
         self,
         workflow: Workflow,
@@ -24,8 +42,12 @@ def __init__(
         call_back: CallBack,
         get_start_node: Callable[[Workflow, WorkflowManage], INode],
         parent_workflow_manage: WorkflowManage,
+        loop_node_id: str = None,
     ):
         self.parent_workflow_manage = parent_workflow_manage
+        self.loop_node_id = loop_node_id
+        loop_start_node = workflow.get_node("loop-start-node")
+        self.loop_variable_fields = get_loop_variable_fields(loop_start_node.properties if loop_start_node else None)
         super().__init__(workflow, parameters, workflow_type, call_back, get_start_node)
 
     def get_parameters(self):
@@ -34,11 +56,48 @@ def get_parameters(self):
     def get_parent_context(self, node_id, key):
         return self.parent_workflow_manage.get_context(node_id, key)
 
-    def generate_prompt(self, prompt):
+    def _loop_variable_owner(self, field):
+        """
+        找到循环变量所属的循环:优先是声明了该变量的最近一层循环,都没有声明则归当前循环
+        @param field: 循环变量名
+        @return: 所属循环的 LoopWorkFlowManage
+        """
+        manage = self
+        while field not in manage.loop_variable_fields and isinstance(
+            manage.parent_workflow_manage, LoopWorkFlowManage
+        ):
+            manage = manage.parent_workflow_manage
+        return manage if field in manage.loop_variable_fields else self
+
+    def write_context(self, node_id, key, value, append=False):
+        """
+        写入上下文:节点属于当前循环体就写自己,否则交给父工作流(父工作流是循环则继续向上)。
+        循环变量(node_id 为 loop)写到所属循环的循环节点上下文里。
+        """
+        if node_id == "loop" and self.loop_node_id:
+            owner = self._loop_variable_owner(key)
+            owner.parent_workflow_manage.write_context(owner.loop_node_id, key, value, append)
+            return
+        if self.workflow.get_node(node_id) is not None:
+            super().write_context(node_id, key, value, append)
+            return
+        self.parent_workflow_manage.write_context(node_id, key, value, append)
+
+    def reset_prompt(self, prompt):
         input_template = self.workflow.reset_prompt(prompt)
-        input_template = self.parent_workflow_manage.workflow.reset_prompt(input_template)
-        context = {**self.context, **self.parent_workflow_manage.context}
-        return render_prompt(input_template, context)
+        if self.loop_node_id:
+            # 长的变量名先替换,避免 loop.ab 被 loop.a 截断
+            for field in sorted(self.loop_variable_fields, key=len, reverse=True):
+                system_variable = f"(context.get('{self.loop_node_id}') or {{}}).get('{field}','')"
+                input_template = input_template.replace(f"循环变量.{field}", system_variable)
+                input_template = input_template.replace(f"loop.{field}", system_variable)
+        return self.parent_workflow_manage.reset_prompt(input_template)
+
+    def get_merged_context(self):
+        """
+        祖先工作流上下文 + 自己的上下文,同名时以自己为准
+        """
+        return {**self.parent_workflow_manage.get_merged_context(), **self.context}
 
     def get_reference_field(self, node_id, fields):
         """
@@ -47,6 +106,10 @@ def get_reference_field(self, node_id, fields):
         @param fields:   字段
         @return: 引用数据
         """
+        if node_id == "loop" and self.loop_node_id and fields:
+            owner = self._loop_variable_owner(fields[0])
+            return owner.parent_workflow_manage.get_reference_field(owner.loop_node_id, fields)
+
         # 先从当前工作流获取
         result = super().get_reference_field(node_id, fields)
         if result is not None:
@@ -57,7 +120,15 @@ def get_reference_field(self, node_id, fields):
 
     @classmethod
     def from_context(
-        cls, get_context, workflow, parameters, workflow_type, call_back, get_start_node, parent_workflow_manage=None
+        cls,
+        get_context,
+        workflow,
+        parameters,
+        workflow_type,
+        call_back,
+        get_start_node,
+        parent_workflow_manage=None,
+        loop_node_id=None,
     ):
         try:
             context = get_context()
@@ -69,6 +140,7 @@ def from_context(
                 call_back=call_back,
                 get_start_node=get_start_node,
                 parent_workflow_manage=parent_workflow_manage,
+                loop_node_id=loop_node_id,
             )
             if context:
                 instance.context = context

apps/application/workflow/nodes/loop_node/loop_node.py

@@ -14,11 +14,24 @@
 
 from application.workflow.common import WorkflowType, new_instance
 from application.workflow.i_node import INode, Signal
+from application.workflow.loop_workflow_manage import LoopWorkFlowManage, get_loop_variable_fields
 from application.workflow.message.struct.content import Position
 from application.workflow.status import Status
 from common.exception.app_exception import AppApiException
 
 MAX_LOOP_COUNT = 500
+# 本节点自身写入上下文的 key;循环变量也存放在本节点的上下文里,不能与它们重名
+LOOP_RESERVED_FIELDS = ("index", "item", "last_context", "error_message", "exception_message")
+
+
+def _get_loop_variable_fields(loop_body):
+    """
+    获取循环体声明的循环变量名(声明在循环体的 loop-start-node 上)
+    """
+    for node in (loop_body or {}).get("nodes") or []:
+        if node.get("id") == "loop-start-node":
+            return get_loop_variable_fields(node.get("properties"))
+    return []
 
 
 class LoopNodeSerializer(serializers.Serializer):
@@ -40,6 +53,10 @@ def is_valid(self, *, raise_exception=False):
             if number is None:
                 message = _("{field}, this field is required.", field="number")
                 raise AppApiException(500, message)
+        # 循环变量存放在循环节点的上下文里,不能和循环节点自身的字段重名
+        for field in _get_loop_variable_fields(self.data.get("loop_body")):
+            if field in LOOP_RESERVED_FIELDS:
+                raise AppApiException(500, _("Loop variable name is reserved: {field}").format(field=field))
 
 
 def _generate_loop_number(number, start_index=0):
@@ -92,6 +109,12 @@ def execute(self):
 
         self.data["params"] = {"loop_type": loop_type, "array": array, "number": number}
 
+        # 循环变量存放在本节点的上下文里;全新执行(非断点续跑)时清空,保证初始为空
+        if position.get("id") != self.node.id:
+            node_context = self.workflow_manage.context.setdefault(self.node.id, {})
+            for field in _get_loop_variable_fields(loop_body):
+                node_context.pop(field, None)
+
         # 根据 start_index 构建迭代器
         if loop_type == "ARRAY":
             iterator = _generate_loop_array(array, start_index=start_index)
@@ -163,7 +186,6 @@ def on_complete(wf_manage, error):
             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)
@@ -210,6 +232,7 @@ def get_context():
             call_back=call_back,
             get_start_node=get_start_node_fn,
             parent_workflow_manage=self.workflow_manage,
+            loop_node_id=self.node.id,
             get_context=get_context,
         )
         if loop_manage:
@@ -243,7 +266,7 @@ def get_details(self, index: int = 0, position: dict = None, old_details: dict =
                 if position.get("children").get("id") == value.get("node_id"):
                     position_index = index
 
-        for index, _loop_details in enumerate(self.data.get("loop_details_list")):
+        for index, _loop_details in enumerate(self.data.get("loop_details_list") or []):
             if position and index == 0:
                 for inner_index, item in enumerate(_loop_details):
                     if position is not None and inner_index == 0 and index == 0:

apps/application/workflow/nodes/loop_start_node/loop_start_node.py

@@ -23,9 +23,6 @@ class LoopStartNode(INode):
     type = "loop-start-node"
 
     def execute(self):
-        loop = self.workflow_manage.context.get("loop")
-        if loop is None:
-            self.write_context("loop", {})
         parameters = self.workflow_manage.get_parameters()
         if parameters is not None:
             index = parameters.get("index", 0)