From 0c709afe5c920322a8646d20e204d0931377af9f Mon Sep 17 00:00:00 2001
From: takatost <takatost@gmail.com>
Date: Wed, 13 Mar 2024 17:10:51 +0800
Subject: [PATCH] add if-else node

---
 api/core/workflow/entities/variable_pool.py   |   2 +-
 api/core/workflow/nodes/if_else/entities.py   |  26 ++
 .../workflow/nodes/if_else/if_else_node.py    | 395 +++++++++++++++++-
 .../core/workflow/nodes/if_else_node.py       | 193 +++++++++
 4 files changed, 614 insertions(+), 2 deletions(-)
 create mode 100644 api/core/workflow/nodes/if_else/entities.py
 create mode 100644 api/tests/unit_tests/core/workflow/nodes/if_else_node.py

diff --git a/api/core/workflow/entities/variable_pool.py b/api/core/workflow/entities/variable_pool.py
index 7a5f58d808..ff96bc3bac 100644
--- a/api/core/workflow/entities/variable_pool.py
+++ b/api/core/workflow/entities/variable_pool.py
@@ -86,6 +86,6 @@ class VariablePool:
                                        ValueType.ARRAY_OBJECT,
                                        ValueType.ARRAY_FILE]:
                 if not isinstance(value, list):
-                    raise ValueError('Invalid value type: array')
+                    raise ValueError(f'Invalid value type: {target_value_type.value}')
 
         return value
diff --git a/api/core/workflow/nodes/if_else/entities.py b/api/core/workflow/nodes/if_else/entities.py
new file mode 100644
index 0000000000..68d51c93be
--- /dev/null
+++ b/api/core/workflow/nodes/if_else/entities.py
@@ -0,0 +1,26 @@
+from typing import Literal, Optional
+
+from pydantic import BaseModel
+
+from core.workflow.entities.base_node_data_entities import BaseNodeData
+
+
+class IfElseNodeData(BaseNodeData):
+    """
+    Answer Node Data.
+    """
+    class Condition(BaseModel):
+        """
+        Condition entity
+        """
+        variable_selector: list[str]
+        comparison_operator: Literal[
+            # for string or array
+            "contains", "not contains", "start with", "end with", "is", "is not", "empty", "not empty",
+            # for number
+            "=", "≠", ">", "<", "≥", "≤", "null", "not null"
+        ]
+        value: Optional[str] = None
+
+    logical_operator: Literal["and", "or"] = "and"
+    conditions: list[Condition]
diff --git a/api/core/workflow/nodes/if_else/if_else_node.py b/api/core/workflow/nodes/if_else/if_else_node.py
index 98a5c85db2..9cb084b116 100644
--- a/api/core/workflow/nodes/if_else/if_else_node.py
+++ b/api/core/workflow/nodes/if_else/if_else_node.py
@@ -1,5 +1,398 @@
+from typing import Optional, cast
+
+from core.workflow.entities.base_node_data_entities import BaseNodeData
+from core.workflow.entities.node_entities import NodeRunResult, NodeType
+from core.workflow.entities.variable_pool import VariablePool
 from core.workflow.nodes.base_node import BaseNode
+from core.workflow.nodes.if_else.entities import IfElseNodeData
+from models.workflow import WorkflowNodeExecutionStatus
 
 
 class IfElseNode(BaseNode):
