refactor(api): rename dify_graph to graphon (#34095)

api/graphon/nodes/variable_assigner/v1/node.py

@@ -0,0 +1,106 @@
+from collections.abc import Generator, Mapping, Sequence
+from typing import TYPE_CHECKING, Any, cast
+
+from graphon.entities import GraphInitParams
+from graphon.entities.graph_config import NodeConfigDict
+from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionStatus
+from graphon.node_events import NodeEventBase, NodeRunResult, StreamCompletedEvent, VariableUpdatedEvent
+from graphon.nodes.base.node import Node
+from graphon.nodes.variable_assigner.common import helpers as common_helpers
+from graphon.nodes.variable_assigner.common.exc import VariableOperatorNodeError
+from graphon.variables import SegmentType, Variable, VariableBase
+
+from .node_data import VariableAssignerData, WriteMode
+
+if TYPE_CHECKING:
+    from graphon.runtime import GraphRuntimeState
+
+
+class VariableAssignerNode(Node[VariableAssignerData]):
+    node_type = BuiltinNodeTypes.VARIABLE_ASSIGNER
+
+    def __init__(
+        self,
+        id: str,
+        config: NodeConfigDict,
+        graph_init_params: "GraphInitParams",
+        graph_runtime_state: "GraphRuntimeState",
+    ):
+        super().__init__(
+            id=id,
+            config=config,
+            graph_init_params=graph_init_params,
+            graph_runtime_state=graph_runtime_state,
+        )
+
+    def blocks_variable_output(self, variable_selectors: set[tuple[str, ...]]) -> bool:
+        """
+        Check if this Variable Assigner node blocks the output of specific variables.
+
+        Returns True if this node updates any of the requested conversation variables.
+        """
+        assigned_selector = tuple(self.node_data.assigned_variable_selector)
+        return assigned_selector in variable_selectors
+
+    @classmethod
+    def version(cls) -> str:
+        return "1"
+
+    @classmethod
+    def _extract_variable_selector_to_variable_mapping(
+        cls,
+        *,
+        graph_config: Mapping[str, Any],
+        node_id: str,
+        node_data: VariableAssignerData,
+    ) -> Mapping[str, Sequence[str]]:
+        mapping = {}
+        selector_key = ".".join(node_data.assigned_variable_selector)
+        key = f"{node_id}.#{selector_key}#"
+        mapping[key] = node_data.assigned_variable_selector
+
+        selector_key = ".".join(node_data.input_variable_selector)
+        key = f"{node_id}.#{selector_key}#"
+        mapping[key] = node_data.input_variable_selector
+        return mapping
+
+    def _run(self) -> Generator[NodeEventBase, None, None]:
+        assigned_variable_selector = self.node_data.assigned_variable_selector
+        # Should be String, Number, Object, ArrayString, ArrayNumber, ArrayObject
+        original_variable = self.graph_runtime_state.variable_pool.get(assigned_variable_selector)
+        if not isinstance(original_variable, VariableBase):
+            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 = SegmentType.get_zero_value(original_variable.value_type)
+                updated_variable = original_variable.model_copy(update={"value": income_value.to_object()})
+
+        updated_variables = [common_helpers.variable_to_processed_data(assigned_variable_selector, updated_variable)]
+        yield VariableUpdatedEvent(variable=cast(Variable, updated_variable))
+        yield StreamCompletedEvent(
+            node_run_result=NodeRunResult(
+                status=WorkflowNodeExecutionStatus.SUCCEEDED,
+                inputs={
+                    "value": income_value.to_object(),
+                },
+                # NOTE(QuantumGhost): although only one variable is updated in `v1.VariableAssignerNode`,
+                # we still set `output_variables` as a list to ensure the schema of output is
+                # compatible with `v2.VariableAssignerNode`.
+                process_data=common_helpers.set_updated_variables({}, updated_variables),
+                outputs={},
+            )
+        )