Feat: upgrade variable assigner (#11285)

Signed-off-by: -LAN- <laipz8200@outlook.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>

api/controllers/console/app/workflow.py

@@ -100,11 +100,11 @@ class DraftWorkflowApi(Resource):
         try:
             environment_variables_list = args.get("environment_variables") or []
             environment_variables = [
-                variable_factory.build_variable_from_mapping(obj) for obj in environment_variables_list
+                variable_factory.build_environment_variable_from_mapping(obj) for obj in environment_variables_list
             ]
             conversation_variables_list = args.get("conversation_variables") or []
             conversation_variables = [
-                variable_factory.build_variable_from_mapping(obj) for obj in conversation_variables_list
+                variable_factory.build_conversation_variable_from_mapping(obj) for obj in conversation_variables_list
             ]
             workflow = workflow_service.sync_draft_workflow(
                 app_model=app_model,
@@ -382,7 +382,7 @@ class DefaultBlockConfigApi(Resource):
         filters = None
         if args.get("q"):
             try:
-                filters = json.loads(args.get("q"))
+                filters = json.loads(args.get("q", ""))
             except json.JSONDecodeError:
                 raise ValueError("Invalid filters")
 

api/core/app/apps/workflow_app_runner.py

@@ -43,7 +43,7 @@ from core.workflow.graph_engine.entities.event import (
 )
 from core.workflow.graph_engine.entities.graph import Graph
 from core.workflow.nodes import NodeType
-from core.workflow.nodes.node_mapping import node_type_classes_mapping
+from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
 from core.workflow.workflow_entry import WorkflowEntry
 from extensions.ext_database import db
 from models.model import App
@@ -138,7 +138,8 @@ class WorkflowBasedAppRunner(AppRunner):
 
         # Get node class
         node_type = NodeType(iteration_node_config.get("data", {}).get("type"))
-        node_cls = node_type_classes_mapping[node_type]
+        node_version = iteration_node_config.get("data", {}).get("version", "1")
+        node_cls = NODE_TYPE_CLASSES_MAPPING[node_type][node_version]
 
         # init variable pool
         variable_pool = VariablePool(

api/core/variables/types.py

@@ -2,16 +2,19 @@ from enum import StrEnum
 
 
 class SegmentType(StrEnum):
-    NONE = "none"
     NUMBER = "number"
     STRING = "string"
+    OBJECT = "object"
     SECRET = "secret"
+
+    FILE = "file"
+
     ARRAY_ANY = "array[any]"
     ARRAY_STRING = "array[string]"
     ARRAY_NUMBER = "array[number]"
     ARRAY_OBJECT = "array[object]"
-    OBJECT = "object"
-    FILE = "file"
     ARRAY_FILE = "array[file]"
 
+    NONE = "none"
+
     GROUP = "group"

api/core/workflow/graph_engine/graph_engine.py

@@ -38,7 +38,7 @@ from core.workflow.nodes.answer.answer_stream_processor import AnswerStreamProce
 from core.workflow.nodes.base import BaseNode
 from core.workflow.nodes.end.end_stream_processor import EndStreamProcessor
 from core.workflow.nodes.event import RunCompletedEvent, RunRetrieverResourceEvent, RunStreamChunkEvent
-from core.workflow.nodes.node_mapping import node_type_classes_mapping
+from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
 from extensions.ext_database import db
 from models.enums import UserFrom
 from models.workflow import WorkflowNodeExecutionStatus, WorkflowType
@@ -227,7 +227,8 @@ class GraphEngine:
 
             # convert to specific node
             node_type = NodeType(node_config.get("data", {}).get("type"))
-            node_cls = node_type_classes_mapping[node_type]
+            node_version = node_config.get("data", {}).get("version", "1")
+            node_cls = NODE_TYPE_CLASSES_MAPPING[node_type][node_version]
 
             previous_node_id = previous_route_node_state.node_id if previous_route_node_state else None
 

api/core/workflow/nodes/answer/answer_stream_generate_router.py

@@ -153,7 +153,7 @@ class AnswerStreamGeneratorRouter:
                 NodeType.IF_ELSE,
                 NodeType.QUESTION_CLASSIFIER,
                 NodeType.ITERATION,
-                NodeType.CONVERSATION_VARIABLE_ASSIGNER,
+                NodeType.VARIABLE_ASSIGNER,
             }:
                 answer_dependencies[answer_node_id].append(source_node_id)
             else:

api/core/workflow/nodes/base/entities.py

@@ -7,6 +7,7 @@ from pydantic import BaseModel
 class BaseNodeData(ABC, BaseModel):
     title: str
     desc: Optional[str] = None
+    version: str = "1"
 
 
 class BaseIterationNodeData(BaseNodeData):

api/core/workflow/nodes/base/node.py

@@ -55,7 +55,9 @@ class BaseNode(Generic[GenericNodeData]):
             raise ValueError("Node ID is required.")
 
         self.node_id = node_id
-        self.node_data: GenericNodeData = cast(GenericNodeData, self._node_data_cls(**config.get("data", {})))
+
+        node_data = self._node_data_cls.model_validate(config.get("data", {}))
+        self.node_data = cast(GenericNodeData, node_data)
 
     @abstractmethod
     def _run(self) -> NodeRunResult | Generator[Union[NodeEvent, "InNodeEvent"], None, None]:

api/core/workflow/nodes/enums.py

@@ -14,11 +14,11 @@ class NodeType(StrEnum):
     HTTP_REQUEST = "http-request"
     TOOL = "tool"
     VARIABLE_AGGREGATOR = "variable-aggregator"
-    VARIABLE_ASSIGNER = "variable-assigner"  # TODO: Merge this into VARIABLE_AGGREGATOR in the database.
+    LEGACY_VARIABLE_AGGREGATOR = "variable-assigner"  # TODO: Merge this into VARIABLE_AGGREGATOR in the database.
     LOOP = "loop"
     ITERATION = "iteration"
     ITERATION_START = "iteration-start"  # Fake start node for iteration.
     PARAMETER_EXTRACTOR = "parameter-extractor"
-    CONVERSATION_VARIABLE_ASSIGNER = "assigner"
+    VARIABLE_ASSIGNER = "assigner"
     DOCUMENT_EXTRACTOR = "document-extractor"
     LIST_OPERATOR = "list-operator"

api/core/workflow/nodes/iteration/iteration_node.py

@@ -298,12 +298,13 @@ class IterationNode(BaseNode[IterationNodeData]):
             # variable selector to variable mapping
             try:
                 # Get node class
-                from core.workflow.nodes.node_mapping import node_type_classes_mapping
+                from core.workflow.nodes.node_mapping import NODE_TYPE_CLASSES_MAPPING
 
                 node_type = NodeType(sub_node_config.get("data", {}).get("type"))
-                node_cls = node_type_classes_mapping.get(node_type)
-                if not node_cls:
+                if node_type not in NODE_TYPE_CLASSES_MAPPING:
                     continue
+                node_version = sub_node_config.get("data", {}).get("version", "1")
+                node_cls = NODE_TYPE_CLASSES_MAPPING[node_type][node_version]
 
                 sub_node_variable_mapping = node_cls.extract_variable_selector_to_variable_mapping(
                     graph_config=graph_config, config=sub_node_config

api/core/workflow/nodes/node_mapping.py

@@ -1,3 +1,5 @@
+from collections.abc import Mapping
+
 from core.workflow.nodes.answer import AnswerNode
 from core.workflow.nodes.base import BaseNode
 from core.workflow.nodes.code import CodeNode
@@ -16,26 +18,87 @@ from core.workflow.nodes.start import StartNode
 from core.workflow.nodes.template_transform import TemplateTransformNode
 from core.workflow.nodes.tool import ToolNode
 from core.workflow.nodes.variable_aggregator import VariableAggregatorNode
-from core.workflow.nodes.variable_assigner import VariableAssignerNode
+from core.workflow.nodes.variable_assigner.v1 import VariableAssignerNode as VariableAssignerNodeV1
+from core.workflow.nodes.variable_assigner.v2 import VariableAssignerNode as VariableAssignerNodeV2
 
-node_type_classes_mapping: dict[NodeType, type[BaseNode]] = {
-    NodeType.START: StartNode,
-    NodeType.END: EndNode,
-    NodeType.ANSWER: AnswerNode,
-    NodeType.LLM: LLMNode,
-    NodeType.KNOWLEDGE_RETRIEVAL: KnowledgeRetrievalNode,
-    NodeType.IF_ELSE: IfElseNode,
-    NodeType.CODE: CodeNode,
-    NodeType.TEMPLATE_TRANSFORM: TemplateTransformNode,
-    NodeType.QUESTION_CLASSIFIER: QuestionClassifierNode,
-    NodeType.HTTP_REQUEST: HttpRequestNode,
-    NodeType.TOOL: ToolNode,
-    NodeType.VARIABLE_AGGREGATOR: VariableAggregatorNode,
-    NodeType.VARIABLE_ASSIGNER: VariableAggregatorNode,  # original name of VARIABLE_AGGREGATOR
-    NodeType.ITERATION: IterationNode,
-    NodeType.ITERATION_START: IterationStartNode,
-    NodeType.PARAMETER_EXTRACTOR: ParameterExtractorNode,
-    NodeType.CONVERSATION_VARIABLE_ASSIGNER: VariableAssignerNode,
-    NodeType.DOCUMENT_EXTRACTOR: DocumentExtractorNode,
-    NodeType.LIST_OPERATOR: ListOperatorNode,
+LATEST_VERSION = "latest"
+
+NODE_TYPE_CLASSES_MAPPING: Mapping[NodeType, Mapping[str, type[BaseNode]]] = {
+    NodeType.START: {
+        LATEST_VERSION: StartNode,
+        "1": StartNode,
+    },
+    NodeType.END: {
+        LATEST_VERSION: EndNode,
+        "1": EndNode,
+    },
+    NodeType.ANSWER: {
+        LATEST_VERSION: AnswerNode,
+        "1": AnswerNode,
+    },
+    NodeType.LLM: {
+        LATEST_VERSION: LLMNode,
+        "1": LLMNode,
+    },
+    NodeType.KNOWLEDGE_RETRIEVAL: {
+        LATEST_VERSION: KnowledgeRetrievalNode,
+        "1": KnowledgeRetrievalNode,
+    },
+    NodeType.IF_ELSE: {
+        LATEST_VERSION: IfElseNode,
+        "1": IfElseNode,
+    },
+    NodeType.CODE: {
+        LATEST_VERSION: CodeNode,
+        "1": CodeNode,
+    },
+    NodeType.TEMPLATE_TRANSFORM: {
+        LATEST_VERSION: TemplateTransformNode,
+        "1": TemplateTransformNode,
+    },
+    NodeType.QUESTION_CLASSIFIER: {
+        LATEST_VERSION: QuestionClassifierNode,
+        "1": QuestionClassifierNode,
+    },
+    NodeType.HTTP_REQUEST: {
+        LATEST_VERSION: HttpRequestNode,
+        "1": HttpRequestNode,
+    },
+    NodeType.TOOL: {
+        LATEST_VERSION: ToolNode,
+        "1": ToolNode,
+    },
+    NodeType.VARIABLE_AGGREGATOR: {
+        LATEST_VERSION: VariableAggregatorNode,
+        "1": VariableAggregatorNode,
+    },
+    NodeType.LEGACY_VARIABLE_AGGREGATOR: {
+        LATEST_VERSION: VariableAggregatorNode,
+        "1": VariableAggregatorNode,
+    },  # original name of VARIABLE_AGGREGATOR
+    NodeType.ITERATION: {
+        LATEST_VERSION: IterationNode,
+        "1": IterationNode,
+    },
+    NodeType.ITERATION_START: {
+        LATEST_VERSION: IterationStartNode,
+        "1": IterationStartNode,
+    },
+    NodeType.PARAMETER_EXTRACTOR: {
+        LATEST_VERSION: ParameterExtractorNode,
+        "1": ParameterExtractorNode,
+    },
+    NodeType.VARIABLE_ASSIGNER: {
+        LATEST_VERSION: VariableAssignerNodeV2,
+        "1": VariableAssignerNodeV1,
+        "2": VariableAssignerNodeV2,
+    },
+    NodeType.DOCUMENT_EXTRACTOR: {
+        LATEST_VERSION: DocumentExtractorNode,
+        "1": DocumentExtractorNode,
+    },
+    NodeType.LIST_OPERATOR: {
+        LATEST_VERSION: ListOperatorNode,
+        "1": ListOperatorNode,
+    },
 }

api/core/workflow/nodes/variable_assigner/__init__.py

@@ -1,8 +0,0 @@
-from .node import VariableAssignerNode
-from .node_data import VariableAssignerData, WriteMode
-
-__all__ = [
-    "VariableAssignerData",
-    "VariableAssignerNode",
-    "WriteMode",
-]

api/core/workflow/nodes/variable_assigner/common/exc.py

@@ -0,0 +1,4 @@
+class VariableOperatorNodeError(Exception):
+    """Base error type, don't use directly."""
+
+    pass

api/core/workflow/nodes/variable_assigner/common/helpers.py

@@ -0,0 +1,19 @@
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from core.variables import Variable
+from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError
+from extensions.ext_database import db
+from models import ConversationVariable
+
+
+def update_conversation_variable(conversation_id: str, variable: Variable):
+    stmt = select(ConversationVariable).where(
+        ConversationVariable.id == variable.id, ConversationVariable.conversation_id == conversation_id
+    )
+    with Session(db.engine) as session:
+        row = session.scalar(stmt)
+        if not row:
+            raise VariableOperatorNodeError("conversation variable not found in the database")
+        row.data = variable.model_dump_json()
+        session.commit()

api/core/workflow/nodes/variable_assigner/exc.py

@@ -1,2 +0,0 @@
-class VariableAssignerNodeError(Exception):
-    pass

api/core/workflow/nodes/variable_assigner/v1/__init__.py

@@ -0,0 +1,3 @@
+from .node import VariableAssignerNode
+
+__all__ = ["VariableAssignerNode"]

api/core/workflow/nodes/variable_assigner/v1/node.py

@@ -0,0 +1,73 @@
+from core.variables import SegmentType, Variable
+from core.workflow.entities.node_entities import NodeRunResult
+from core.workflow.nodes.base import BaseNode, BaseNodeData
+from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.variable_assigner.common import helpers as common_helpers
+from core.workflow.nodes.variable_assigner.common.exc import VariableOperatorNodeError
+from factories import variable_factory
+from models.workflow import WorkflowNodeExecutionStatus
+
+from .node_data import VariableAssignerData, WriteMode
+
+
+class VariableAssignerNode(BaseNode[VariableAssignerData]):
+    _node_data_cls: type[BaseNodeData] = VariableAssignerData
+    _node_type = NodeType.VARIABLE_ASSIGNER
+
+    def _run(self) -> NodeRunResult:
+        # Should be String, Number, Object, ArrayString, ArrayNumber, ArrayObject
+        original_variable = self.graph_runtime_state.variable_pool.get(self.node_data.assigned_variable_selector)
+        if not isinstance(original_variable, Variable):
+            raise VariableOperatorNodeError("assigned variable not found")
+
+        match self.node_data.write_mode:
+            case WriteMode.OVER_WRITE:
+                income_value = self.graph_runtime_state.variable_pool.get(self.node_data.input_variable_selector)
+                if not income_value:
+                    raise VariableOperatorNodeError("input value not found")
+                updated_variable = original_variable.model_copy(update={"value": income_value.value})
+
+            case WriteMode.APPEND:
+                income_value = self.graph_runtime_state.variable_pool.get(self.node_data.input_variable_selector)
+                if not income_value:
+                    raise VariableOperatorNodeError("input value not found")
+                updated_value = original_variable.value + [income_value.value]
+                updated_variable = original_variable.model_copy(update={"value": updated_value})
+
+            case WriteMode.CLEAR:
+                income_value = get_zero_value(original_variable.value_type)
+                updated_variable = original_variable.model_copy(update={"value": income_value.to_object()})
+
+            case _:
+                raise VariableOperatorNodeError(f"unsupported write mode: {self.node_data.write_mode}")
+
+        # Over write the variable.
+        self.graph_runtime_state.variable_pool.add(self.node_data.assigned_variable_selector, updated_variable)
+
+        # TODO: Move database operation to the pipeline.
+        # Update conversation variable.
+        conversation_id = self.graph_runtime_state.variable_pool.get(["sys", "conversation_id"])
+        if not conversation_id:
+            raise VariableOperatorNodeError("conversation_id not found")
+        common_helpers.update_conversation_variable(conversation_id=conversation_id.text, variable=updated_variable)
+
+        return NodeRunResult(
+            status=WorkflowNodeExecutionStatus.SUCCEEDED,
+            inputs={
+                "value": income_value.to_object(),
+            },
+        )
+
+
+def get_zero_value(t: SegmentType):
+    match t:
+        case SegmentType.ARRAY_OBJECT | SegmentType.ARRAY_STRING | SegmentType.ARRAY_NUMBER:
+            return variable_factory.build_segment([])
+        case SegmentType.OBJECT:
+            return variable_factory.build_segment({})
+        case SegmentType.STRING:
+            return variable_factory.build_segment("")
+        case SegmentType.NUMBER:
+            return variable_factory.build_segment(0)
+        case _:
+            raise VariableOperatorNodeError(f"unsupported variable type: {t}")

api/core/workflow/nodes/variable_assigner/v1/node_data.py

@@ -0,0 +1,16 @@
+from collections.abc import Sequence
+from enum import StrEnum
+
+from core.workflow.nodes.base import BaseNodeData
+
+
+class WriteMode(StrEnum):
+    OVER_WRITE = "over-write"
+    APPEND = "append"
+    CLEAR = "clear"
+
+
+class VariableAssignerData(BaseNodeData):
+    assigned_variable_selector: Sequence[str]
+    write_mode: WriteMode
+    input_variable_selector: Sequence[str]

api/core/workflow/nodes/variable_assigner/v2/__init__.py

@@ -0,0 +1,3 @@
+from .node import VariableAssignerNode
+
+__all__ = ["VariableAssignerNode"]

api/core/workflow/nodes/variable_assigner/v2/constants.py

@@ -0,0 +1,11 @@
+from core.variables import SegmentType
+
+EMPTY_VALUE_MAPPING = {
+    SegmentType.STRING: "",
+    SegmentType.NUMBER: 0,
+    SegmentType.OBJECT: {},
+    SegmentType.ARRAY_ANY: [],
+    SegmentType.ARRAY_STRING: [],
+    SegmentType.ARRAY_NUMBER: [],
+    SegmentType.ARRAY_OBJECT: [],
+}