mirror of
https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git
synced 2025-08-15 12:45:52 +08:00
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- <laipz8200@outlook.com>
This commit is contained in:
-LAN-
parent
62c412f26d
commit
8c3d835b45
@ -11,6 +11,8 @@ class Operation(StrEnum):
|
||||
SUBTRACT = "-="
|
||||
MULTIPLY = "*="
|
||||
DIVIDE = "/="
|
||||
REMOVE_FIRST = "remove-first"
|
||||
REMOVE_LAST = "remove-last"
|
||||
|
||||
|
||||
class InputType(StrEnum):
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -0,0 +1 @@
|
||||
|
||||
@ -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() == []
|
||||
Loading…
x
Reference in New Issue
Block a user