From bcc95e520b4570789d5037959848c0ab5bfd5388 Mon Sep 17 00:00:00 2001 From: -LAN- Date: Wed, 30 Apr 2025 15:50:00 +0800 Subject: [PATCH] feat: support remove first and remove last in variable assigner (#19144) 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 ++++++++++++++++++ .../assigner/components/var-list/index.tsx | 1 + .../workflow/nodes/assigner/default.ts | 2 +- .../workflow/nodes/assigner/types.ts | 2 + .../workflow/nodes/assigner/use-config.ts | 2 +- web/i18n/en-US/workflow.ts | 2 + web/i18n/zh-Hans/workflow.ts | 2 + web/i18n/zh-Hant/workflow.ts | 2 + 12 files changed, 425 insertions(+), 4 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() == [] diff --git a/web/app/components/workflow/nodes/assigner/components/var-list/index.tsx b/web/app/components/workflow/nodes/assigner/components/var-list/index.tsx index 170e6484c4..f34a1435ad 100644 --- a/web/app/components/workflow/nodes/assigner/components/var-list/index.tsx +++ b/web/app/components/workflow/nodes/assigner/components/var-list/index.tsx @@ -152,6 +152,7 @@ const VarList: FC = ({ /> {item.operation !== WriteMode.clear && item.operation !== WriteMode.set + && item.operation !== WriteMode.removeFirst && item.operation !== WriteMode.removeLast && !writeModeTypesNum?.includes(item.operation) && ( = { if (!errorMessages && !value.variable_selector?.length) errorMessages = t(`${i18nPrefix}.fieldRequired`, { field: t('workflow.nodes.assigner.assignedVariable') }) - if (!errorMessages && value.operation !== WriteMode.clear) { + if (!errorMessages && value.operation !== WriteMode.clear && value.operation !== WriteMode.removeFirst && value.operation !== WriteMode.removeLast) { if (value.operation === WriteMode.set || value.operation === WriteMode.increment || value.operation === WriteMode.decrement || value.operation === WriteMode.multiply || value.operation === WriteMode.divide) { diff --git a/web/app/components/workflow/nodes/assigner/types.ts b/web/app/components/workflow/nodes/assigner/types.ts index 9f06f65dc3..85d2b2850f 100644 --- a/web/app/components/workflow/nodes/assigner/types.ts +++ b/web/app/components/workflow/nodes/assigner/types.ts @@ -10,6 +10,8 @@ export enum WriteMode { decrement = '-=', multiply = '*=', divide = '/=', + removeFirst = 'remove-first', + removeLast = 'remove-last', } export enum AssignerNodeInputType { diff --git a/web/app/components/workflow/nodes/assigner/use-config.ts b/web/app/components/workflow/nodes/assigner/use-config.ts index e7beb1f37a..cbd5475483 100644 --- a/web/app/components/workflow/nodes/assigner/use-config.ts +++ b/web/app/components/workflow/nodes/assigner/use-config.ts @@ -69,7 +69,7 @@ const useConfig = (id: string, rawPayload: AssignerNodeType) => { newSetInputs(newInputs) }, [inputs, newSetInputs]) - const writeModeTypesArr = [WriteMode.overwrite, WriteMode.clear, WriteMode.append, WriteMode.extend] + const writeModeTypesArr = [WriteMode.overwrite, WriteMode.clear, WriteMode.append, WriteMode.extend, WriteMode.removeFirst, WriteMode.removeLast] const writeModeTypes = [WriteMode.overwrite, WriteMode.clear, WriteMode.set] const writeModeTypesNum = [WriteMode.increment, WriteMode.decrement, WriteMode.multiply, WriteMode.divide] diff --git a/web/i18n/en-US/workflow.ts b/web/i18n/en-US/workflow.ts index 543f689611..ab0c6a5879 100644 --- a/web/i18n/en-US/workflow.ts +++ b/web/i18n/en-US/workflow.ts @@ -638,6 +638,8 @@ const translation = { 'clear': 'Clear', 'extend': 'Extend', 'append': 'Append', + 'remove-first': 'Remove First', + 'remove-last': 'Remove Last', '+=': '+=', '-=': '-=', '*=': '*=', diff --git a/web/i18n/zh-Hans/workflow.ts b/web/i18n/zh-Hans/workflow.ts index ab56c468ce..9f8d2c6964 100644 --- a/web/i18n/zh-Hans/workflow.ts +++ b/web/i18n/zh-Hans/workflow.ts @@ -638,6 +638,8 @@ const translation = { 'clear': '清空', 'extend': '扩展', 'append': '追加', + 'remove-first': '移除首项', + 'remove-last': '移除末项', '+=': '+=', '-=': '-=', '*=': '*=', diff --git a/web/i18n/zh-Hant/workflow.ts b/web/i18n/zh-Hant/workflow.ts index 1e4fd2ef21..0930316047 100644 --- a/web/i18n/zh-Hant/workflow.ts +++ b/web/i18n/zh-Hant/workflow.ts @@ -564,6 +564,8 @@ const translation = { '-=': '-=', 'append': '附加', 'clear': '清除', + 'remove-first': '移除首項', + 'remove-last': '移除末項', }, 'noAssignedVars': '沒有可用的已分配變數', 'variables': '變數',