Feat: upgrade variable assigner (#11285)

Signed-off-by: -LAN- <laipz8200@outlook.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
Yi Xiao
2024-12-03 13:56:40 +08:00
committed by GitHub
parent e79eac688a
commit e135ffc2c1
62 changed files with 1565 additions and 301 deletions

View File

@@ -1,4 +1,4 @@
from core.rag.datasource.vdb.analyticdb.analyticdb_vector import AnalyticdbConfig, AnalyticdbVector
from core.rag.datasource.vdb.analyticdb.analyticdb_vector import AnalyticdbVector
from core.rag.datasource.vdb.analyticdb.analyticdb_vector_openapi import AnalyticdbVectorOpenAPIConfig
from core.rag.datasource.vdb.analyticdb.analyticdb_vector_sql import AnalyticdbVectorBySqlConfig
from tests.integration_tests.vdb.test_vector_store import AbstractVectorTest, setup_mock_redis

View File

@@ -19,36 +19,36 @@ from factories import variable_factory
def test_string_variable():
test_data = {"value_type": "string", "name": "test_text", "value": "Hello, World!"}
result = variable_factory.build_variable_from_mapping(test_data)
result = variable_factory.build_conversation_variable_from_mapping(test_data)
assert isinstance(result, StringVariable)
def test_integer_variable():
test_data = {"value_type": "number", "name": "test_int", "value": 42}
result = variable_factory.build_variable_from_mapping(test_data)
result = variable_factory.build_conversation_variable_from_mapping(test_data)
assert isinstance(result, IntegerVariable)
def test_float_variable():
test_data = {"value_type": "number", "name": "test_float", "value": 3.14}
result = variable_factory.build_variable_from_mapping(test_data)
result = variable_factory.build_conversation_variable_from_mapping(test_data)
assert isinstance(result, FloatVariable)
def test_secret_variable():
test_data = {"value_type": "secret", "name": "test_secret", "value": "secret_value"}
result = variable_factory.build_variable_from_mapping(test_data)
result = variable_factory.build_conversation_variable_from_mapping(test_data)
assert isinstance(result, SecretVariable)
def test_invalid_value_type():
test_data = {"value_type": "unknown", "name": "test_invalid", "value": "value"}
with pytest.raises(VariableError):
variable_factory.build_variable_from_mapping(test_data)
variable_factory.build_conversation_variable_from_mapping(test_data)
def test_build_a_blank_string():
result = variable_factory.build_variable_from_mapping(
result = variable_factory.build_conversation_variable_from_mapping(
{
"value_type": "string",
"name": "blank",
@@ -80,7 +80,7 @@ def test_object_variable():
"key2": 2,
},
}
variable = variable_factory.build_variable_from_mapping(mapping)
variable = variable_factory.build_conversation_variable_from_mapping(mapping)
assert isinstance(variable, ObjectSegment)
assert isinstance(variable.value["key1"], str)
assert isinstance(variable.value["key2"], int)
@@ -97,7 +97,7 @@ def test_array_string_variable():
"text",
],
}
variable = variable_factory.build_variable_from_mapping(mapping)
variable = variable_factory.build_conversation_variable_from_mapping(mapping)
assert isinstance(variable, ArrayStringVariable)
assert isinstance(variable.value[0], str)
assert isinstance(variable.value[1], str)
@@ -114,7 +114,7 @@ def test_array_number_variable():
2.0,
],
}
variable = variable_factory.build_variable_from_mapping(mapping)
variable = variable_factory.build_conversation_variable_from_mapping(mapping)
assert isinstance(variable, ArrayNumberVariable)
assert isinstance(variable.value[0], int)
assert isinstance(variable.value[1], float)
@@ -137,7 +137,7 @@ def test_array_object_variable():
},
],
}
variable = variable_factory.build_variable_from_mapping(mapping)
variable = variable_factory.build_conversation_variable_from_mapping(mapping)
assert isinstance(variable, ArrayObjectVariable)
assert isinstance(variable.value[0], dict)
assert isinstance(variable.value[1], dict)
@@ -149,7 +149,7 @@ def test_array_object_variable():
def test_variable_cannot_large_than_200_kb():
with pytest.raises(VariableError):
variable_factory.build_variable_from_mapping(
variable_factory.build_conversation_variable_from_mapping(
{
"id": str(uuid4()),
"value_type": "string",

View File

@@ -10,7 +10,8 @@ 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.variable_assigner import VariableAssignerNode, WriteMode
from core.workflow.nodes.variable_assigner.v1 import VariableAssignerNode
from core.workflow.nodes.variable_assigner.v1.node_data import WriteMode
from models.enums import UserFrom
from models.workflow import WorkflowType
@@ -84,6 +85,7 @@ def test_overwrite_string_variable():
config={
"id": "node_id",
"data": {
"title": "test",
"assigned_variable_selector": ["conversation", conversation_variable.name],
"write_mode": WriteMode.OVER_WRITE.value,
"input_variable_selector": [DEFAULT_NODE_ID, input_variable.name],
@@ -91,7 +93,7 @@ def test_overwrite_string_variable():
},
)
with mock.patch("core.workflow.nodes.variable_assigner.node.update_conversation_variable") as mock_run:
with mock.patch("core.workflow.nodes.variable_assigner.common.helpers.update_conversation_variable") as mock_run:
list(node.run())
mock_run.assert_called_once()
@@ -166,6 +168,7 @@ def test_append_variable_to_array():
config={
"id": "node_id",
"data": {
"title": "test",
"assigned_variable_selector": ["conversation", conversation_variable.name],
"write_mode": WriteMode.APPEND.value,
"input_variable_selector": [DEFAULT_NODE_ID, input_variable.name],
@@ -173,7 +176,7 @@ def test_append_variable_to_array():
},
)
with mock.patch("core.workflow.nodes.variable_assigner.node.update_conversation_variable") as mock_run:
with mock.patch("core.workflow.nodes.variable_assigner.common.helpers.update_conversation_variable") as mock_run:
list(node.run())
mock_run.assert_called_once()
@@ -237,6 +240,7 @@ def test_clear_array():
config={
"id": "node_id",
"data": {
"title": "test",
"assigned_variable_selector": ["conversation", conversation_variable.name],
"write_mode": WriteMode.CLEAR.value,
"input_variable_selector": [],
@@ -244,7 +248,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.common.helpers.update_conversation_variable") as mock_run:
list(node.run())
mock_run.assert_called_once()

View File

@@ -0,0 +1,24 @@
import pytest
from core.variables import SegmentType
from core.workflow.nodes.variable_assigner.v2.enums import Operation
from core.workflow.nodes.variable_assigner.v2.helpers import is_input_value_valid
def test_is_input_value_valid_overwrite_array_string():
# Valid cases
assert is_input_value_valid(
variable_type=SegmentType.ARRAY_STRING, operation=Operation.OVER_WRITE, value=["hello", "world"]
)
assert is_input_value_valid(variable_type=SegmentType.ARRAY_STRING, operation=Operation.OVER_WRITE, value=[])
# Invalid cases
assert not is_input_value_valid(
variable_type=SegmentType.ARRAY_STRING, operation=Operation.OVER_WRITE, value="not an array"
)
assert not is_input_value_valid(
variable_type=SegmentType.ARRAY_STRING, operation=Operation.OVER_WRITE, value=[1, 2, 3]
)
assert not is_input_value_valid(
variable_type=SegmentType.ARRAY_STRING, operation=Operation.OVER_WRITE, value=["valid", 123, "invalid"]
)

View File

@@ -6,7 +6,7 @@ from models import ConversationVariable
def test_from_variable_and_to_variable():
variable = variable_factory.build_variable_from_mapping(
variable = variable_factory.build_conversation_variable_from_mapping(
{
"id": str(uuid4()),
"name": "name",

View File

@@ -24,10 +24,18 @@ def test_environment_variables():
)
# Create some EnvironmentVariable instances
variable1 = StringVariable.model_validate({"name": "var1", "value": "value1", "id": str(uuid4())})
variable2 = IntegerVariable.model_validate({"name": "var2", "value": 123, "id": str(uuid4())})
variable3 = SecretVariable.model_validate({"name": "var3", "value": "secret", "id": str(uuid4())})
variable4 = FloatVariable.model_validate({"name": "var4", "value": 3.14, "id": str(uuid4())})
variable1 = StringVariable.model_validate(
{"name": "var1", "value": "value1", "id": str(uuid4()), "selector": ["env", "var1"]}
)
variable2 = IntegerVariable.model_validate(
{"name": "var2", "value": 123, "id": str(uuid4()), "selector": ["env", "var2"]}
)
variable3 = SecretVariable.model_validate(
{"name": "var3", "value": "secret", "id": str(uuid4()), "selector": ["env", "var3"]}
)
variable4 = FloatVariable.model_validate(
{"name": "var4", "value": 3.14, "id": str(uuid4()), "selector": ["env", "var4"]}
)
with (
mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"),
@@ -58,10 +66,18 @@ def test_update_environment_variables():
)
# Create some EnvironmentVariable instances
variable1 = StringVariable.model_validate({"name": "var1", "value": "value1", "id": str(uuid4())})
variable2 = IntegerVariable.model_validate({"name": "var2", "value": 123, "id": str(uuid4())})
variable3 = SecretVariable.model_validate({"name": "var3", "value": "secret", "id": str(uuid4())})
variable4 = FloatVariable.model_validate({"name": "var4", "value": 3.14, "id": str(uuid4())})
variable1 = StringVariable.model_validate(
{"name": "var1", "value": "value1", "id": str(uuid4()), "selector": ["env", "var1"]}
)
variable2 = IntegerVariable.model_validate(
{"name": "var2", "value": 123, "id": str(uuid4()), "selector": ["env", "var2"]}
)
variable3 = SecretVariable.model_validate(
{"name": "var3", "value": "secret", "id": str(uuid4()), "selector": ["env", "var3"]}
)
variable4 = FloatVariable.model_validate(
{"name": "var4", "value": 3.14, "id": str(uuid4()), "selector": ["env", "var4"]}
)
with (
mock.patch("core.helper.encrypter.encrypt_token", return_value="encrypted_token"),