diff --git a/api/core/workflow/nodes/http_request/http_request_node.py b/api/core/workflow/nodes/http_request/http_request_node.py index 6db9c43555..185f896166 100644 --- a/api/core/workflow/nodes/http_request/http_request_node.py +++ b/api/core/workflow/nodes/http_request/http_request_node.py @@ -53,7 +53,10 @@ class HttpRequestNode(BaseNode): node_data: HttpRequestNodeData = cast(HttpRequestNodeData, self.node_data) # TODO: Switch to use segment directly if node_data.authorization.config and node_data.authorization.config.api_key: - node_data.authorization.config.api_key = parser.convert_template(template=node_data.authorization.config.api_key, variable_pool=variable_pool).text + node_data.authorization.config.api_key = parser.convert_template( + template=node_data.authorization.config.api_key, + variable_pool=self.graph_runtime_state.variable_pool + ).text # init http executor http_executor = None diff --git a/api/tests/integration_tests/workflow/nodes/test_code.py b/api/tests/integration_tests/workflow/nodes/test_code.py index 8b77174472..952c90674d 100644 --- a/api/tests/integration_tests/workflow/nodes/test_code.py +++ b/api/tests/integration_tests/workflow/nodes/test_code.py @@ -1,17 +1,72 @@ +import time +import uuid from os import getenv +from typing import cast import pytest from core.app.entities.app_invoke_entities import InvokeFrom -from core.workflow.entities.node_entities import UserFrom +from core.workflow.entities.node_entities import NodeRunResult, UserFrom from core.workflow.entities.variable_pool import VariablePool +from core.workflow.enums import SystemVariableKey +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState from core.workflow.nodes.code.code_node import CodeNode -from models.workflow import WorkflowNodeExecutionStatus +from core.workflow.nodes.code.entities import CodeNodeData +from models.workflow import WorkflowNodeExecutionStatus, WorkflowType from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock CODE_MAX_STRING_LENGTH = int(getenv("CODE_MAX_STRING_LENGTH", "10000")) +def init_code_node(code_config: dict): + graph_config = { + "edges": [ + { + "id": "start-source-code-target", + "source": "start", + "target": "code", + }, + ], + "nodes": [{"data": {"type": "start"}, "id": "start"}, code_config], + } + + graph = Graph.init(graph_config=graph_config) + + init_params = GraphInitParams( + tenant_id="1", + app_id="1", + workflow_type=WorkflowType.WORKFLOW, + workflow_id="1", + graph_config=graph_config, + user_id="1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + # construct variable pool + variable_pool = VariablePool( + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + variable_pool.add(["code", "123", "args1"], 1) + variable_pool.add(["code", "123", "args2"], 2) + + node = CodeNode( + id=str(uuid.uuid4()), + graph_init_params=init_params, + graph=graph, + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), + config=code_config, + ) + + return node + + @pytest.mark.parametrize("setup_code_executor_mock", [["none"]], indirect=True) def test_execute_code(setup_code_executor_mock): code = """ @@ -22,45 +77,36 @@ def test_execute_code(setup_code_executor_mock): """ # trim first 4 spaces at the beginning of each line code = "\n".join([line[4:] for line in code.split("\n")]) - node = CodeNode( - id='test', - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.WEB_APP, - config={ - "id": "1", - "data": { - "outputs": { - "result": { - "type": "number", - }, - }, - "title": "123", - "variables": [ - { - "variable": "args1", - "value_selector": ["1", "123", "args1"], - }, - {"variable": "args2", "value_selector": ["1", "123", "args2"]}, - ], - "answer": "123", - "code_language": "python3", - "code": code, - }, - }, - ) - # construct variable pool - pool = VariablePool(system_variables={}, user_inputs={}, environment_variables=[]) - pool.add(["1", "123", "args1"], 1) - pool.add(["1", "123", "args2"], 2) + code_config = { + "id": "code", + "data": { + "outputs": { + "result": { + "type": "number", + }, + }, + "title": "123", + "variables": [ + { + "variable": "args1", + "value_selector": ["1", "123", "args1"], + }, + {"variable": "args2", "value_selector": ["1", "123", "args2"]}, + ], + "answer": "123", + "code_language": "python3", + "code": code, + }, + } + + node = init_code_node(code_config) # execute node - result = node.run(pool) + result = node._run() + assert isinstance(result, NodeRunResult) assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert result.outputs["result"] == 3 assert result.error is None @@ -75,45 +121,34 @@ def test_execute_code_output_validator(setup_code_executor_mock): """ # trim first 4 spaces at the beginning of each line code = "\n".join([line[4:] for line in code.split("\n")]) - node = CodeNode( - id='test', - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.WEB_APP, - config={ - "id": "1", - "data": { - "outputs": { - "result": { - "type": "string", - }, - }, - "title": "123", - "variables": [ - { - "variable": "args1", - "value_selector": ["1", "123", "args1"], - }, - {"variable": "args2", "value_selector": ["1", "123", "args2"]}, - ], - "answer": "123", - "code_language": "python3", - "code": code, - }, - }, - ) - # construct variable pool - pool = VariablePool(system_variables={}, user_inputs={}, environment_variables=[]) - pool.add(["1", "123", "args1"], 1) - pool.add(["1", "123", "args2"], 2) + code_config = { + "id": "code", + "data": { + "outputs": { + "result": { + "type": "string", + }, + }, + "title": "123", + "variables": [ + { + "variable": "args1", + "value_selector": ["1", "123", "args1"], + }, + {"variable": "args2", "value_selector": ["1", "123", "args2"]}, + ], + "answer": "123", + "code_language": "python3", + "code": code, + }, + } + + node = init_code_node(code_config) # execute node - result = node.run(pool) - + result = node._run() + assert isinstance(result, NodeRunResult) assert result.status == WorkflowNodeExecutionStatus.FAILED assert result.error == "Output variable `result` must be a string" @@ -129,65 +164,60 @@ def test_execute_code_output_validator_depth(): """ # trim first 4 spaces at the beginning of each line code = "\n".join([line[4:] for line in code.split("\n")]) - node = CodeNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - user_from=UserFrom.ACCOUNT, - invoke_from=InvokeFrom.WEB_APP, - config={ - "id": "1", - "data": { - "outputs": { - "string_validator": { - "type": "string", - }, - "number_validator": { - "type": "number", - }, - "number_array_validator": { - "type": "array[number]", - }, - "string_array_validator": { - "type": "array[string]", - }, - "object_validator": { - "type": "object", - "children": { - "result": { - "type": "number", - }, - "depth": { - "type": "object", - "children": { - "depth": { - "type": "object", - "children": { - "depth": { - "type": "number", - } - }, - } - }, + + code_config = { + "id": "code", + "data": { + "outputs": { + "string_validator": { + "type": "string", + }, + "number_validator": { + "type": "number", + }, + "number_array_validator": { + "type": "array[number]", + }, + "string_array_validator": { + "type": "array[string]", + }, + "object_validator": { + "type": "object", + "children": { + "result": { + "type": "number", + }, + "depth": { + "type": "object", + "children": { + "depth": { + "type": "object", + "children": { + "depth": { + "type": "number", + } + }, + } }, }, }, }, - "title": "123", - "variables": [ - { - "variable": "args1", - "value_selector": ["1", "123", "args1"], - }, - {"variable": "args2", "value_selector": ["1", "123", "args2"]}, - ], - "answer": "123", - "code_language": "python3", - "code": code, }, + "title": "123", + "variables": [ + { + "variable": "args1", + "value_selector": ["1", "123", "args1"], + }, + {"variable": "args2", "value_selector": ["1", "123", "args2"]}, + ], + "answer": "123", + "code_language": "python3", + "code": code, }, - ) + } + + node = init_code_node(code_config) # construct result result = { @@ -198,6 +228,8 @@ def test_execute_code_output_validator_depth(): "object_validator": {"result": 1, "depth": {"depth": {"depth": 1}}}, } + node.node_data = cast(CodeNodeData, node.node_data) + # validate node._transform_result(result, node.node_data.outputs) @@ -252,35 +284,30 @@ def test_execute_code_output_object_list(): """ # trim first 4 spaces at the beginning of each line code = "\n".join([line[4:] for line in code.split("\n")]) - node = CodeNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, - config={ - "id": "1", - "data": { - "outputs": { - "object_list": { - "type": "array[object]", - }, + + code_config = { + "id": "code", + "data": { + "outputs": { + "object_list": { + "type": "array[object]", }, - "title": "123", - "variables": [ - { - "variable": "args1", - "value_selector": ["1", "123", "args1"], - }, - {"variable": "args2", "value_selector": ["1", "123", "args2"]}, - ], - "answer": "123", - "code_language": "python3", - "code": code, }, + "title": "123", + "variables": [ + { + "variable": "args1", + "value_selector": ["1", "123", "args1"], + }, + {"variable": "args2", "value_selector": ["1", "123", "args2"]}, + ], + "answer": "123", + "code_language": "python3", + "code": code, }, - ) + } + + node = init_code_node(code_config) # construct result result = { @@ -297,6 +324,8 @@ def test_execute_code_output_object_list(): ] } + node.node_data = cast(CodeNodeData, node.node_data) + # validate node._transform_result(result, node.node_data.outputs) diff --git a/api/tests/integration_tests/workflow/nodes/test_http.py b/api/tests/integration_tests/workflow/nodes/test_http.py index 5a7ae2d9da..65aaa0bddd 100644 --- a/api/tests/integration_tests/workflow/nodes/test_http.py +++ b/api/tests/integration_tests/workflow/nodes/test_http.py @@ -1,3 +1,5 @@ +import time +import uuid from urllib.parse import urlencode import pytest @@ -5,27 +7,63 @@ import pytest from core.app.entities.app_invoke_entities import InvokeFrom from core.workflow.entities.node_entities import UserFrom from core.workflow.entities.variable_pool import VariablePool +from core.workflow.enums import SystemVariableKey +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState from core.workflow.nodes.http_request.http_request_node import HttpRequestNode +from models.workflow import WorkflowType from tests.integration_tests.workflow.nodes.__mock.http import setup_http_mock -BASIC_NODE_DATA = { - "tenant_id": "1", - "app_id": "1", - "workflow_id": "1", - "user_id": "1", - "user_from": UserFrom.ACCOUNT, - "invoke_from": InvokeFrom.WEB_APP, -} -# construct variable pool -pool = VariablePool(system_variables={}, user_inputs={}, environment_variables=[]) -pool.add(["a", "b123", "args1"], 1) -pool.add(["a", "b123", "args2"], 2) +def init_http_node(config: dict): + graph_config = { + "edges": [ + { + "id": "start-source-next-target", + "source": "start", + "target": "1", + }, + ], + "nodes": [{"data": {"type": "start"}, "id": "start"}, config], + } + + graph = Graph.init(graph_config=graph_config) + + init_params = GraphInitParams( + tenant_id="1", + app_id="1", + workflow_type=WorkflowType.WORKFLOW, + workflow_id="1", + graph_config=graph_config, + user_id="1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + # construct variable pool + variable_pool = VariablePool( + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + variable_pool.add(["a", "b123", "args1"], 1) + variable_pool.add(["a", "b123", "args2"], 2) + + return HttpRequestNode( + id=str(uuid.uuid4()), + graph_init_params=init_params, + graph=graph, + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), + config=config, + ) @pytest.mark.parametrize("setup_http_mock", [["none"]], indirect=True) def test_get(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -45,12 +83,11 @@ def test_get(setup_http_mock): "params": "A:b", "body": None, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) - + result = node._run() + assert result.process_data is not None data = result.process_data.get("request", "") assert "?A=b" in data @@ -59,7 +96,7 @@ def test_get(setup_http_mock): @pytest.mark.parametrize("setup_http_mock", [["none"]], indirect=True) def test_no_auth(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -75,12 +112,11 @@ def test_no_auth(setup_http_mock): "params": "A:b", "body": None, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) - + result = node._run() + assert result.process_data is not None data = result.process_data.get("request", "") assert "?A=b" in data @@ -89,7 +125,7 @@ def test_no_auth(setup_http_mock): @pytest.mark.parametrize("setup_http_mock", [["none"]], indirect=True) def test_custom_authorization_header(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -109,12 +145,11 @@ def test_custom_authorization_header(setup_http_mock): "params": "A:b", "body": None, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) - + result = node._run() + assert result.process_data is not None data = result.process_data.get("request", "") assert "?A=b" in data @@ -123,7 +158,7 @@ def test_custom_authorization_header(setup_http_mock): @pytest.mark.parametrize("setup_http_mock", [["none"]], indirect=True) def test_template(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -143,11 +178,11 @@ def test_template(setup_http_mock): "params": "A:b\nTemplate:{{#a.b123.args2#}}", "body": None, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) + result = node._run() + assert result.process_data is not None data = result.process_data.get("request", "") assert "?A=b" in data @@ -158,7 +193,7 @@ def test_template(setup_http_mock): @pytest.mark.parametrize("setup_http_mock", [["none"]], indirect=True) def test_json(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -178,11 +213,11 @@ def test_json(setup_http_mock): "params": "A:b", "body": {"type": "json", "data": '{"a": "{{#a.b123.args1#}}"}'}, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) + result = node._run() + assert result.process_data is not None data = result.process_data.get("request", "") assert '{"a": "1"}' in data @@ -190,7 +225,7 @@ def test_json(setup_http_mock): def test_x_www_form_urlencoded(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -210,11 +245,11 @@ def test_x_www_form_urlencoded(setup_http_mock): "params": "A:b", "body": {"type": "x-www-form-urlencoded", "data": "a:{{#a.b123.args1#}}\nb:{{#a.b123.args2#}}"}, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) + result = node._run() + assert result.process_data is not None data = result.process_data.get("request", "") assert "a=1&b=2" in data @@ -222,7 +257,7 @@ def test_x_www_form_urlencoded(setup_http_mock): def test_form_data(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -242,11 +277,11 @@ def test_form_data(setup_http_mock): "params": "A:b", "body": {"type": "form-data", "data": "a:{{#a.b123.args1#}}\nb:{{#a.b123.args2#}}"}, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) + result = node._run() + assert result.process_data is not None data = result.process_data.get("request", "") assert 'form-data; name="a"' in data @@ -257,7 +292,7 @@ def test_form_data(setup_http_mock): def test_none_data(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -277,11 +312,11 @@ def test_none_data(setup_http_mock): "params": "A:b", "body": {"type": "none", "data": "123123123"}, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) + result = node._run() + assert result.process_data is not None data = result.process_data.get("request", "") assert "X-Header: 123" in data @@ -289,7 +324,7 @@ def test_none_data(setup_http_mock): def test_mock_404(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -305,19 +340,19 @@ def test_mock_404(setup_http_mock): "params": "", "headers": "X-Header:123", }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) + result = node._run() + assert result.outputs is not None resp = result.outputs assert 404 == resp.get("status_code") - assert "Not Found" in resp.get("body") + assert "Not Found" in resp.get("body", "") def test_multi_colons_parse(setup_http_mock): - node = HttpRequestNode( + node = init_http_node( config={ "id": "1", "data": { @@ -333,13 +368,14 @@ def test_multi_colons_parse(setup_http_mock): "headers": "Referer:http://example3.com\nRedirect:http://example4.com", "body": {"type": "form-data", "data": "Referer:http://example5.com\nRedirect:http://example6.com"}, }, - }, - **BASIC_NODE_DATA, + } ) - result = node.run(pool) + result = node._run() + assert result.process_data is not None + assert result.outputs is not None resp = result.outputs - assert urlencode({"Redirect": "http://example2.com"}) in result.process_data.get("request") - assert 'form-data; name="Redirect"\n\nhttp://example6.com' in result.process_data.get("request") - assert "http://example3.com" == resp.get("headers").get("referer") + assert urlencode({"Redirect": "http://example2.com"}) in result.process_data.get("request", "") + assert 'form-data; name="Redirect"\n\nhttp://example6.com' in result.process_data.get("request", "") + assert "http://example3.com" == resp.get("headers", {}).get("referer") diff --git a/api/tests/integration_tests/workflow/nodes/test_llm.py b/api/tests/integration_tests/workflow/nodes/test_llm.py index 371d916578..dfb43650d2 100644 --- a/api/tests/integration_tests/workflow/nodes/test_llm.py +++ b/api/tests/integration_tests/workflow/nodes/test_llm.py @@ -1,5 +1,8 @@ import json import os +import time +import uuid +from collections.abc import Generator from unittest.mock import MagicMock import pytest @@ -13,25 +16,74 @@ from core.model_runtime.model_providers import ModelProviderFactory from core.workflow.entities.node_entities import UserFrom from core.workflow.entities.variable_pool import VariablePool from core.workflow.enums import SystemVariableKey +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState +from core.workflow.nodes.event import RunCompletedEvent from core.workflow.nodes.llm.llm_node import LLMNode from extensions.ext_database import db from models.provider import ProviderType -from models.workflow import WorkflowNodeExecutionStatus +from models.workflow import WorkflowNodeExecutionStatus, WorkflowType """FOR MOCK FIXTURES, DO NOT REMOVE""" from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock -@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True) -def test_execute_llm(setup_openai_mock): - node = LLMNode( +def init_llm_node(config: dict) -> LLMNode: + graph_config = { + "edges": [ + { + "id": "start-source-next-target", + "source": "start", + "target": "llm", + }, + ], + "nodes": [{"data": {"type": "start"}, "id": "start"}, config], + } + + graph = Graph.init(graph_config=graph_config) + + init_params = GraphInitParams( tenant_id="1", app_id="1", + workflow_type=WorkflowType.WORKFLOW, workflow_id="1", + graph_config=graph_config, user_id="1", - invoke_from=InvokeFrom.WEB_APP, user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + # construct variable pool + variable_pool = VariablePool( + system_variables={ + SystemVariableKey.QUERY: "what's the weather today?", + SystemVariableKey.FILES: [], + SystemVariableKey.CONVERSATION_ID: "abababa", + SystemVariableKey.USER_ID: "aaa", + }, + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + variable_pool.add(["abc", "output"], "sunny") + + node = LLMNode( + id=str(uuid.uuid4()), + graph_init_params=init_params, + graph=graph, + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), + config=config, + ) + + return node + + +@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True) +def test_execute_llm(setup_openai_mock): + node = init_llm_node( config={ "id": "llm", "data": { @@ -49,19 +101,6 @@ def test_execute_llm(setup_openai_mock): }, ) - # construct variable pool - pool = VariablePool( - system_variables={ - SystemVariableKey.QUERY: "what's the weather today?", - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: "abababa", - SystemVariableKey.USER_ID: "aaa", - }, - user_inputs={}, - environment_variables=[], - ) - pool.add(["abc", "output"], "sunny") - credentials = {"openai_api_key": os.environ.get("OPENAI_API_KEY")} provider_instance = ModelProviderFactory().get_provider_instance("openai") @@ -80,13 +119,15 @@ def test_execute_llm(setup_openai_mock): model_type_instance=model_type_instance, ) model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo") + model_schema = model_type_instance.get_model_schema("gpt-3.5-turbo") + assert model_schema is not None model_config = ModelConfigWithCredentialsEntity( model="gpt-3.5-turbo", provider="openai", mode="chat", credentials=credentials, parameters={}, - model_schema=model_type_instance.get_model_schema("gpt-3.5-turbo"), + model_schema=model_schema, provider_model_bundle=provider_model_bundle, ) @@ -96,11 +137,16 @@ def test_execute_llm(setup_openai_mock): node._fetch_model_config = MagicMock(return_value=(model_instance, model_config)) # execute node - result = node.run(pool) + result = node._run() + assert isinstance(result, Generator) - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs["text"] is not None - assert result.outputs["usage"]["total_tokens"] > 0 + for item in result: + if isinstance(item, RunCompletedEvent): + assert item.run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert item.run_result.process_data is not None + assert item.run_result.outputs is not None + assert item.run_result.outputs.get("text") is not None + assert item.run_result.outputs.get("usage", {})["total_tokens"] > 0 @pytest.mark.parametrize("setup_code_executor_mock", [["none"]], indirect=True) @@ -109,13 +155,7 @@ def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_openai_mock): """ Test execute LLM node with jinja2 """ - node = LLMNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_llm_node( config={ "id": "llm", "data": { @@ -149,19 +189,6 @@ def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_openai_mock): }, ) - # construct variable pool - pool = VariablePool( - system_variables={ - SystemVariableKey.QUERY: "what's the weather today?", - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: "abababa", - SystemVariableKey.USER_ID: "aaa", - }, - user_inputs={}, - environment_variables=[], - ) - pool.add(["abc", "output"], "sunny") - credentials = {"openai_api_key": os.environ.get("OPENAI_API_KEY")} provider_instance = ModelProviderFactory().get_provider_instance("openai") @@ -181,14 +208,15 @@ def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_openai_mock): ) model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model="gpt-3.5-turbo") - + model_schema = model_type_instance.get_model_schema("gpt-3.5-turbo") + assert model_schema is not None model_config = ModelConfigWithCredentialsEntity( model="gpt-3.5-turbo", provider="openai", mode="chat", credentials=credentials, parameters={}, - model_schema=model_type_instance.get_model_schema("gpt-3.5-turbo"), + model_schema=model_schema, provider_model_bundle=provider_model_bundle, ) @@ -198,8 +226,11 @@ def test_execute_llm_with_jinja2(setup_code_executor_mock, setup_openai_mock): node._fetch_model_config = MagicMock(return_value=(model_instance, model_config)) # execute node - result = node.run(pool) + result = node._run() - assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert "sunny" in json.dumps(result.process_data) - assert "what's the weather today?" in json.dumps(result.process_data) + for item in result: + if isinstance(item, RunCompletedEvent): + assert item.run_result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert item.run_result.process_data is not None + assert "sunny" in json.dumps(item.run_result.process_data) + assert "what's the weather today?" in json.dumps(item.run_result.process_data) diff --git a/api/tests/integration_tests/workflow/nodes/test_parameter_extractor.py b/api/tests/integration_tests/workflow/nodes/test_parameter_extractor.py index 954dd1a305..cbe9c5914f 100644 --- a/api/tests/integration_tests/workflow/nodes/test_parameter_extractor.py +++ b/api/tests/integration_tests/workflow/nodes/test_parameter_extractor.py @@ -1,5 +1,7 @@ import json import os +import time +import uuid from typing import Optional from unittest.mock import MagicMock @@ -14,12 +16,15 @@ from core.model_runtime.model_providers.model_provider_factory import ModelProvi from core.workflow.entities.node_entities import UserFrom from core.workflow.entities.variable_pool import VariablePool from core.workflow.enums import SystemVariableKey +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState from core.workflow.nodes.parameter_extractor.parameter_extractor_node import ParameterExtractorNode from extensions.ext_database import db from models.provider import ProviderType """FOR MOCK FIXTURES, DO NOT REMOVE""" -from models.workflow import WorkflowNodeExecutionStatus +from models.workflow import WorkflowNodeExecutionStatus, WorkflowType from tests.integration_tests.model_runtime.__mock.anthropic import setup_anthropic_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock @@ -46,13 +51,15 @@ def get_mocked_fetch_model_config( model_type_instance=model_type_instance, ) model_instance = ModelInstance(provider_model_bundle=provider_model_bundle, model=model) + model_schema = model_type_instance.get_model_schema(model) + assert model_schema is not None model_config = ModelConfigWithCredentialsEntity( model=model, provider=provider, mode=mode, credentials=credentials, parameters={}, - model_schema=model_type_instance.get_model_schema(model), + model_schema=model_schema, provider_model_bundle=provider_model_bundle, ) @@ -73,18 +80,62 @@ def get_mocked_fetch_memory(memory_text: str): return MagicMock(return_value=MemoryMock()) +def init_parameter_extractor_node(config: dict): + graph_config = { + "edges": [ + { + "id": "start-source-next-target", + "source": "start", + "target": "llm", + }, + ], + "nodes": [{"data": {"type": "start"}, "id": "start"}, config], + } + + graph = Graph.init(graph_config=graph_config) + + init_params = GraphInitParams( + tenant_id="1", + app_id="1", + workflow_type=WorkflowType.WORKFLOW, + workflow_id="1", + graph_config=graph_config, + user_id="1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + # construct variable pool + variable_pool = VariablePool( + system_variables={ + SystemVariableKey.QUERY: "what's the weather in SF", + SystemVariableKey.FILES: [], + SystemVariableKey.CONVERSATION_ID: "abababa", + SystemVariableKey.USER_ID: "aaa", + }, + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + variable_pool.add(["a", "b123", "args1"], 1) + variable_pool.add(["a", "b123", "args2"], 2) + + return ParameterExtractorNode( + id=str(uuid.uuid4()), + graph_init_params=init_params, + graph=graph, + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), + config=config, + ) + + @pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True) def test_function_calling_parameter_extractor(setup_openai_mock): """ Test function calling for parameter extractor. """ - node = ParameterExtractorNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_parameter_extractor_node( config={ "id": "llm", "data": { @@ -97,7 +148,7 @@ def test_function_calling_parameter_extractor(setup_openai_mock): "reasoning_mode": "function_call", "memory": None, }, - }, + } ) node._fetch_model_config = get_mocked_fetch_model_config( @@ -120,9 +171,10 @@ def test_function_calling_parameter_extractor(setup_openai_mock): environment_variables=[], ) - result = node.run(pool) + result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert result.outputs.get("location") == "kawaii" assert result.outputs.get("__reason") == None @@ -132,13 +184,7 @@ def test_instructions(setup_openai_mock): """ Test chat parameter extractor. """ - node = ParameterExtractorNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_parameter_extractor_node( config={ "id": "llm", "data": { @@ -162,29 +208,19 @@ def test_instructions(setup_openai_mock): ) db.session.close = MagicMock() - # construct variable pool - pool = VariablePool( - system_variables={ - SystemVariableKey.QUERY: "what's the weather in SF", - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: "abababa", - SystemVariableKey.USER_ID: "aaa", - }, - user_inputs={}, - environment_variables=[], - ) - - result = node.run(pool) + result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert result.outputs.get("location") == "kawaii" assert result.outputs.get("__reason") == None process_data = result.process_data + assert process_data is not None process_data.get("prompts") - for prompt in process_data.get("prompts"): + for prompt in process_data.get("prompts", []): if prompt.get("role") == "system": assert "what's the weather in SF" in prompt.get("text") @@ -194,13 +230,7 @@ def test_chat_parameter_extractor(setup_anthropic_mock): """ Test chat parameter extractor. """ - node = ParameterExtractorNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_parameter_extractor_node( config={ "id": "llm", "data": { @@ -224,27 +254,17 @@ def test_chat_parameter_extractor(setup_anthropic_mock): ) db.session.close = MagicMock() - # construct variable pool - pool = VariablePool( - system_variables={ - SystemVariableKey.QUERY: "what's the weather in SF", - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: "abababa", - SystemVariableKey.USER_ID: "aaa", - }, - user_inputs={}, - environment_variables=[], - ) - - result = node.run(pool) + result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert result.outputs.get("location") == "" assert ( result.outputs.get("__reason") == "Failed to extract result from function call or text response, using empty result." ) - prompts = result.process_data.get("prompts") + assert result.process_data is not None + prompts = result.process_data.get("prompts", []) for prompt in prompts: if prompt.get("role") == "user": @@ -257,13 +277,7 @@ def test_completion_parameter_extractor(setup_openai_mock): """ Test completion parameter extractor. """ - node = ParameterExtractorNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_parameter_extractor_node( config={ "id": "llm", "data": { @@ -292,28 +306,18 @@ def test_completion_parameter_extractor(setup_openai_mock): ) db.session.close = MagicMock() - # construct variable pool - pool = VariablePool( - system_variables={ - SystemVariableKey.QUERY: "what's the weather in SF", - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: "abababa", - SystemVariableKey.USER_ID: "aaa", - }, - user_inputs={}, - environment_variables=[], - ) - - result = node.run(pool) + result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert result.outputs.get("location") == "" assert ( result.outputs.get("__reason") == "Failed to extract result from function call or text response, using empty result." ) - assert len(result.process_data.get("prompts")) == 1 - assert "SF" in result.process_data.get("prompts")[0].get("text") + assert result.process_data is not None + assert len(result.process_data.get("prompts", [])) == 1 + assert "SF" in result.process_data.get("prompts", [])[0].get("text") def test_extract_json_response(): @@ -321,13 +325,7 @@ def test_extract_json_response(): Test extract json response. """ - node = ParameterExtractorNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_parameter_extractor_node( config={ "id": "llm", "data": { @@ -356,6 +354,7 @@ def test_extract_json_response(): hello world. """) + assert result is not None assert result["location"] == "kawaii" @@ -364,13 +363,7 @@ def test_chat_parameter_extractor_with_memory(setup_anthropic_mock): """ Test chat parameter extractor with memory. """ - node = ParameterExtractorNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_parameter_extractor_node( config={ "id": "llm", "data": { @@ -395,27 +388,17 @@ def test_chat_parameter_extractor_with_memory(setup_anthropic_mock): node._fetch_memory = get_mocked_fetch_memory("customized memory") db.session.close = MagicMock() - # construct variable pool - pool = VariablePool( - system_variables={ - SystemVariableKey.QUERY: "what's the weather in SF", - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: "abababa", - SystemVariableKey.USER_ID: "aaa", - }, - user_inputs={}, - environment_variables=[], - ) - - result = node.run(pool) + result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert result.outputs.get("location") == "" assert ( result.outputs.get("__reason") == "Failed to extract result from function call or text response, using empty result." ) - prompts = result.process_data.get("prompts") + assert result.process_data is not None + prompts = result.process_data.get("prompts", []) latest_role = None for prompt in prompts: diff --git a/api/tests/integration_tests/workflow/nodes/test_template_transform.py b/api/tests/integration_tests/workflow/nodes/test_template_transform.py index 534e044ab3..073c4bb799 100644 --- a/api/tests/integration_tests/workflow/nodes/test_template_transform.py +++ b/api/tests/integration_tests/workflow/nodes/test_template_transform.py @@ -1,46 +1,84 @@ +import time +import uuid + import pytest from core.app.entities.app_invoke_entities import InvokeFrom from core.workflow.entities.node_entities import UserFrom from core.workflow.entities.variable_pool import VariablePool +from core.workflow.enums import SystemVariableKey +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState from core.workflow.nodes.template_transform.template_transform_node import TemplateTransformNode -from models.workflow import WorkflowNodeExecutionStatus +from models.workflow import WorkflowNodeExecutionStatus, WorkflowType from tests.integration_tests.workflow.nodes.__mock.code_executor import setup_code_executor_mock @pytest.mark.parametrize("setup_code_executor_mock", [["none"]], indirect=True) def test_execute_code(setup_code_executor_mock): code = """{{args2}}""" - node = TemplateTransformNode( + config = { + "id": "1", + "data": { + "title": "123", + "variables": [ + { + "variable": "args1", + "value_selector": ["1", "123", "args1"], + }, + {"variable": "args2", "value_selector": ["1", "123", "args2"]}, + ], + "template": code, + }, + } + + graph_config = { + "edges": [ + { + "id": "start-source-next-target", + "source": "start", + "target": "1", + }, + ], + "nodes": [{"data": {"type": "start"}, "id": "start"}, config], + } + + graph = Graph.init(graph_config=graph_config) + + init_params = GraphInitParams( tenant_id="1", app_id="1", + workflow_type=WorkflowType.WORKFLOW, workflow_id="1", + graph_config=graph_config, user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.END_USER, - config={ - "id": "1", - "data": { - "title": "123", - "variables": [ - { - "variable": "args1", - "value_selector": ["1", "123", "args1"], - }, - {"variable": "args2", "value_selector": ["1", "123", "args2"]}, - ], - "template": code, - }, - }, + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, ) # construct variable pool - pool = VariablePool(system_variables={}, user_inputs={}, environment_variables=[]) - pool.add(["1", "123", "args1"], 1) - pool.add(["1", "123", "args2"], 3) + variable_pool = VariablePool( + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + variable_pool.add(["1", "123", "args1"], 1) + variable_pool.add(["1", "123", "args2"], 3) + + node = TemplateTransformNode( + id=str(uuid.uuid4()), + graph_init_params=init_params, + graph=graph, + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), + config=config, + ) # execute node - result = node.run(pool) + result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert result.outputs["output"] == "3" diff --git a/api/tests/integration_tests/workflow/nodes/test_tool.py b/api/tests/integration_tests/workflow/nodes/test_tool.py index 778075a2de..4d94cdb28a 100644 --- a/api/tests/integration_tests/workflow/nodes/test_tool.py +++ b/api/tests/integration_tests/workflow/nodes/test_tool.py @@ -1,21 +1,62 @@ +import time +import uuid + from core.app.entities.app_invoke_entities import InvokeFrom -from core.workflow.entities.node_entities import UserFrom +from core.workflow.entities.node_entities import NodeRunResult, UserFrom from core.workflow.entities.variable_pool import VariablePool +from core.workflow.enums import SystemVariableKey +from core.workflow.graph_engine.entities.graph import Graph +from core.workflow.graph_engine.entities.graph_init_params import GraphInitParams +from core.workflow.graph_engine.entities.graph_runtime_state import GraphRuntimeState from core.workflow.nodes.tool.tool_node import ToolNode -from models.workflow import WorkflowNodeExecutionStatus +from models.workflow import WorkflowNodeExecutionStatus, WorkflowType + + +def init_tool_node(config: dict): + graph_config = { + "edges": [ + { + "id": "start-source-next-target", + "source": "start", + "target": "1", + }, + ], + "nodes": [{"data": {"type": "start"}, "id": "start"}, config], + } + + graph = Graph.init(graph_config=graph_config) + + init_params = GraphInitParams( + tenant_id="1", + app_id="1", + workflow_type=WorkflowType.WORKFLOW, + workflow_id="1", + graph_config=graph_config, + user_id="1", + user_from=UserFrom.ACCOUNT, + invoke_from=InvokeFrom.DEBUGGER, + call_depth=0, + ) + + # construct variable pool + variable_pool = VariablePool( + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, + user_inputs={}, + environment_variables=[], + conversation_variables=[], + ) + + return ToolNode( + id=str(uuid.uuid4()), + graph_init_params=init_params, + graph=graph, + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), + config=config, + ) def test_tool_variable_invoke(): - pool = VariablePool(system_variables={}, user_inputs={}, environment_variables=[]) - pool.add(["1", "123", "args1"], "1+1") - - node = ToolNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_tool_node( config={ "id": "1", "data": { @@ -34,28 +75,22 @@ def test_tool_variable_invoke(): } }, }, - }, + } ) - # execute node - result = node.run(pool) + node.graph_runtime_state.variable_pool.add(["1", "123", "args1"], "1+1") + # execute node + result = node._run() + assert isinstance(result, NodeRunResult) assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert "2" in result.outputs["text"] assert result.outputs["files"] == [] def test_tool_mixed_invoke(): - pool = VariablePool(system_variables={}, user_inputs={}, environment_variables=[]) - pool.add(["1", "args1"], "1+1") - - node = ToolNode( - tenant_id="1", - app_id="1", - workflow_id="1", - user_id="1", - invoke_from=InvokeFrom.WEB_APP, - user_from=UserFrom.ACCOUNT, + node = init_tool_node( config={ "id": "1", "data": { @@ -74,12 +109,15 @@ def test_tool_mixed_invoke(): } }, }, - }, + } ) - # execute node - result = node.run(pool) + node.graph_runtime_state.variable_pool.add(["1", "args1"], "1+1") + # execute node + result = node._run() + assert isinstance(result, NodeRunResult) assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED + assert result.outputs is not None assert "2" in result.outputs["text"] assert result.outputs["files"] == [] diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_graph.py b/api/tests/unit_tests/core/workflow/graph_engine/test_graph.py index a29ba16d9f..e9637839a8 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_graph.py +++ b/api/tests/unit_tests/core/workflow/graph_engine/test_graph.py @@ -32,33 +32,22 @@ def test_init(): "id": "http-source-answer2-target", "source": "http", "target": "answer2", - } + }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "llm", }, - "id": "llm" + "id": "llm", }, { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer", }, { - "data": { - "type": "question-classifier" - }, + "data": {"type": "question-classifier"}, "id": "qc", }, { @@ -68,19 +57,13 @@ def test_init(): "id": "http", }, { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer2", - } + }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) start_node_id = "start" @@ -127,7 +110,7 @@ def test__init_iteration_graph(): "source": "code", "sourceHandle": "source", "target": "iteration", - } + }, ], "nodes": [ { @@ -143,17 +126,11 @@ def test__init_iteration_graph(): "id": "llm", }, { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer", }, { - "data": { - "type": "iteration" - }, + "data": {"type": "iteration"}, "id": "iteration", }, { @@ -171,11 +148,7 @@ def test__init_iteration_graph(): "parentId": "iteration", }, { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer-in-iteration", "parentId": "iteration", }, @@ -184,27 +157,18 @@ def test__init_iteration_graph(): "type": "code", }, "id": "code", - } - ] + }, + ], } - graph = Graph.init( - graph_config=graph_config, - root_node_id="template-transform-in-iteration" - ) + graph = Graph.init(graph_config=graph_config, root_node_id="template-transform-in-iteration") graph.add_extra_edge( source_node_id="answer-in-iteration", target_node_id="template-transform-in-iteration", run_condition=RunCondition( type="condition", - conditions=[ - Condition( - variable_selector=["iteration", "index"], - comparison_operator="≤", - value="5" - ) - ] - ) + conditions=[Condition(variable_selector=["iteration", "index"], comparison_operator="≤", value="5")], + ), ) # iteration: @@ -248,47 +212,36 @@ def test_parallels_graph(): "id": "llm3-source-answer-target", "source": "llm3", "target": "answer", - } + }, ], "nodes": [ + {"data": {"type": "start"}, "id": "start"}, { "data": { - "type": "start" + "type": "llm", }, - "id": "start" + "id": "llm1", }, { "data": { "type": "llm", }, - "id": "llm1" + "id": "llm2", }, { "data": { "type": "llm", }, - "id": "llm2" + "id": "llm3", }, { - "data": { - "type": "llm", - }, - "id": "llm3" - }, - { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer", }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) assert graph.root_node_id == "start" for i in range(3): @@ -330,47 +283,36 @@ def test_parallels_graph2(): "id": "llm2-source-answer-target", "source": "llm2", "target": "answer", - } + }, ], "nodes": [ + {"data": {"type": "start"}, "id": "start"}, { "data": { - "type": "start" + "type": "llm", }, - "id": "start" + "id": "llm1", }, { "data": { "type": "llm", }, - "id": "llm1" + "id": "llm2", }, { "data": { "type": "llm", }, - "id": "llm2" + "id": "llm3", }, { - "data": { - "type": "llm", - }, - "id": "llm3" - }, - { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer", }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) assert graph.root_node_id == "start" for i in range(3): @@ -407,44 +349,33 @@ def test_parallels_graph3(): }, ], "nodes": [ + {"data": {"type": "start"}, "id": "start"}, { "data": { - "type": "start" + "type": "llm", }, - "id": "start" + "id": "llm1", }, { "data": { "type": "llm", }, - "id": "llm1" + "id": "llm2", }, { "data": { "type": "llm", }, - "id": "llm2" + "id": "llm3", }, { - "data": { - "type": "llm", - }, - "id": "llm3" - }, - { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer", }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) assert graph.root_node_id == "start" for i in range(3): @@ -504,65 +435,54 @@ def test_parallels_graph4(): "id": "code3-source-answer-target", "source": "code3", "target": "answer", - } + }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "llm", }, - "id": "llm1" + "id": "llm1", }, { "data": { "type": "code", }, - "id": "code1" + "id": "code1", }, { "data": { "type": "llm", }, - "id": "llm2" + "id": "llm2", }, { "data": { "type": "code", }, - "id": "code2" + "id": "code2", }, { "data": { "type": "llm", }, - "id": "llm3" + "id": "llm3", }, { "data": { "type": "code", }, - "id": "code3" + "id": "code3", }, { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer", }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) assert graph.root_node_id == "start" for i in range(3): @@ -641,77 +561,66 @@ def test_parallels_graph5(): "id": "code2-source-answer-target", "source": "code2", "target": "answer", - } + }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "llm", }, - "id": "llm1" + "id": "llm1", }, { "data": { "type": "code", }, - "id": "code1" + "id": "code1", }, { "data": { "type": "llm", }, - "id": "llm2" + "id": "llm2", }, { "data": { "type": "code", }, - "id": "code2" + "id": "code2", }, { "data": { "type": "llm", }, - "id": "llm3" + "id": "llm3", }, { "data": { "type": "code", }, - "id": "code3" + "id": "code3", }, { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer", }, { "data": { "type": "llm", }, - "id": "llm4" + "id": "llm4", }, { "data": { "type": "llm", }, - "id": "llm5" + "id": "llm5", }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) assert graph.root_node_id == "start" for i in range(5): @@ -786,65 +695,54 @@ def test_parallels_graph6(): "id": "code3-source-answer-target", "source": "code3", "target": "answer", - } + }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "llm", }, - "id": "llm1" + "id": "llm1", }, { "data": { "type": "code", }, - "id": "code1" + "id": "code1", }, { "data": { "type": "llm", }, - "id": "llm2" + "id": "llm2", }, { "data": { "type": "code", }, - "id": "code2" + "id": "code2", }, { "data": { "type": "llm", }, - "id": "llm3" + "id": "llm3", }, { "data": { "type": "code", }, - "id": "code3" + "id": "code3", }, { - "data": { - "type": "answer", - "title": "answer", - "answer": "1" - }, + "data": {"type": "answer", "title": "answer", "answer": "1"}, "id": "answer", }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) assert graph.root_node_id == "start" for i in range(3): diff --git a/api/tests/unit_tests/core/workflow/graph_engine/test_graph_engine.py b/api/tests/unit_tests/core/workflow/graph_engine/test_graph_engine.py index aa341f065b..a2d71d61fc 100644 --- a/api/tests/unit_tests/core/workflow/graph_engine/test_graph_engine.py +++ b/api/tests/unit_tests/core/workflow/graph_engine/test_graph_engine.py @@ -22,8 +22,8 @@ from core.workflow.nodes.llm.llm_node import LLMNode from models.workflow import WorkflowNodeExecutionStatus, WorkflowType -@patch('extensions.ext_database.db.session.remove') -@patch('extensions.ext_database.db.session.close') +@patch("extensions.ext_database.db.session.remove") +@patch("extensions.ext_database.db.session.close") def test_run_parallel_in_workflow(mock_close, mock_remove): graph_config = { "edges": [ @@ -51,85 +51,61 @@ def test_run_parallel_in_workflow(mock_close, mock_remove): "id": "5", "source": "llm3", "target": "end2", - } + }, ], "nodes": [ { "data": { "type": "start", "title": "start", - "variables": [{ - "label": "query", - "max_length": 48, - "options": [], - "required": True, - "type": "text-input", - "variable": "query" - }] + "variables": [ + { + "label": "query", + "max_length": 48, + "options": [], + "required": True, + "type": "text-input", + "variable": "query", + } + ], }, - "id": "start" + "id": "start", }, { "data": { "type": "llm", "title": "llm1", - "context": { - "enabled": False, - "variable_selector": [] - }, + "context": {"enabled": False, "variable_selector": []}, "model": { - "completion_params": { - "temperature": 0.7 - }, + "completion_params": {"temperature": 0.7}, "mode": "chat", "name": "gpt-4o", - "provider": "openai" + "provider": "openai", }, - "prompt_template": [{ - "role": "system", - "text": "say hi" - }, { - "role": "user", - "text": "{{#start.query#}}" - }], - "vision": { - "configs": { - "detail": "high" - }, - "enabled": False - } + "prompt_template": [ + {"role": "system", "text": "say hi"}, + {"role": "user", "text": "{{#start.query#}}"}, + ], + "vision": {"configs": {"detail": "high"}, "enabled": False}, }, - "id": "llm1" + "id": "llm1", }, { "data": { "type": "llm", "title": "llm2", - "context": { - "enabled": False, - "variable_selector": [] - }, + "context": {"enabled": False, "variable_selector": []}, "model": { - "completion_params": { - "temperature": 0.7 - }, + "completion_params": {"temperature": 0.7}, "mode": "chat", "name": "gpt-4o", - "provider": "openai" + "provider": "openai", }, - "prompt_template": [{ - "role": "system", - "text": "say bye" - }, { - "role": "user", - "text": "{{#start.query#}}" - }], - "vision": { - "configs": { - "detail": "high" - }, - "enabled": False - } + "prompt_template": [ + {"role": "system", "text": "say bye"}, + {"role": "user", "text": "{{#start.query#}}"}, + ], + "vision": {"configs": {"detail": "high"}, "enabled": False}, }, "id": "llm2", }, @@ -137,31 +113,18 @@ def test_run_parallel_in_workflow(mock_close, mock_remove): "data": { "type": "llm", "title": "llm3", - "context": { - "enabled": False, - "variable_selector": [] - }, + "context": {"enabled": False, "variable_selector": []}, "model": { - "completion_params": { - "temperature": 0.7 - }, + "completion_params": {"temperature": 0.7}, "mode": "chat", "name": "gpt-4o", - "provider": "openai" + "provider": "openai", }, - "prompt_template": [{ - "role": "system", - "text": "say good morning" - }, { - "role": "user", - "text": "{{#start.query#}}" - }], - "vision": { - "configs": { - "detail": "high" - }, - "enabled": False - } + "prompt_template": [ + {"role": "system", "text": "say good morning"}, + {"role": "user", "text": "{{#start.query#}}"}, + ], + "vision": {"configs": {"detail": "high"}, "enabled": False}, }, "id": "llm3", }, @@ -169,13 +132,10 @@ def test_run_parallel_in_workflow(mock_close, mock_remove): "data": { "type": "end", "title": "end1", - "outputs": [{ - "value_selector": ["llm2", "text"], - "variable": "result2" - }, { - "value_selector": ["start", "query"], - "variable": "query" - }], + "outputs": [ + {"value_selector": ["llm2", "text"], "variable": "result2"}, + {"value_selector": ["start", "query"], "variable": "query"}, + ], }, "id": "end1", }, @@ -183,29 +143,21 @@ def test_run_parallel_in_workflow(mock_close, mock_remove): "data": { "type": "end", "title": "end2", - "outputs": [{ - "value_selector": ["llm1", "text"], - "variable": "result1" - }, { - "value_selector": ["llm3", "text"], - "variable": "result3" - }], + "outputs": [ + {"value_selector": ["llm1", "text"], "variable": "result1"}, + {"value_selector": ["llm3", "text"], "variable": "result3"}, + ], }, "id": "end2", - } + }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) - variable_pool = VariablePool(system_variables={ - SystemVariableKey.FILES: [], - SystemVariableKey.USER_ID: 'aaa' - }, user_inputs={ - "query": "hi" - }) + variable_pool = VariablePool( + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, user_inputs={"query": "hi"} + ) graph_engine = GraphEngine( tenant_id="111", @@ -220,19 +172,14 @@ def test_run_parallel_in_workflow(mock_close, mock_remove): graph=graph, variable_pool=variable_pool, max_execution_steps=500, - max_execution_time=1200 + max_execution_time=1200, ) def llm_generator(self): - contents = [ - 'hi', - 'bye', - 'good morning' - ] + contents = ["hi", "bye", "good morning"] yield RunStreamChunkEvent( - chunk_content=contents[int(self.node_id[-1]) - 1], - from_variable_selector=[self.node_id, 'text'] + chunk_content=contents[int(self.node_id[-1]) - 1], from_variable_selector=[self.node_id, "text"] ) yield RunCompletedEvent( @@ -244,14 +191,14 @@ def test_run_parallel_in_workflow(mock_close, mock_remove): metadata={ NodeRunMetadataKey.TOTAL_TOKENS: 1, NodeRunMetadataKey.TOTAL_PRICE: 1, - NodeRunMetadataKey.CURRENCY: 'USD' - } + NodeRunMetadataKey.CURRENCY: "USD", + }, ) ) # print("") - with patch.object(LLMNode, '_run', new=llm_generator): + with patch.object(LLMNode, "_run", new=llm_generator): items = [] generator = graph_engine.run() for item in generator: @@ -263,21 +210,19 @@ def test_run_parallel_in_workflow(mock_close, mock_remove): assert not isinstance(item, NodeRunFailedEvent) assert not isinstance(item, GraphRunFailedEvent) - if isinstance(item, BaseNodeEvent) and item.route_node_state.node_id in [ - 'llm2', 'llm3', 'end1', 'end2' - ]: + if isinstance(item, BaseNodeEvent) and item.route_node_state.node_id in ["llm2", "llm3", "end1", "end2"]: assert item.parallel_id is not None assert len(items) == 18 assert isinstance(items[0], GraphRunStartedEvent) assert isinstance(items[1], NodeRunStartedEvent) - assert items[1].route_node_state.node_id == 'start' + assert items[1].route_node_state.node_id == "start" assert isinstance(items[2], NodeRunSucceededEvent) - assert items[2].route_node_state.node_id == 'start' + assert items[2].route_node_state.node_id == "start" -@patch('extensions.ext_database.db.session.remove') -@patch('extensions.ext_database.db.session.close') +@patch("extensions.ext_database.db.session.remove") +@patch("extensions.ext_database.db.session.close") def test_run_parallel_in_chatflow(mock_close, mock_remove): graph_config = { "edges": [ @@ -305,69 +250,41 @@ def test_run_parallel_in_chatflow(mock_close, mock_remove): "id": "5", "source": "answer3", "target": "answer5", - } + }, ], "nodes": [ + {"data": {"type": "start", "title": "start"}, "id": "start"}, + {"data": {"type": "answer", "title": "answer1", "answer": "1"}, "id": "answer1"}, { - "data": { - "type": "start", - "title": "start" - }, - "id": "start" - }, - { - "data": { - "type": "answer", - "title": "answer1", - "answer": "1" - }, - "id": "answer1" - }, - { - "data": { - "type": "answer", - "title": "answer2", - "answer": "2" - }, + "data": {"type": "answer", "title": "answer2", "answer": "2"}, "id": "answer2", }, { - "data": { - "type": "answer", - "title": "answer3", - "answer": "3" - }, + "data": {"type": "answer", "title": "answer3", "answer": "3"}, "id": "answer3", }, { - "data": { - "type": "answer", - "title": "answer4", - "answer": "4" - }, + "data": {"type": "answer", "title": "answer4", "answer": "4"}, "id": "answer4", }, { - "data": { - "type": "answer", - "title": "answer5", - "answer": "5" - }, + "data": {"type": "answer", "title": "answer5", "answer": "5"}, "id": "answer5", - } + }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) - variable_pool = VariablePool(system_variables={ - SystemVariableKey.QUERY: 'what\'s the weather in SF', - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: 'abababa', - SystemVariableKey.USER_ID: 'aaa' - }, user_inputs={}) + variable_pool = VariablePool( + system_variables={ + SystemVariableKey.QUERY: "what's the weather in SF", + SystemVariableKey.FILES: [], + SystemVariableKey.CONVERSATION_ID: "abababa", + SystemVariableKey.USER_ID: "aaa", + }, + user_inputs={}, + ) graph_engine = GraphEngine( tenant_id="111", @@ -382,7 +299,7 @@ def test_run_parallel_in_chatflow(mock_close, mock_remove): graph=graph, variable_pool=variable_pool, max_execution_steps=500, - max_execution_time=1200 + max_execution_time=1200, ) # print("") @@ -399,135 +316,155 @@ def test_run_parallel_in_chatflow(mock_close, mock_remove): assert not isinstance(item, GraphRunFailedEvent) if isinstance(item, BaseNodeEvent) and item.route_node_state.node_id in [ - 'answer2', 'answer3', 'answer4', 'answer5' + "answer2", + "answer3", + "answer4", + "answer5", ]: assert item.parallel_id is not None assert len(items) == 23 assert isinstance(items[0], GraphRunStartedEvent) assert isinstance(items[1], NodeRunStartedEvent) - assert items[1].route_node_state.node_id == 'start' + assert items[1].route_node_state.node_id == "start" assert isinstance(items[2], NodeRunSucceededEvent) - assert items[2].route_node_state.node_id == 'start' + assert items[2].route_node_state.node_id == "start" -@patch('extensions.ext_database.db.session.remove') -@patch('extensions.ext_database.db.session.close') +@patch("extensions.ext_database.db.session.remove") +@patch("extensions.ext_database.db.session.close") def test_run_branch(mock_close, mock_remove): graph_config = { - "edges": [{ - "id": "1", - "source": "start", - "target": "if-else-1", - }, { - "id": "2", - "source": "if-else-1", - "sourceHandle": "true", - "target": "answer-1", - }, { - "id": "3", - "source": "if-else-1", - "sourceHandle": "false", - "target": "if-else-2", - }, { - "id": "4", - "source": "if-else-2", - "sourceHandle": "true", - "target": "answer-2", - }, { - "id": "5", - "source": "if-else-2", - "sourceHandle": "false", - "target": "answer-3", - }], - "nodes": [{ - "data": { - "title": "Start", - "type": "start", - "variables": [{ - "label": "uid", - "max_length": 48, - "options": [], - "required": True, - "type": "text-input", - "variable": "uid" - }] + "edges": [ + { + "id": "1", + "source": "start", + "target": "if-else-1", }, - "id": "start" - }, { - "data": { - "answer": "1 {{#start.uid#}}", - "title": "Answer", - "type": "answer", - "variables": [] + { + "id": "2", + "source": "if-else-1", + "sourceHandle": "true", + "target": "answer-1", }, - "id": "answer-1", - }, { - "data": { - "cases": [{ - "case_id": "true", - "conditions": [{ - "comparison_operator": "contains", - "id": "b0f02473-08b6-4a81-af91-15345dcb2ec8", - "value": "hi", - "varType": "string", - "variable_selector": ["sys", "query"] - }], - "id": "true", - "logical_operator": "and" - }], - "desc": "", - "title": "IF/ELSE", - "type": "if-else" + { + "id": "3", + "source": "if-else-1", + "sourceHandle": "false", + "target": "if-else-2", }, - "id": "if-else-1", - }, { - "data": { - "cases": [{ - "case_id": "true", - "conditions": [{ - "comparison_operator": "contains", - "id": "ae895199-5608-433b-b5f0-0997ae1431e4", - "value": "takatost", - "varType": "string", - "variable_selector": ["sys", "query"] - }], - "id": "true", - "logical_operator": "and" - }], - "title": "IF/ELSE 2", - "type": "if-else" + { + "id": "4", + "source": "if-else-2", + "sourceHandle": "true", + "target": "answer-2", }, - "id": "if-else-2", - }, { - "data": { - "answer": "2", - "title": "Answer 2", - "type": "answer", + { + "id": "5", + "source": "if-else-2", + "sourceHandle": "false", + "target": "answer-3", }, - "id": "answer-2", - }, { - "data": { - "answer": "3", - "title": "Answer 3", - "type": "answer", + ], + "nodes": [ + { + "data": { + "title": "Start", + "type": "start", + "variables": [ + { + "label": "uid", + "max_length": 48, + "options": [], + "required": True, + "type": "text-input", + "variable": "uid", + } + ], + }, + "id": "start", }, - "id": "answer-3", - }] + { + "data": {"answer": "1 {{#start.uid#}}", "title": "Answer", "type": "answer", "variables": []}, + "id": "answer-1", + }, + { + "data": { + "cases": [ + { + "case_id": "true", + "conditions": [ + { + "comparison_operator": "contains", + "id": "b0f02473-08b6-4a81-af91-15345dcb2ec8", + "value": "hi", + "varType": "string", + "variable_selector": ["sys", "query"], + } + ], + "id": "true", + "logical_operator": "and", + } + ], + "desc": "", + "title": "IF/ELSE", + "type": "if-else", + }, + "id": "if-else-1", + }, + { + "data": { + "cases": [ + { + "case_id": "true", + "conditions": [ + { + "comparison_operator": "contains", + "id": "ae895199-5608-433b-b5f0-0997ae1431e4", + "value": "takatost", + "varType": "string", + "variable_selector": ["sys", "query"], + } + ], + "id": "true", + "logical_operator": "and", + } + ], + "title": "IF/ELSE 2", + "type": "if-else", + }, + "id": "if-else-2", + }, + { + "data": { + "answer": "2", + "title": "Answer 2", + "type": "answer", + }, + "id": "answer-2", + }, + { + "data": { + "answer": "3", + "title": "Answer 3", + "type": "answer", + }, + "id": "answer-3", + }, + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) - variable_pool = VariablePool(system_variables={ - SystemVariableKey.QUERY: 'hi', - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: 'abababa', - SystemVariableKey.USER_ID: 'aaa' - }, user_inputs={ - "uid": "takato" - }) + variable_pool = VariablePool( + system_variables={ + SystemVariableKey.QUERY: "hi", + SystemVariableKey.FILES: [], + SystemVariableKey.CONVERSATION_ID: "abababa", + SystemVariableKey.USER_ID: "aaa", + }, + user_inputs={"uid": "takato"}, + ) graph_engine = GraphEngine( tenant_id="111", @@ -542,7 +479,7 @@ def test_run_branch(mock_close, mock_remove): graph=graph, variable_pool=variable_pool, max_execution_steps=500, - max_execution_time=1200 + max_execution_time=1200, ) # print("") @@ -554,15 +491,15 @@ def test_run_branch(mock_close, mock_remove): items.append(item) assert len(items) == 10 - assert items[3].route_node_state.node_id == 'if-else-1' - assert items[4].route_node_state.node_id == 'if-else-1' + assert items[3].route_node_state.node_id == "if-else-1" + assert items[4].route_node_state.node_id == "if-else-1" assert isinstance(items[5], NodeRunStreamChunkEvent) - assert items[5].chunk_content == '1 ' + assert items[5].chunk_content == "1 " assert isinstance(items[6], NodeRunStreamChunkEvent) - assert items[6].chunk_content == 'takato' - assert items[7].route_node_state.node_id == 'answer-1' - assert items[8].route_node_state.node_id == 'answer-1' - assert items[8].route_node_state.node_run_result.outputs['answer'] == '1 takato' + assert items[6].chunk_content == "takato" + assert items[7].route_node_state.node_id == "answer-1" + assert items[8].route_node_state.node_id == "answer-1" + assert items[8].route_node_state.node_run_result.outputs["answer"] == "1 takato" assert isinstance(items[9], GraphRunSucceededEvent) # print(graph_engine.graph_runtime_state.model_dump_json(indent=2)) diff --git a/api/tests/unit_tests/core/workflow/nodes/answer/test_answer.py b/api/tests/unit_tests/core/workflow/nodes/answer/test_answer.py index fbcb209d07..fe4ede6335 100644 --- a/api/tests/unit_tests/core/workflow/nodes/answer/test_answer.py +++ b/api/tests/unit_tests/core/workflow/nodes/answer/test_answer.py @@ -24,61 +24,52 @@ def test_execute_answer(): }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "llm", }, - "id": "llm" + "id": "llm", }, - ] + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.WORKFLOW, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) # construct variable pool - pool = VariablePool(system_variables={ - SystemVariableKey.FILES: [], - SystemVariableKey.USER_ID: 'aaa' - }, user_inputs={}, environment_variables=[]) - pool.add(['start', 'weather'], 'sunny') - pool.add(['llm', 'text'], 'You are a helpful AI.') + pool = VariablePool( + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, + user_inputs={}, + environment_variables=[], + ) + pool.add(["start", "weather"], "sunny") + pool.add(["llm", "text"], "You are a helpful AI.") node = AnswerNode( id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=pool, start_at=time.perf_counter()), config={ - 'id': 'answer', - 'data': { - 'title': '123', - 'type': 'answer', - 'answer': 'Today\'s weather is {{#start.weather#}}\n{{#llm.text#}}\n{{img}}\nFin.' - } - } + "id": "answer", + "data": { + "title": "123", + "type": "answer", + "answer": "Today's weather is {{#start.weather#}}\n{{#llm.text#}}\n{{img}}\nFin.", + }, + }, ) # Mock db.session.close() @@ -88,4 +79,4 @@ def test_execute_answer(): result = node._run() assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED - assert result.outputs['answer'] == "Today's weather is sunny\nYou are a helpful AI.\n{{img}}\nFin." + assert result.outputs["answer"] == "Today's weather is sunny\nYou are a helpful AI.\n{{img}}\nFin." diff --git a/api/tests/unit_tests/core/workflow/nodes/answer/test_answer_stream_generate_router.py b/api/tests/unit_tests/core/workflow/nodes/answer/test_answer_stream_generate_router.py index 7a48c3548b..bce87536d8 100644 --- a/api/tests/unit_tests/core/workflow/nodes/answer/test_answer_stream_generate_router.py +++ b/api/tests/unit_tests/core/workflow/nodes/answer/test_answer_stream_generate_router.py @@ -54,72 +54,56 @@ def test_init(): "id": "llm1-source-answer-target", "source": "llm1", "target": "answer", - } + }, ], "nodes": [ + {"data": {"type": "start"}, "id": "start"}, { "data": { - "type": "start" + "type": "llm", }, - "id": "start" + "id": "llm1", }, { "data": { "type": "llm", }, - "id": "llm1" + "id": "llm2", }, { "data": { "type": "llm", }, - "id": "llm2" + "id": "llm3", }, { "data": { "type": "llm", }, - "id": "llm3" + "id": "llm4", }, { "data": { "type": "llm", }, - "id": "llm4" + "id": "llm5", }, { - "data": { - "type": "llm", - }, - "id": "llm5" - }, - { - "data": { - "type": "answer", - "title": "answer", - "answer": "1{{#llm2.text#}}2" - }, + "data": {"type": "answer", "title": "answer", "answer": "1{{#llm2.text#}}2"}, "id": "answer", }, { - "data": { - "type": "answer", - "title": "answer2", - "answer": "1{{#llm3.text#}}2" - }, + "data": {"type": "answer", "title": "answer2", "answer": "1{{#llm3.text#}}2"}, "id": "answer2", }, ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) answer_stream_generate_route = AnswerStreamGeneratorRouter.init( - node_id_config_mapping=graph.node_id_config_mapping, - reverse_edge_mapping=graph.reverse_edge_mapping + node_id_config_mapping=graph.node_id_config_mapping, reverse_edge_mapping=graph.reverse_edge_mapping ) - assert answer_stream_generate_route.answer_dependencies['answer'] == ['answer2'] - assert answer_stream_generate_route.answer_dependencies['answer2'] == [] + assert answer_stream_generate_route.answer_dependencies["answer"] == ["answer2"] + assert answer_stream_generate_route.answer_dependencies["answer2"] == [] diff --git a/api/tests/unit_tests/core/workflow/nodes/answer/test_answer_stream_processor.py b/api/tests/unit_tests/core/workflow/nodes/answer/test_answer_stream_processor.py index 13b74d65ac..6b1d1e9070 100644 --- a/api/tests/unit_tests/core/workflow/nodes/answer/test_answer_stream_processor.py +++ b/api/tests/unit_tests/core/workflow/nodes/answer/test_answer_stream_processor.py @@ -18,7 +18,7 @@ from core.workflow.nodes.start.entities import StartNodeData def _recursive_process(graph: Graph, next_node_id: str) -> Generator[GraphEngineEvent, None, None]: - if next_node_id == 'start': + if next_node_id == "start": yield from _publish_events(graph, next_node_id) for edge in graph.edge_mapping.get(next_node_id, []): @@ -29,10 +29,7 @@ def _recursive_process(graph: Graph, next_node_id: str) -> Generator[GraphEngine def _publish_events(graph: Graph, next_node_id: str) -> Generator[GraphEngineEvent, None, None]: - route_node_state = RouteNodeState( - node_id=next_node_id, - start_at=datetime.now(timezone.utc).replace(tzinfo=None) - ) + route_node_state = RouteNodeState(node_id=next_node_id, start_at=datetime.now(timezone.utc).replace(tzinfo=None)) parallel_id = graph.node_parallel_mapping.get(next_node_id) parallel_start_node_id = None @@ -52,10 +49,10 @@ def _publish_events(graph: Graph, next_node_id: str) -> Generator[GraphEngineEve node_data=mock_node_data, route_node_state=route_node_state, parallel_id=graph.node_parallel_mapping.get(next_node_id), - parallel_start_node_id=parallel_start_node_id + parallel_start_node_id=parallel_start_node_id, ) - if 'llm' in next_node_id: + if "llm" in next_node_id: length = int(next_node_id[-1]) for i in range(0, length): yield NodeRunStreamChunkEvent( @@ -67,7 +64,7 @@ def _publish_events(graph: Graph, next_node_id: str) -> Generator[GraphEngineEve route_node_state=route_node_state, from_variable_selector=[next_node_id, "text"], parallel_id=parallel_id, - parallel_start_node_id=parallel_start_node_id + parallel_start_node_id=parallel_start_node_id, ) route_node_state.status = RouteNodeState.Status.SUCCESS @@ -79,7 +76,7 @@ def _publish_events(graph: Graph, next_node_id: str) -> Generator[GraphEngineEve node_data=mock_node_data, route_node_state=route_node_state, parallel_id=parallel_id, - parallel_start_node_id=parallel_start_node_id + parallel_start_node_id=parallel_start_node_id, ) @@ -135,79 +132,64 @@ def test_process(): "id": "llm1-source-answer-target", "source": "llm1", "target": "answer", - } + }, ], "nodes": [ + {"data": {"type": "start"}, "id": "start"}, { "data": { - "type": "start" + "type": "llm", }, - "id": "start" + "id": "llm1", }, { "data": { "type": "llm", }, - "id": "llm1" + "id": "llm2", }, { "data": { "type": "llm", }, - "id": "llm2" + "id": "llm3", }, { "data": { "type": "llm", }, - "id": "llm3" + "id": "llm4", }, { "data": { "type": "llm", }, - "id": "llm4" + "id": "llm5", }, { - "data": { - "type": "llm", - }, - "id": "llm5" - }, - { - "data": { - "type": "answer", - "title": "answer", - "answer": "a{{#llm2.text#}}b" - }, + "data": {"type": "answer", "title": "answer", "answer": "a{{#llm2.text#}}b"}, "id": "answer", }, { - "data": { - "type": "answer", - "title": "answer2", - "answer": "c{{#llm3.text#}}d" - }, + "data": {"type": "answer", "title": "answer2", "answer": "c{{#llm3.text#}}d"}, "id": "answer2", }, ], } - graph = Graph.init( - graph_config=graph_config + graph = Graph.init(graph_config=graph_config) + + variable_pool = VariablePool( + system_variables={ + SystemVariableKey.QUERY: "what's the weather in SF", + SystemVariableKey.FILES: [], + SystemVariableKey.CONVERSATION_ID: "abababa", + SystemVariableKey.USER_ID: "aaa", + }, + user_inputs={}, ) - variable_pool = VariablePool(system_variables={ - SystemVariableKey.QUERY: 'what\'s the weather in SF', - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: 'abababa', - SystemVariableKey.USER_ID: 'aaa' - }, user_inputs={}) - - answer_stream_processor = AnswerStreamProcessor( - graph=graph, - variable_pool=variable_pool - ) + answer_stream_processor = AnswerStreamProcessor(graph=graph, variable_pool=variable_pool) def graph_generator() -> Generator[GraphEngineEvent, None, None]: # print("") @@ -215,10 +197,10 @@ def test_process(): # print("[ORIGIN]", event.__class__.__name__ + ":", event.route_node_state.node_id, # " " + (event.chunk_content if isinstance(event, NodeRunStreamChunkEvent) else "")) if isinstance(event, NodeRunSucceededEvent): - if 'llm' in event.route_node_state.node_id: + if "llm" in event.route_node_state.node_id: variable_pool.add( [event.route_node_state.node_id, "text"], - "".join(str(i) for i in range(0, int(event.route_node_state.node_id[-1]))) + "".join(str(i) for i in range(0, int(event.route_node_state.node_id[-1]))), ) yield event diff --git a/api/tests/unit_tests/core/workflow/nodes/iteration/test_iteration.py b/api/tests/unit_tests/core/workflow/nodes/iteration/test_iteration.py index 344553f344..89855df6a7 100644 --- a/api/tests/unit_tests/core/workflow/nodes/iteration/test_iteration.py +++ b/api/tests/unit_tests/core/workflow/nodes/iteration/test_iteration.py @@ -17,159 +17,152 @@ from models.workflow import WorkflowNodeExecutionStatus, WorkflowType def test_run(): graph_config = { - "edges": [{ - "id": "start-source-pe-target", - "source": "start", - "target": "pe", - }, { - "id": "iteration-1-source-answer-3-target", - "source": "iteration-1", - "target": "answer-3", - }, { - "id": "tt-source-if-else-target", - "source": "tt", - "target": "if-else", - }, { - "id": "if-else-true-answer-2-target", - "source": "if-else", - "sourceHandle": "true", - "target": "answer-2", - }, { - "id": "if-else-false-answer-4-target", - "source": "if-else", - "sourceHandle": "false", - "target": "answer-4", - }, { - "id": "pe-source-iteration-1-target", - "source": "pe", - "target": "iteration-1", - }], - "nodes": [{ - "data": { - "title": "Start", - "type": "start", - "variables": [] + "edges": [ + { + "id": "start-source-pe-target", + "source": "start", + "target": "pe", }, - "id": "start" - }, { - "data": { - "iterator_selector": ["pe", "list_output"], - "output_selector": ["tt", "output"], - "output_type": "array[string]", - "startNodeType": "template-transform", - "start_node_id": "tt", - "title": "iteration", - "type": "iteration", + { + "id": "iteration-1-source-answer-3-target", + "source": "iteration-1", + "target": "answer-3", }, - "id": "iteration-1", - }, { - "data": { - "answer": "{{#tt.output#}}", - "iteration_id": "iteration-1", - "title": "answer 2", - "type": "answer" + { + "id": "tt-source-if-else-target", + "source": "tt", + "target": "if-else", }, - "id": "answer-2" - }, { - "data": { - "iteration_id": "iteration-1", - "template": "{{ arg1 }} 123", - "title": "template transform", - "type": "template-transform", - "variables": [{ - "value_selector": ["sys", "query"], - "variable": "arg1" - }] + { + "id": "if-else-true-answer-2-target", + "source": "if-else", + "sourceHandle": "true", + "target": "answer-2", }, - "id": "tt", - }, { - "data": { - "answer": "{{#iteration-1.output#}}88888", - "title": "answer 3", - "type": "answer" + { + "id": "if-else-false-answer-4-target", + "source": "if-else", + "sourceHandle": "false", + "target": "answer-4", }, - "id": "answer-3", - }, { - "data": { - "conditions": [{ - "comparison_operator": "is", - "id": "1721916275284", - "value": "hi", - "variable_selector": ["sys", "query"] - }], - "iteration_id": "iteration-1", - "logical_operator": "and", - "title": "if", - "type": "if-else" + { + "id": "pe-source-iteration-1-target", + "source": "pe", + "target": "iteration-1", }, - "id": "if-else", - }, { - "data": { - "answer": "no hi", - "iteration_id": "iteration-1", - "title": "answer 4", - "type": "answer" - }, - "id": "answer-4", - }, { - "data": { - "instruction": "test1", - "model": { - "completion_params": { - "temperature": 0.7 - }, - "mode": "chat", - "name": "gpt-4o", - "provider": "openai" + ], + "nodes": [ + {"data": {"title": "Start", "type": "start", "variables": []}, "id": "start"}, + { + "data": { + "iterator_selector": ["pe", "list_output"], + "output_selector": ["tt", "output"], + "output_type": "array[string]", + "startNodeType": "template-transform", + "start_node_id": "tt", + "title": "iteration", + "type": "iteration", }, - "parameters": [{ - "description": "test", - "name": "list_output", - "required": False, - "type": "array[string]" - }], - "query": ["sys", "query"], - "reasoning_mode": "prompt", - "title": "pe", - "type": "parameter-extractor" + "id": "iteration-1", }, - "id": "pe", - }] + { + "data": { + "answer": "{{#tt.output#}}", + "iteration_id": "iteration-1", + "title": "answer 2", + "type": "answer", + }, + "id": "answer-2", + }, + { + "data": { + "iteration_id": "iteration-1", + "template": "{{ arg1 }} 123", + "title": "template transform", + "type": "template-transform", + "variables": [{"value_selector": ["sys", "query"], "variable": "arg1"}], + }, + "id": "tt", + }, + { + "data": {"answer": "{{#iteration-1.output#}}88888", "title": "answer 3", "type": "answer"}, + "id": "answer-3", + }, + { + "data": { + "conditions": [ + { + "comparison_operator": "is", + "id": "1721916275284", + "value": "hi", + "variable_selector": ["sys", "query"], + } + ], + "iteration_id": "iteration-1", + "logical_operator": "and", + "title": "if", + "type": "if-else", + }, + "id": "if-else", + }, + { + "data": {"answer": "no hi", "iteration_id": "iteration-1", "title": "answer 4", "type": "answer"}, + "id": "answer-4", + }, + { + "data": { + "instruction": "test1", + "model": { + "completion_params": {"temperature": 0.7}, + "mode": "chat", + "name": "gpt-4o", + "provider": "openai", + }, + "parameters": [ + {"description": "test", "name": "list_output", "required": False, "type": "array[string]"} + ], + "query": ["sys", "query"], + "reasoning_mode": "prompt", + "title": "pe", + "type": "parameter-extractor", + }, + "id": "pe", + }, + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.CHAT, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) # construct variable pool - pool = VariablePool(system_variables={ - SystemVariableKey.QUERY: 'dify', - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: 'abababa', - SystemVariableKey.USER_ID: '1' - }, user_inputs={}, environment_variables=[]) - pool.add(['pe', 'list_output'], ["dify-1", "dify-2"]) + pool = VariablePool( + system_variables={ + SystemVariableKey.QUERY: "dify", + SystemVariableKey.FILES: [], + SystemVariableKey.CONVERSATION_ID: "abababa", + SystemVariableKey.USER_ID: "1", + }, + user_inputs={}, + environment_variables=[], + ) + pool.add(["pe", "list_output"], ["dify-1", "dify-2"]) iteration_node = IterationNode( id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=pool, start_at=time.perf_counter()), config={ "data": { "iterator_selector": ["pe", "list_output"], @@ -181,23 +174,19 @@ def test_run(): "type": "iteration", }, "id": "iteration-1", - } + }, ) def tt_generator(self): return NodeRunResult( status=WorkflowNodeExecutionStatus.SUCCEEDED, - inputs={ - 'iterator_selector': 'dify' - }, - outputs={ - 'output': 'dify 123' - } + inputs={"iterator_selector": "dify"}, + outputs={"output": "dify 123"}, ) # print("") - with patch.object(TemplateTransformNode, '_run', new=tt_generator): + with patch.object(TemplateTransformNode, "_run", new=tt_generator): # execute node result = iteration_node._run() @@ -214,177 +203,169 @@ def test_run(): def test_run_parallel(): graph_config = { - "edges": [{ - "id": "start-source-pe-target", - "source": "start", - "target": "pe", - }, { - "id": "iteration-1-source-answer-3-target", - "source": "iteration-1", - "target": "answer-3", - }, { - "id": "tt-source-if-else-target", - "source": "tt", - "target": "if-else", - }, { - "id": "tt-2-source-if-else-target", - "source": "tt-2", - "target": "if-else", - }, { - "id": "if-else-true-answer-2-target", - "source": "if-else", - "sourceHandle": "true", - "target": "answer-2", - }, { - "id": "if-else-false-answer-4-target", - "source": "if-else", - "sourceHandle": "false", - "target": "answer-4", - }, { - "id": "pe-source-iteration-1-target", - "source": "pe", - "target": "iteration-1", - }], - "nodes": [{ - "data": { - "title": "Start", - "type": "start", - "variables": [] + "edges": [ + { + "id": "start-source-pe-target", + "source": "start", + "target": "pe", }, - "id": "start" - }, { - "data": { - "iterator_selector": ["pe", "list_output"], - "output_selector": ["tt", "output"], - "output_type": "array[string]", - "startNodeType": "template-transform", - "start_node_id": "tt", - "title": "iteration", - "type": "iteration", + { + "id": "iteration-1-source-answer-3-target", + "source": "iteration-1", + "target": "answer-3", }, - "id": "iteration-1", - }, { - "data": { - "answer": "{{#tt.output#}}", - "iteration_id": "iteration-1", - "title": "answer 2", - "type": "answer" + { + "id": "tt-source-if-else-target", + "source": "tt", + "target": "if-else", }, - "id": "answer-2" - }, { - "data": { - "iteration_id": "iteration-1", - "start_node_in_iteration": True, - "template": "{{ arg1 }} 123", - "title": "template transform", - "type": "template-transform", - "variables": [{ - "value_selector": ["sys", "query"], - "variable": "arg1" - }] + { + "id": "tt-2-source-if-else-target", + "source": "tt-2", + "target": "if-else", }, - "id": "tt", - },{ - "data": { - "iteration_id": "iteration-1", - "start_node_in_iteration": True, - "template": "{{ arg1 }} 321", - "title": "template transform", - "type": "template-transform", - "variables": [{ - "value_selector": ["sys", "query"], - "variable": "arg1" - }] + { + "id": "if-else-true-answer-2-target", + "source": "if-else", + "sourceHandle": "true", + "target": "answer-2", }, - "id": "tt-2", - }, { - "data": { - "answer": "{{#iteration-1.output#}}88888", - "title": "answer 3", - "type": "answer" + { + "id": "if-else-false-answer-4-target", + "source": "if-else", + "sourceHandle": "false", + "target": "answer-4", }, - "id": "answer-3", - }, { - "data": { - "conditions": [{ - "comparison_operator": "is", - "id": "1721916275284", - "value": "hi", - "variable_selector": ["sys", "query"] - }], - "iteration_id": "iteration-1", - "logical_operator": "and", - "title": "if", - "type": "if-else" + { + "id": "pe-source-iteration-1-target", + "source": "pe", + "target": "iteration-1", }, - "id": "if-else", - }, { - "data": { - "answer": "no hi", - "iteration_id": "iteration-1", - "title": "answer 4", - "type": "answer" - }, - "id": "answer-4", - }, { - "data": { - "instruction": "test1", - "model": { - "completion_params": { - "temperature": 0.7 - }, - "mode": "chat", - "name": "gpt-4o", - "provider": "openai" + ], + "nodes": [ + {"data": {"title": "Start", "type": "start", "variables": []}, "id": "start"}, + { + "data": { + "iterator_selector": ["pe", "list_output"], + "output_selector": ["tt", "output"], + "output_type": "array[string]", + "startNodeType": "template-transform", + "start_node_id": "tt", + "title": "iteration", + "type": "iteration", }, - "parameters": [{ - "description": "test", - "name": "list_output", - "required": False, - "type": "array[string]" - }], - "query": ["sys", "query"], - "reasoning_mode": "prompt", - "title": "pe", - "type": "parameter-extractor" + "id": "iteration-1", }, - "id": "pe", - }] + { + "data": { + "answer": "{{#tt.output#}}", + "iteration_id": "iteration-1", + "title": "answer 2", + "type": "answer", + }, + "id": "answer-2", + }, + { + "data": { + "iteration_id": "iteration-1", + "start_node_in_iteration": True, + "template": "{{ arg1 }} 123", + "title": "template transform", + "type": "template-transform", + "variables": [{"value_selector": ["sys", "query"], "variable": "arg1"}], + }, + "id": "tt", + }, + { + "data": { + "iteration_id": "iteration-1", + "start_node_in_iteration": True, + "template": "{{ arg1 }} 321", + "title": "template transform", + "type": "template-transform", + "variables": [{"value_selector": ["sys", "query"], "variable": "arg1"}], + }, + "id": "tt-2", + }, + { + "data": {"answer": "{{#iteration-1.output#}}88888", "title": "answer 3", "type": "answer"}, + "id": "answer-3", + }, + { + "data": { + "conditions": [ + { + "comparison_operator": "is", + "id": "1721916275284", + "value": "hi", + "variable_selector": ["sys", "query"], + } + ], + "iteration_id": "iteration-1", + "logical_operator": "and", + "title": "if", + "type": "if-else", + }, + "id": "if-else", + }, + { + "data": {"answer": "no hi", "iteration_id": "iteration-1", "title": "answer 4", "type": "answer"}, + "id": "answer-4", + }, + { + "data": { + "instruction": "test1", + "model": { + "completion_params": {"temperature": 0.7}, + "mode": "chat", + "name": "gpt-4o", + "provider": "openai", + }, + "parameters": [ + {"description": "test", "name": "list_output", "required": False, "type": "array[string]"} + ], + "query": ["sys", "query"], + "reasoning_mode": "prompt", + "title": "pe", + "type": "parameter-extractor", + }, + "id": "pe", + }, + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.CHAT, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) # construct variable pool - pool = VariablePool(system_variables={ - SystemVariableKey.QUERY: 'dify', - SystemVariableKey.FILES: [], - SystemVariableKey.CONVERSATION_ID: 'abababa', - SystemVariableKey.USER_ID: '1' - }, user_inputs={}, environment_variables=[]) - pool.add(['pe', 'list_output'], ["dify-1", "dify-2"]) + pool = VariablePool( + system_variables={ + SystemVariableKey.QUERY: "dify", + SystemVariableKey.FILES: [], + SystemVariableKey.CONVERSATION_ID: "abababa", + SystemVariableKey.USER_ID: "1", + }, + user_inputs={}, + environment_variables=[], + ) + pool.add(["pe", "list_output"], ["dify-1", "dify-2"]) iteration_node = IterationNode( id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=pool, start_at=time.perf_counter()), config={ "data": { "iterator_selector": ["pe", "list_output"], @@ -395,23 +376,19 @@ def test_run_parallel(): "type": "iteration", }, "id": "iteration-1", - } + }, ) def tt_generator(self): return NodeRunResult( status=WorkflowNodeExecutionStatus.SUCCEEDED, - inputs={ - 'iterator_selector': 'dify' - }, - outputs={ - 'output': 'dify 123' - } + inputs={"iterator_selector": "dify"}, + outputs={"output": "dify 123"}, ) # print("") - with patch.object(TemplateTransformNode, '_run', new=tt_generator): + with patch.object(TemplateTransformNode, "_run", new=tt_generator): # execute node result = iteration_node._run() diff --git a/api/tests/unit_tests/core/workflow/nodes/test_answer.py b/api/tests/unit_tests/core/workflow/nodes/test_answer.py index 06a1509bb3..cb2e99a854 100644 --- a/api/tests/unit_tests/core/workflow/nodes/test_answer.py +++ b/api/tests/unit_tests/core/workflow/nodes/test_answer.py @@ -24,60 +24,47 @@ def test_execute_answer(): }, ], "nodes": [ + {"data": {"type": "start"}, "id": "start"}, { "data": { - "type": "start" + "title": "123", + "type": "answer", + "answer": "Today's weather is {{#start.weather#}}\n{{#llm.text#}}\n{{img}}\nFin.", }, - "id": "start" + "id": "answer", }, - { - "data": { - 'title': '123', - 'type': 'answer', - 'answer': 'Today\'s weather is {{#start.weather#}}\n{{#llm.text#}}\n{{img}}\nFin.' - }, - "id": "answer" - }, - ] + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.WORKFLOW, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) # construct variable pool variable_pool = VariablePool( - system_variables={ - SystemVariableKey.FILES: [], - SystemVariableKey.USER_ID: 'aaa' - }, + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, user_inputs={}, environment_variables=[], conversation_variables=[], ) - variable_pool.add(['start', 'weather'], 'sunny') - variable_pool.add(['llm', 'text'], 'You are a helpful AI.') + variable_pool.add(["start", "weather"], "sunny") + variable_pool.add(["llm", "text"], "You are a helpful AI.") node = AnswerNode( id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=variable_pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), config={ "id": "answer", "data": { diff --git a/api/tests/unit_tests/core/workflow/nodes/test_if_else.py b/api/tests/unit_tests/core/workflow/nodes/test_if_else.py index c2bc90b355..0795f134d0 100644 --- a/api/tests/unit_tests/core/workflow/nodes/test_if_else.py +++ b/api/tests/unit_tests/core/workflow/nodes/test_if_else.py @@ -15,65 +15,49 @@ from models.workflow import WorkflowNodeExecutionStatus, WorkflowType def test_execute_if_else_result_true(): - graph_config = { - "edges": [], - "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - } - ] - } + graph_config = {"edges": [], "nodes": [{"data": {"type": "start"}, "id": "start"}]} - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.WORKFLOW, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) # construct variable pool - pool = VariablePool(system_variables={ - SystemVariableKey.FILES: [], - SystemVariableKey.USER_ID: 'aaa' - }, user_inputs={}) - pool.add(['start', 'array_contains'], ['ab', 'def']) - pool.add(['start', 'array_not_contains'], ['ac', 'def']) - pool.add(['start', 'contains'], 'cabcde') - pool.add(['start', 'not_contains'], 'zacde') - pool.add(['start', 'start_with'], 'abc') - pool.add(['start', 'end_with'], 'zzab') - pool.add(['start', 'is'], 'ab') - pool.add(['start', 'is_not'], 'aab') - pool.add(['start', 'empty'], '') - pool.add(['start', 'not_empty'], 'aaa') - pool.add(['start', 'equals'], 22) - pool.add(['start', 'not_equals'], 23) - pool.add(['start', 'greater_than'], 23) - pool.add(['start', 'less_than'], 21) - pool.add(['start', 'greater_than_or_equal'], 22) - pool.add(['start', 'less_than_or_equal'], 21) - pool.add(['start', 'not_null'], '1212') + pool = VariablePool( + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, user_inputs={} + ) + pool.add(["start", "array_contains"], ["ab", "def"]) + pool.add(["start", "array_not_contains"], ["ac", "def"]) + pool.add(["start", "contains"], "cabcde") + pool.add(["start", "not_contains"], "zacde") + pool.add(["start", "start_with"], "abc") + pool.add(["start", "end_with"], "zzab") + pool.add(["start", "is"], "ab") + pool.add(["start", "is_not"], "aab") + pool.add(["start", "empty"], "") + pool.add(["start", "not_empty"], "aaa") + pool.add(["start", "equals"], 22) + pool.add(["start", "not_equals"], 23) + pool.add(["start", "greater_than"], 23) + pool.add(["start", "less_than"], 21) + pool.add(["start", "greater_than_or_equal"], 22) + pool.add(["start", "less_than_or_equal"], 21) + pool.add(["start", "not_null"], "1212") node = IfElseNode( id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=pool, start_at=time.perf_counter()), config={ "id": "if-else", "data": { @@ -140,53 +124,44 @@ def test_execute_if_else_result_false(): }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "llm", }, - "id": "llm" + "id": "llm", }, - ] + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.WORKFLOW, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) # construct variable pool - pool = VariablePool(system_variables={ - SystemVariableKey.FILES: [], - SystemVariableKey.USER_ID: 'aaa' - }, user_inputs={}, environment_variables=[]) - pool.add(['start', 'array_contains'], ['1ab', 'def']) - pool.add(['start', 'array_not_contains'], ['ab', 'def']) + pool = VariablePool( + system_variables={SystemVariableKey.FILES: [], SystemVariableKey.USER_ID: "aaa"}, + user_inputs={}, + environment_variables=[], + ) + pool.add(["start", "array_contains"], ["1ab", "def"]) + pool.add(["start", "array_not_contains"], ["ab", "def"]) node = IfElseNode( id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=pool, start_at=time.perf_counter()), config={ "id": "if-else", "data": { diff --git a/api/tests/unit_tests/core/workflow/nodes/test_variable_assigner.py b/api/tests/unit_tests/core/workflow/nodes/test_variable_assigner.py index 61a715b2ee..f45a93f1be 100644 --- a/api/tests/unit_tests/core/workflow/nodes/test_variable_assigner.py +++ b/api/tests/unit_tests/core/workflow/nodes/test_variable_assigner.py @@ -27,35 +27,28 @@ def test_overwrite_string_variable(): }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "assigner", }, - "id": "assigner" + "id": "assigner", }, - ] + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.WORKFLOW, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) conversation_variable = StringVariable( @@ -87,21 +80,18 @@ def test_overwrite_string_variable(): id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=variable_pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), config={ - 'id': 'node_id', - 'data': { - 'assigned_variable_selector': ['conversation', conversation_variable.name], - 'write_mode': WriteMode.OVER_WRITE.value, - 'input_variable_selector': [DEFAULT_NODE_ID, input_variable.name], + "id": "node_id", + "data": { + "assigned_variable_selector": ["conversation", conversation_variable.name], + "write_mode": WriteMode.OVER_WRITE.value, + "input_variable_selector": [DEFAULT_NODE_ID, input_variable.name], }, }, ) - with mock.patch('core.workflow.nodes.variable_assigner.node.update_conversation_variable') as mock_run: + with mock.patch("core.workflow.nodes.variable_assigner.node.update_conversation_variable") as mock_run: list(node.run()) mock_run.assert_called_once() @@ -121,35 +111,28 @@ def test_append_variable_to_array(): }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "assigner", }, - "id": "assigner" + "id": "assigner", }, - ] + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.WORKFLOW, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) conversation_variable = ArrayStringVariable( @@ -179,21 +162,18 @@ def test_append_variable_to_array(): id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=variable_pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), config={ - 'id': 'node_id', - 'data': { - 'assigned_variable_selector': ['conversation', conversation_variable.name], - 'write_mode': WriteMode.APPEND.value, - 'input_variable_selector': [DEFAULT_NODE_ID, input_variable.name], + "id": "node_id", + "data": { + "assigned_variable_selector": ["conversation", conversation_variable.name], + "write_mode": WriteMode.APPEND.value, + "input_variable_selector": [DEFAULT_NODE_ID, input_variable.name], }, }, ) - with mock.patch('core.workflow.nodes.variable_assigner.node.update_conversation_variable') as mock_run: + with mock.patch("core.workflow.nodes.variable_assigner.node.update_conversation_variable") as mock_run: list(node.run()) mock_run.assert_called_once() @@ -212,35 +192,28 @@ def test_clear_array(): }, ], "nodes": [ - { - "data": { - "type": "start" - }, - "id": "start" - }, + {"data": {"type": "start"}, "id": "start"}, { "data": { "type": "assigner", }, - "id": "assigner" + "id": "assigner", }, - ] + ], } - graph = Graph.init( - graph_config=graph_config - ) + graph = Graph.init(graph_config=graph_config) init_params = GraphInitParams( - tenant_id='1', - app_id='1', + tenant_id="1", + app_id="1", workflow_type=WorkflowType.WORKFLOW, - workflow_id='1', + workflow_id="1", graph_config=graph_config, - user_id='1', + user_id="1", user_from=UserFrom.ACCOUNT, invoke_from=InvokeFrom.DEBUGGER, - call_depth=0 + call_depth=0, ) conversation_variable = ArrayStringVariable( @@ -250,7 +223,7 @@ def test_clear_array(): ) variable_pool = VariablePool( - system_variables={SystemVariableKey.CONVERSATION_ID: 'conversation_id'}, + system_variables={SystemVariableKey.CONVERSATION_ID: "conversation_id"}, user_inputs={}, environment_variables=[], conversation_variables=[conversation_variable], @@ -260,10 +233,7 @@ def test_clear_array(): id=str(uuid.uuid4()), graph_init_params=init_params, graph=graph, - graph_runtime_state=GraphRuntimeState( - variable_pool=variable_pool, - start_at=time.perf_counter() - ), + graph_runtime_state=GraphRuntimeState(variable_pool=variable_pool, start_at=time.perf_counter()), config={ "id": "node_id", "data": { @@ -274,7 +244,7 @@ def test_clear_array(): }, ) - with mock.patch('core.workflow.nodes.variable_assigner.node.update_conversation_variable') as mock_run: + with mock.patch("core.workflow.nodes.variable_assigner.node.update_conversation_variable") as mock_run: list(node.run()) mock_run.assert_called_once()