From 8c3d835b45e3e0cf64e0861c1d33aee4247649ec Mon Sep 17 00:00:00 2001 From: -LAN- Date: Wed, 30 Apr 2025 15:02:23 +0800 Subject: [PATCH] feat(api/variable_assigner_v2): Add operations to remove first/last elements from arrays Introduces 'remove-first' and 'remove-last' operations for array variables, allowing for removal of the first or last element respectively. Ensures these operations are supported only for array types. Includes unit tests to verify the correct behavior when applied to arrays, including edge cases with empty arrays. Signed-off-by: -LAN- --- .../nodes/variable_assigner/v2/enums.py | 2 + .../nodes/variable_assigner/v2/helpers.py | 11 +- .../nodes/variable_assigner/v2/node.py | 12 +- .../nodes/variable_assigner/v2/__init__.py | 1 + .../v2/test_variable_assigner_v2.py | 390 ++++++++++++++++++ 5 files changed, 414 insertions(+), 2 deletions(-) create mode 100644 api/tests/unit_tests/core/workflow/nodes/variable_assigner/v2/__init__.py create mode 100644 api/tests/unit_tests/core/workflow/nodes/variable_assigner/v2/test_variable_assigner_v2.py diff --git a/api/core/workflow/nodes/variable_assigner/v2/enums.py b/api/core/workflow/nodes/variable_assigner/v2/enums.py index 36cf68aa19..291b1208d4 100644 --- a/api/core/workflow/nodes/variable_assigner/v2/enums.py +++ b/api/core/workflow/nodes/variable_assigner/v2/enums.py @@ -11,6 +11,8 @@ class Operation(StrEnum): SUBTRACT = "-=" MULTIPLY = "*=" DIVIDE = "/=" + REMOVE_FIRST = "remove-first" + REMOVE_LAST = "remove-last" class InputType(StrEnum): diff --git a/api/core/workflow/nodes/variable_assigner/v2/helpers.py b/api/core/workflow/nodes/variable_assigner/v2/helpers.py index a86c7eb94a..8fb2a27388 100644 --- a/api/core/workflow/nodes/variable_assigner/v2/helpers.py +++ b/api/core/workflow/nodes/variable_assigner/v2/helpers.py @@ -23,6 +23,15 @@ def is_operation_supported(*, variable_type: SegmentType, operation: Operation): SegmentType.ARRAY_NUMBER, SegmentType.ARRAY_FILE, } + case Operation.REMOVE_FIRST | Operation.REMOVE_LAST: + # Only array variable can have elements removed + return variable_type in { + SegmentType.ARRAY_ANY, + SegmentType.ARRAY_OBJECT, + SegmentType.ARRAY_STRING, + SegmentType.ARRAY_NUMBER, + SegmentType.ARRAY_FILE, + } case _: return False @@ -51,7 +60,7 @@ def is_constant_input_supported(*, variable_type: SegmentType, operation: Operat def is_input_value_valid(*, variable_type: SegmentType, operation: Operation, value: Any): - if operation == Operation.CLEAR: + if operation in {Operation.CLEAR, Operation.REMOVE_FIRST, Operation.REMOVE_LAST}: return True match variable_type: case SegmentType.STRING: diff --git a/api/core/workflow/nodes/variable_assigner/v2/node.py b/api/core/workflow/nodes/variable_assigner/v2/node.py index 0305eb7f41..6a7ad86b51 100644 --- a/api/core/workflow/nodes/variable_assigner/v2/node.py +++ b/api/core/workflow/nodes/variable_assigner/v2/node.py @@ -64,7 +64,7 @@ class VariableAssignerNode(BaseNode[VariableAssignerNodeData]): # Get value from variable pool if ( item.input_type == InputType.VARIABLE - and item.operation != Operation.CLEAR + and item.operation not in {Operation.CLEAR, Operation.REMOVE_FIRST, Operation.REMOVE_LAST} and item.value is not None ): value = self.graph_runtime_state.variable_pool.get(item.value) @@ -165,5 +165,15 @@ class VariableAssignerNode(BaseNode[VariableAssignerNodeData]): return variable.value * value case Operation.DIVIDE: return variable.value / value + case Operation.REMOVE_FIRST: + # If array is empty, do nothing + if not variable.value: + return variable.value + return variable.value[1:] + case Operation.REMOVE_LAST: + # If array is empty, do nothing + if not variable.value: + return variable.value + return variable.value[:-1] case _: raise OperationNotSupportedError(operation=operation, variable_type=variable.value_type) diff --git a/api/tests/unit_tests/core/workflow/nodes/variable_assigner/v2/__init__.py b/api/tests/unit_tests/core/workflow/nodes/variable_assigner/v2/__init__.py new file mode 100644 index 0000000000..8b13789179 --- /dev/null +++ b/api/tests/unit_tests/core/workflow/nodes/variable_assigner/v2/__init__.py @@ -0,0 +1 @@ + diff --git a/api/tests/unit_tests/core/workflow/nodes/variable_assigner/v2/test_variable_assigner_v2.py b/api/tests/unit_tests/core/workflow/nodes/variable_assigner/v2/test_variable_assigner_v2.py new file mode 100644 index 0000000000..7c5597dd89 --- /dev/null +++ b/api/tests/unit_tests/core/workflow/nodes/variable_assigner/v2/test_variable_assigner_v2.py @@ -0,0 +1,390 @@ +import time +import uuid +from uuid import uuid4 + +from core.app.entities.app_invoke_entities import InvokeFrom +from core.variables import ArrayStringVariable +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.variable_assigner.v2 import VariableAssignerNode +from core.workflow.nodes.variable_assigner.v2.enums import InputType, Operation +from models.enums import UserFrom +from models.workflow import WorkflowType + +DEFAULT_NODE_ID = "node_id" + + +def test_handle_item_directly(): + """Test the _handle_item method directly for remove operations.""" + # Create variables + variable1 = ArrayStringVariable( + id=str(uuid4()), + name="test_variable1", + value=["first", "second", "third"], + ) + + variable2 = ArrayStringVariable( + id=str(uuid4()), + name="test_variable2", + value=["first", "second", "third"], + ) + + # Create a mock class with just the _handle_item method + class MockNode: + def _handle_item(self, *, variable, operation, value): + match operation: + case Operation.REMOVE_FIRST: + if not variable.value: + return variable.value + return variable.value[1:] + case Operation.REMOVE_LAST: + if not variable.value: + return variable.value + return variable.value[:-1] + + node = MockNode() + + # Test remove-first + result1 = node._handle_item( + variable=variable1, + operation=Operation.REMOVE_FIRST, + value=None, + ) + + # Test remove-last + result2 = node._handle_item( + variable=variable2, + operation=Operation.REMOVE_LAST, + value=None, + ) + + # Check the results + assert result1 == ["second", "third"] + assert result2 == ["first", "second"] + + +def test_remove_first_from_array(): + """Test removing the first element from an array.""" + graph_config = { + "edges": [ + { + "id": "start-source-assigner-target", + "source": "start", + "target": "assigner", + }, + ], + "nodes": [ + {"data": {"type": "start"}, "id": "start"}, + { + "data": { + "type": "assigner", + }, + "id": "assigner", + }, + ], + } + + 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, + ) + + conversation_variable = ArrayStringVariable( + id=str(uuid4()), + name="test_conversation_variable", + value=["first", "second", "third"], + selector=["conversation", "test_conversation_variable"], + ) + + variable_pool = VariablePool( + system_variables={SystemVariableKey.CONVERSATION_ID: "conversation_id"}, + user_inputs={}, + environment_variables=[], + conversation_variables=[conversation_variable], + ) + + node = VariableAssignerNode( + 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={ + "id": "node_id", + "data": { + "title": "test", + "version": "2", + "items": [ + { + "variable_selector": ["conversation", conversation_variable.name], + "input_type": InputType.VARIABLE, + "operation": Operation.REMOVE_FIRST, + "value": None, + } + ], + }, + }, + ) + + # Skip the mock assertion since we're in a test environment + # Print the variable before running + print(f"Before: {variable_pool.get(['conversation', conversation_variable.name]).to_object()}") + + # Run the node + result = list(node.run()) + + # Print the variable after running and the result + print(f"After: {variable_pool.get(['conversation', conversation_variable.name]).to_object()}") + print(f"Result: {result}") + + got = variable_pool.get(["conversation", conversation_variable.name]) + assert got is not None + assert got.to_object() == ["second", "third"] + + +def test_remove_last_from_array(): + """Test removing the last element from an array.""" + graph_config = { + "edges": [ + { + "id": "start-source-assigner-target", + "source": "start", + "target": "assigner", + }, + ], + "nodes": [ + {"data": {"type": "start"}, "id": "start"}, + { + "data": { + "type": "assigner", + }, + "id": "assigner", + }, + ], + } + + 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, + ) + + conversation_variable = ArrayStringVariable( + id=str(uuid4()), + name="test_conversation_variable", + value=["first", "second", "third"], + selector=["conversation", "test_conversation_variable"], + ) + + variable_pool = VariablePool( + system_variables={SystemVariableKey.CONVERSATION_ID: "conversation_id"}, + user_inputs={}, + environment_variables=[], + conversation_variables=[conversation_variable], + ) + + node = VariableAssignerNode( + 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={ + "id": "node_id", + "data": { + "title": "test", + "version": "2", + "items": [ + { + "variable_selector": ["conversation", conversation_variable.name], + "input_type": InputType.VARIABLE, + "operation": Operation.REMOVE_LAST, + "value": None, + } + ], + }, + }, + ) + + # Skip the mock assertion since we're in a test environment + list(node.run()) + + got = variable_pool.get(["conversation", conversation_variable.name]) + assert got is not None + assert got.to_object() == ["first", "second"] + + +def test_remove_first_from_empty_array(): + """Test removing the first element from an empty array (should do nothing).""" + graph_config = { + "edges": [ + { + "id": "start-source-assigner-target", + "source": "start", + "target": "assigner", + }, + ], + "nodes": [ + {"data": {"type": "start"}, "id": "start"}, + { + "data": { + "type": "assigner", + }, + "id": "assigner", + }, + ], + } + + 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, + ) + + conversation_variable = ArrayStringVariable( + id=str(uuid4()), + name="test_conversation_variable", + value=[], + selector=["conversation", "test_conversation_variable"], + ) + + variable_pool = VariablePool( + system_variables={SystemVariableKey.CONVERSATION_ID: "conversation_id"}, + user_inputs={}, + environment_variables=[], + conversation_variables=[conversation_variable], + ) + + node = VariableAssignerNode( + 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={ + "id": "node_id", + "data": { + "title": "test", + "version": "2", + "items": [ + { + "variable_selector": ["conversation", conversation_variable.name], + "input_type": InputType.VARIABLE, + "operation": Operation.REMOVE_FIRST, + "value": None, + } + ], + }, + }, + ) + + # Skip the mock assertion since we're in a test environment + list(node.run()) + + got = variable_pool.get(["conversation", conversation_variable.name]) + assert got is not None + assert got.to_object() == [] + + +def test_remove_last_from_empty_array(): + """Test removing the last element from an empty array (should do nothing).""" + graph_config = { + "edges": [ + { + "id": "start-source-assigner-target", + "source": "start", + "target": "assigner", + }, + ], + "nodes": [ + {"data": {"type": "start"}, "id": "start"}, + { + "data": { + "type": "assigner", + }, + "id": "assigner", + }, + ], + } + + 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, + ) + + conversation_variable = ArrayStringVariable( + id=str(uuid4()), + name="test_conversation_variable", + value=[], + selector=["conversation", "test_conversation_variable"], + ) + + variable_pool = VariablePool( + system_variables={SystemVariableKey.CONVERSATION_ID: "conversation_id"}, + user_inputs={}, + environment_variables=[], + conversation_variables=[conversation_variable], + ) + + node = VariableAssignerNode( + 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={ + "id": "node_id", + "data": { + "title": "test", + "version": "2", + "items": [ + { + "variable_selector": ["conversation", conversation_variable.name], + "input_type": InputType.VARIABLE, + "operation": Operation.REMOVE_LAST, + "value": None, + } + ], + }, + }, + ) + + # Skip the mock assertion since we're in a test environment + list(node.run()) + + got = variable_pool.get(["conversation", conversation_variable.name]) + assert got is not None + assert got.to_object() == []