This commit is contained in:
takatost 2024-08-25 22:02:21 +08:00
parent 4771e85630
commit 6c61776ee1
16 changed files with 1304 additions and 1445 deletions

View File

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

View File

@ -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)

View File

@ -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")

View File

@ -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)

View File

@ -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:

View File

@ -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"

View File

@ -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"] == []

View File

@ -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):

View File

@ -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))

View File

@ -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."

View File

@ -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"] == []

View File

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

View File

@ -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()

View File

@ -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": {

View File

@ -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": {

View File

@ -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()