-    pass
+    _node_data_cls = IfElseNodeData
+    node_type = NodeType.IF_ELSE
+
+    def _run(self, variable_pool: VariablePool) -> NodeRunResult:
+        """
+        Run node
+        :param variable_pool: variable pool
+        :return:
+        """
+        node_data = self.node_data
+        node_data = cast(self._node_data_cls, node_data)
+
+        node_inputs = {
+            "conditions": []
+        }
+
+        process_datas = {
+            "condition_results": []
+        }
+
+        try:
+            logical_operator = node_data.logical_operator
+            input_conditions = []
+            for condition in node_data.conditions:
+                actual_value = variable_pool.get_variable_value(
+                    variable_selector=condition.variable_selector
+                )
+
+                expected_value = condition.value
+
+                input_conditions.append({
+                    "actual_value": actual_value,
+                    "expected_value": expected_value,
+                    "comparison_operator": condition.comparison_operator
+                })
+
+            node_inputs["conditions"] = input_conditions
+
+            for input_condition in input_conditions:
+                actual_value = input_condition["actual_value"]
+                expected_value = input_condition["expected_value"]
+                comparison_operator = input_condition["comparison_operator"]
+
+                if comparison_operator == "contains":
+                    compare_result = self._assert_contains(actual_value, expected_value)
+                elif comparison_operator == "not contains":
+                    compare_result = self._assert_not_contains(actual_value, expected_value)
+                elif comparison_operator == "start with":
+                    compare_result = self._assert_start_with(actual_value, expected_value)
+                elif comparison_operator == "end with":
+                    compare_result = self._assert_end_with(actual_value, expected_value)
+                elif comparison_operator == "is":
+                    compare_result = self._assert_is(actual_value, expected_value)
+                elif comparison_operator == "is not":
+                    compare_result = self._assert_is_not(actual_value, expected_value)
+                elif comparison_operator == "empty":
+                    compare_result = self._assert_empty(actual_value)
+                elif comparison_operator == "not empty":
+                    compare_result = self._assert_not_empty(actual_value)
+                elif comparison_operator == "=":
+                    compare_result = self._assert_equal(actual_value, expected_value)
+                elif comparison_operator == "≠":
+                    compare_result = self._assert_not_equal(actual_value, expected_value)
+                elif comparison_operator == ">":
+                    compare_result = self._assert_greater_than(actual_value, expected_value)
+                elif comparison_operator == "<":
+                    compare_result = self._assert_less_than(actual_value, expected_value)
+                elif comparison_operator == "≥":
+                    compare_result = self._assert_greater_than_or_equal(actual_value, expected_value)
+                elif comparison_operator == "≤":
+                    compare_result = self._assert_less_than_or_equal(actual_value, expected_value)
+                elif comparison_operator == "null":
+                    compare_result = self._assert_null(actual_value)
+                elif comparison_operator == "not null":
+                    compare_result = self._assert_not_null(actual_value)
+                else:
+                    continue
+
+                process_datas["condition_results"].append({
+                    **input_condition,
+                    "result": compare_result
+                })
+        except Exception as e:
+            return NodeRunResult(
+                status=WorkflowNodeExecutionStatus.FAILED,
+                inputs=node_inputs,
+                process_datas=process_datas,
+                error=str(e)
+            )
+
+        if logical_operator == "and":
+            compare_result = False not in [condition["result"] for condition in process_datas["condition_results"]]
+        else:
+            compare_result = True in [condition["result"] for condition in process_datas["condition_results"]]
+
+        return NodeRunResult(
+            status=WorkflowNodeExecutionStatus.SUCCEEDED,
+            inputs=node_inputs,
+            process_datas=process_datas,
+            edge_source_handle="false" if not compare_result else "true",
+            outputs={
+                "result": compare_result
+            }
+        )
+
+    def _assert_contains(self, actual_value: Optional[str | list], expected_value: str) -> bool:
+        """
+        Assert contains
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if not actual_value:
+            return False
+
+        if not isinstance(actual_value, str | list):
+            raise ValueError('Invalid actual value type: string or array')
+
+        if expected_value not in actual_value:
+            return False
+        return True
+
+    def _assert_not_contains(self, actual_value: Optional[str | list], expected_value: str) -> bool:
+        """
+        Assert not contains
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if not actual_value:
+            return True
+
+        if not isinstance(actual_value, str | list):
+            raise ValueError('Invalid actual value type: string or array')
+
+        if expected_value in actual_value:
+            return False
+        return True
+
+    def _assert_start_with(self, actual_value: Optional[str], expected_value: str) -> bool:
+        """
+        Assert start with
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if not actual_value:
+            return False
+
+        if not isinstance(actual_value, str):
+            raise ValueError('Invalid actual value type: string')
+
+        if not actual_value.startswith(expected_value):
+            return False
+        return True
+
+    def _assert_end_with(self, actual_value: Optional[str], expected_value: str) -> bool:
+        """
+        Assert end with
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if not actual_value:
+            return False
+
+        if not isinstance(actual_value, str):
+            raise ValueError('Invalid actual value type: string')
+
+        if not actual_value.endswith(expected_value):
+            return False
+        return True
+
+    def _assert_is(self, actual_value: Optional[str], expected_value: str) -> bool:
+        """
+        Assert is
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if actual_value is None:
+            return False
+
+        if not isinstance(actual_value, str):
+            raise ValueError('Invalid actual value type: string')
+
+        if actual_value != expected_value:
+            return False
+        return True
+
+    def _assert_is_not(self, actual_value: Optional[str], expected_value: str) -> bool:
+        """
+        Assert is not
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if actual_value is None:
+            return False
+
+        if not isinstance(actual_value, str):
+            raise ValueError('Invalid actual value type: string')
+
+        if actual_value == expected_value:
+            return False
+        return True
+
+    def _assert_empty(self, actual_value: Optional[str]) -> bool:
+        """
+        Assert empty
+        :param actual_value: actual value
+        :return:
+        """
+        if not actual_value:
+            return True
+        return False
+
+    def _assert_not_empty(self, actual_value: Optional[str]) -> bool:
+        """
+        Assert not empty
+        :param actual_value: actual value
+        :return:
+        """
+        if actual_value:
+            return True
+        return False
+
+    def _assert_equal(self, actual_value: Optional[int | float], expected_value: str) -> bool:
+        """
+        Assert equal
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if actual_value is None:
+            return False
+
+        if not isinstance(actual_value, int | float):
+            raise ValueError('Invalid actual value type: number')
+
+        if isinstance(actual_value, int):
+            expected_value = int(expected_value)
+        else:
+            expected_value = float(expected_value)
+
+        if actual_value != expected_value:
+            return False
+        return True
+
+    def _assert_not_equal(self, actual_value: Optional[int | float], expected_value: str) -> bool:
+        """
+        Assert not equal
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if actual_value is None:
+            return False
+
+        if not isinstance(actual_value, int | float):
+            raise ValueError('Invalid actual value type: number')
+
+        if isinstance(actual_value, int):
+            expected_value = int(expected_value)
+        else:
+            expected_value = float(expected_value)
+
+        if actual_value == expected_value:
+            return False
+        return True
+
+    def _assert_greater_than(self, actual_value: Optional[int | float], expected_value: str) -> bool:
+        """
+        Assert greater than
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if actual_value is None:
+            return False
+
+        if not isinstance(actual_value, int | float):
+            raise ValueError('Invalid actual value type: number')
+
+        if isinstance(actual_value, int):
+            expected_value = int(expected_value)
+        else:
+            expected_value = float(expected_value)
+
+        if actual_value <= expected_value:
+            return False
+        return True
+
+    def _assert_less_than(self, actual_value: Optional[int | float], expected_value: str) -> bool:
+        """
+        Assert less than
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if actual_value is None:
+            return False
+
+        if not isinstance(actual_value, int | float):
+            raise ValueError('Invalid actual value type: number')
+
+        if isinstance(actual_value, int):
+            expected_value = int(expected_value)
+        else:
+            expected_value = float(expected_value)
+
+        if actual_value >= expected_value:
+            return False
+        return True
+
+    def _assert_greater_than_or_equal(self, actual_value: Optional[int | float], expected_value: str) -> bool:
+        """
+        Assert greater than or equal
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if actual_value is None:
+            return False
+
+        if not isinstance(actual_value, int | float):
+            raise ValueError('Invalid actual value type: number')
+
+        if isinstance(actual_value, int):
+            expected_value = int(expected_value)
+        else:
+            expected_value = float(expected_value)
+
+        if actual_value < expected_value:
+            return False
+        return True
+
+    def _assert_less_than_or_equal(self, actual_value: Optional[int | float], expected_value: str) -> bool:
+        """
+        Assert less than or equal
+        :param actual_value: actual value
+        :param expected_value: expected value
+        :return:
+        """
+        if actual_value is None:
+            return False
+
+        if not isinstance(actual_value, int | float):
+            raise ValueError('Invalid actual value type: number')
+
+        if isinstance(actual_value, int):
+            expected_value = int(expected_value)
+        else:
+            expected_value = float(expected_value)
+
+        if actual_value > expected_value:
+            return False
+        return True
+
+    def _assert_null(self, actual_value: Optional[int | float]) -> bool:
+        """
+        Assert null
+        :param actual_value: actual value
+        :return:
+        """
+        if actual_value is None:
+            return True
+        return False
+
+    def _assert_not_null(self, actual_value: Optional[int | float]) -> bool:
+        """
+        Assert not null
+        :param actual_value: actual value
+        :return:
+        """
+        if actual_value is not None:
+            return True
+        return False
+
+    @classmethod
+    def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[str, list[str]]:
+        """
+        Extract variable selector to variable mapping
+        :param node_data: node data
+        :return:
+        """
+        return {}
diff --git a/api/tests/unit_tests/core/workflow/nodes/if_else_node.py b/api/tests/unit_tests/core/workflow/nodes/if_else_node.py
new file mode 100644
index 0000000000..7b402ad0a0
--- /dev/null
+++ b/api/tests/unit_tests/core/workflow/nodes/if_else_node.py
@@ -0,0 +1,193 @@
+from unittest.mock import MagicMock
+
+from core.workflow.entities.node_entities import SystemVariable
+from core.workflow.entities.variable_pool import VariablePool
+from core.workflow.nodes.base_node import UserFrom
+from core.workflow.nodes.if_else.if_else_node import IfElseNode
+from extensions.ext_database import db
+from models.workflow import WorkflowNodeExecutionStatus
+
+
+def test_execute_if_else_result_true():
+    node = IfElseNode(
+        tenant_id='1',
+        app_id='1',
+        workflow_id='1',
+        user_id='1',
+        user_from=UserFrom.ACCOUNT,
+        config={
+            'id': 'if-else',
+            'data': {
+                'title': '123',
+                'type': 'if-else',
+                'logical_operator': 'and',
+                'conditions': [
+                    {
+                        'comparison_operator': 'contains',
+                        'variable_selector': ['start', 'array_contains'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'not contains',
+                        'variable_selector': ['start', 'array_not_contains'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'contains',
+                        'variable_selector': ['start', 'contains'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'not contains',
+                        'variable_selector': ['start', 'not_contains'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'start with',
+                        'variable_selector': ['start', 'start_with'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'end with',
+                        'variable_selector': ['start', 'end_with'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'is',
+                        'variable_selector': ['start', 'is'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'is not',
+                        'variable_selector': ['start', 'is_not'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'empty',
+                        'variable_selector': ['start', 'empty'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'not empty',
+                        'variable_selector': ['start', 'not_empty'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': '=',
+                        'variable_selector': ['start', 'equals'],
+                        'value': '22'
+                    },
+                    {
+                        'comparison_operator': '≠',
+                        'variable_selector': ['start', 'not_equals'],
+                        'value': '22'
+                    },
+                    {
+                        'comparison_operator': '>',
+                        'variable_selector': ['start', 'greater_than'],
+                        'value': '22'
+                    },
+                    {
+                        'comparison_operator': '<',
+                        'variable_selector': ['start', 'less_than'],
+                        'value': '22'
+                    },
+                    {
+                        'comparison_operator': '≥',
+                        'variable_selector': ['start', 'greater_than_or_equal'],
+                        'value': '22'
+                    },
+                    {
+                        'comparison_operator': '≤',
+                        'variable_selector': ['start', 'less_than_or_equal'],
+                        'value': '22'
+                    },
+                    {
+                        'comparison_operator': 'null',
+                        'variable_selector': ['start', 'null']
+                    },
+                    {
+                        'comparison_operator': 'not null',
+                        'variable_selector': ['start', 'not_null']
+                    },
+                ]
+            }
+        }
+    )
+
+    # construct variable pool
+    pool = VariablePool(system_variables={
+        SystemVariable.FILES: [],
+    }, user_inputs={})
+    pool.append_variable(node_id='start', variable_key_list=['array_contains'], value=['ab', 'def'])
+    pool.append_variable(node_id='start', variable_key_list=['array_not_contains'], value=['ac', 'def'])
+    pool.append_variable(node_id='start', variable_key_list=['contains'], value='cabcde')
+    pool.append_variable(node_id='start', variable_key_list=['not_contains'], value='zacde')
+    pool.append_variable(node_id='start', variable_key_list=['start_with'], value='abc')
+    pool.append_variable(node_id='start', variable_key_list=['end_with'], value='zzab')
+    pool.append_variable(node_id='start', variable_key_list=['is'], value='ab')
+    pool.append_variable(node_id='start', variable_key_list=['is_not'], value='aab')
+    pool.append_variable(node_id='start', variable_key_list=['empty'], value='')
+    pool.append_variable(node_id='start', variable_key_list=['not_empty'], value='aaa')
+    pool.append_variable(node_id='start', variable_key_list=['equals'], value=22)
+    pool.append_variable(node_id='start', variable_key_list=['not_equals'], value=23)
+    pool.append_variable(node_id='start', variable_key_list=['greater_than'], value=23)
+    pool.append_variable(node_id='start', variable_key_list=['less_than'], value=21)
+    pool.append_variable(node_id='start', variable_key_list=['greater_than_or_equal'], value=22)
+    pool.append_variable(node_id='start', variable_key_list=['less_than_or_equal'], value=21)
+    pool.append_variable(node_id='start', variable_key_list=['not_null'], value='1212')
+
+    # Mock db.session.close()
+    db.session.close = MagicMock()
+
+    # execute node
+    result = node._run(pool)
+
+    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
+    assert result.outputs['result'] is True
+
+
+def test_execute_if_else_result_false():
+    node = IfElseNode(
+        tenant_id='1',
+        app_id='1',
+        workflow_id='1',
+        user_id='1',
+        user_from=UserFrom.ACCOUNT,
+        config={
+            'id': 'if-else',
+            'data': {
+                'title': '123',
+                'type': 'if-else',
+                'logical_operator': 'or',
+                'conditions': [
+                    {
+                        'comparison_operator': 'contains',
+                        'variable_selector': ['start', 'array_contains'],
+                        'value': 'ab'
+                    },
+                    {
+                        'comparison_operator': 'not contains',
+                        'variable_selector': ['start', 'array_not_contains'],
+                        'value': 'ab'
+                    }
+                ]
+            }
+        }
+    )
+
+    # construct variable pool
+    pool = VariablePool(system_variables={
+        SystemVariable.FILES: [],
+    }, user_inputs={})
+    pool.append_variable(node_id='start', variable_key_list=['array_contains'], value=['1ab', 'def'])
+    pool.append_variable(node_id='start', variable_key_list=['array_not_contains'], value=['ab', 'def'])
+
+    # Mock db.session.close()
+    db.session.close = MagicMock()
+
+    # execute node
+    result = node._run(pool)
+
+    assert result.status == WorkflowNodeExecutionStatus.SUCCEEDED
+    assert result.outputs['result'] is False