Underground-Digital/Workflow-Engine

1import time2import uuid3 4from core.app.entities.app_invoke_entities import InvokeFrom5from core.workflow.entities.node_entities import NodeRunResult6from core.workflow.entities.variable_pool import VariablePool7from core.workflow.enums import SystemVariableKey8from core.workflow.graph_engine.entities.graph import Graph9from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams10from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState11from core.workflow.nodes.tool.tool_node import ToolNode12from models.enums import UserFrom13from models.workflow import WorkflowNodeExecutionStatus, WorkflowType14 15 16def init_tool_node(config: dict):17    graph_config = {18        "edges": [19            {20                "id": "start-source-next-target",21                "source": "start",22                "target": "1",23            },24        ],25        "nodes": [{"data": {"type": "start"}, "id": "start"}, config],26    }27 28    graph = Graph.init(graph_config=graph_config)29 30    init_params = GraphInitParams(31        tenant_id="1",32        app_id="1",33        workflow_type=WorkflowType.WORKFLOW,34        workflow_id="1",35        graph_config=graph_config,36        user_id="1",37        user_from=UserFrom.ACCOUNT,38        invoke_from=InvokeFrom.DEBUGGER,39        call_depth=0,40    )41 42    # construct variable pool43    variable_pool = VariablePool(44        system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"},45        user_inputs={},46        environment_variables=[],47        conversation_variables=[],48    )49 50    return ToolNode(51        id=str(uuid.uuid4()),52        graph_init_params=init_params,53        graph=graph,54        graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()),55        config=config,56    )57 58 59def test_tool_variable_invoke():60    node = init_tool_node(61        config={62            "id": "1",63            "data": {64                "title": "a",65                "desc": "a",66                "provider_id": "maths",67                "provider_type": "builtin",68                "provider_name": "maths",69                "tool_name": "eval_expression",70                "tool_label": "eval_expression",71                "tool_configurations": {},72                "tool_parameters": {73                    "expression": {74                        "type": "variable",75                        "value": ["1", "123", "args1"],76                    }77                },78            },79        }80    )81 82    node.graph_runtime_state.variable_pool.add(["1", "123", "args1"], "1+1")83 84    # execute node85    result = node._run()86    assert isinstance(result, NodeRunResult)87    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED88    assert result.outputs is not None89    assert "2" in result.outputs["text"]90    assert result.outputs["files"] == []91 92 93def test_tool_mixed_invoke():94    node = init_tool_node(95        config={96            "id": "1",97            "data": {98                "title": "a",99                "desc": "a",100                "provider_id": "maths",101                "provider_type": "builtin",102                "provider_name": "maths",103                "tool_name": "eval_expression",104                "tool_label": "eval_expression",105                "tool_configurations": {},106                "tool_parameters": {107                    "expression": {108                        "type": "mixed",109                        "value": "{{#1.args1#}}",110                    }111                },112            },113        }114    )115 116    node.graph_runtime_state.variable_pool.add(["1", "args1"], "1+1")117 118    # execute node119    result = node._run()120    assert isinstance(result, NodeRunResult)121    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED122    assert result.outputs is not None123    assert "2" in result.outputs["text"]124    assert result.outputs["files"] == []125