refactor(variables): clarify base vs union type naming (#30634)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

api/core/app/apps/advanced_chat/app_runner.py

@@ -24,7 +24,7 @@ from core.app.layers.conversation_variable_persist_layer import ConversationVari
 from core.db.session_factory import session_factory
 from core.moderation.base import ModerationError
 from core.moderation.input_moderation import InputModeration
-from core.variables.variables import VariableUnion
+from core.variables.variables import Variable
 from core.workflow.enums import WorkflowType
 from core.workflow.graph_engine.command_channels.redis_channel import RedisChannel
 from core.workflow.graph_engine.layers.base import GraphEngineLayer
@@ -149,8 +149,8 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
                 system_variables=system_inputs,
                 user_inputs=inputs,
                 environment_variables=self._workflow.environment_variables,
-                # Based on the definition of `VariableUnion`,
-                # `list[Variable]` can be safely used as `list[VariableUnion]` since they are compatible.
+                # Based on the definition of `Variable`,
+                # `VariableBase` instances can be safely used as `Variable` since they are compatible.
                 conversation_variables=conversation_variables,
             )
 
@@ -318,7 +318,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
             trace_manager=app_generate_entity.trace_manager,
         )
 
-    def _initialize_conversation_variables(self) -> list[VariableUnion]:
+    def _initialize_conversation_variables(self) -> list[Variable]:
         """
         Initialize conversation variables for the current conversation.
 
@@ -343,7 +343,7 @@ class AdvancedChatAppRunner(WorkflowBasedAppRunner):
             conversation_variables = [var.to_variable() for var in existing_variables]
 
             session.commit()
-            return cast(list[VariableUnion], conversation_variables)
+            return cast(list[Variable], conversation_variables)
 
     def _load_existing_conversation_variables(self, session: Session) -> list[ConversationVariable]:
         """

api/core/app/layers/conversation_variable_persist_layer.py

@@ -1,6 +1,6 @@
 import logging
 
-from core.variables import Variable
+from core.variables import VariableBase
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
 from core.workflow.conversation_variable_updater import ConversationVariableUpdater
 from core.workflow.enums import NodeType
@@ -44,7 +44,7 @@ class ConversationVariablePersistenceLayer(GraphEngineLayer):
             if selector[0] != CONVERSATION_VARIABLE_NODE_ID:
                 continue
             variable = self.graph_runtime_state.variable_pool.get(selector)
-            if not isinstance(variable, Variable):
+            if not isinstance(variable, VariableBase):
                 logger.warning(
                     "Conversation variable not found in variable pool. selector=%s",
                     selector,

api/core/variables/__init__.py

@@ -30,6 +30,7 @@ from .variables import (
     SecretVariable,
     StringVariable,
     Variable,
+    VariableBase,
 )
 
 __all__ = [
@@ -62,4 +63,5 @@ __all__ = [
     "StringSegment",
     "StringVariable",
     "Variable",
+    "VariableBase",
 ]

api/core/variables/segments.py

@@ -232,7 +232,7 @@ def get_segment_discriminator(v: Any) -> SegmentType | None:
 # - All variants in `SegmentUnion` must inherit from the `Segment` class.
 # - The union must include all non-abstract subclasses of `Segment`, except:
 #   - `SegmentGroup`, which is not added to the variable pool.
-#   - `Variable` and its subclasses, which are handled by `VariableUnion`.
+#   - `VariableBase` and its subclasses, which are handled by `Variable`.
 SegmentUnion: TypeAlias = Annotated[
     (
         Annotated[NoneSegment, Tag(SegmentType.NONE)]

api/core/variables/variables.py

@@ -27,7 +27,7 @@ from .segments import (
 from .types import SegmentType
 
 
-class Variable(Segment):
+class VariableBase(Segment):
     """
     A variable is a segment that has a name.
 
@@ -45,23 +45,23 @@ class Variable(Segment):
     selector: Sequence[str] = Field(default_factory=list)
 
 
-class StringVariable(StringSegment, Variable):
+class StringVariable(StringSegment, VariableBase):
     pass
 
 
-class FloatVariable(FloatSegment, Variable):
+class FloatVariable(FloatSegment, VariableBase):
     pass
 
 
-class IntegerVariable(IntegerSegment, Variable):
+class IntegerVariable(IntegerSegment, VariableBase):
     pass
 
 
-class ObjectVariable(ObjectSegment, Variable):
+class ObjectVariable(ObjectSegment, VariableBase):
     pass
 
 
-class ArrayVariable(ArraySegment, Variable):
+class ArrayVariable(ArraySegment, VariableBase):
     pass
 
 
@@ -89,16 +89,16 @@ class SecretVariable(StringVariable):
         return encrypter.obfuscated_token(self.value)
 
 
-class NoneVariable(NoneSegment, Variable):
+class NoneVariable(NoneSegment, VariableBase):
     value_type: SegmentType = SegmentType.NONE
     value: None = None
 
 
-class FileVariable(FileSegment, Variable):
+class FileVariable(FileSegment, VariableBase):
     pass
 
 
-class BooleanVariable(BooleanSegment, Variable):
+class BooleanVariable(BooleanSegment, VariableBase):
     pass
 
 
@@ -139,13 +139,13 @@ class RAGPipelineVariableInput(BaseModel):
     value: Any
 
 
-# The `VariableUnion`` type is used to enable serialization and deserialization with Pydantic.
-# Use `Variable` for type hinting when serialization is not required.
+# The `Variable` type is used to enable serialization and deserialization with Pydantic.
+# Use `VariableBase` for type hinting when serialization is not required.
 #
 # Note:
-# - All variants in `VariableUnion` must inherit from the `Variable` class.
-# - The union must include all non-abstract subclasses of `Segment`, except:
-VariableUnion: TypeAlias = Annotated[
+# - All variants in `Variable` must inherit from the `VariableBase` class.
+# - The union must include all non-abstract subclasses of `VariableBase`.
+Variable: TypeAlias = Annotated[
     (
         Annotated[NoneVariable, Tag(SegmentType.NONE)]
         | Annotated[StringVariable, Tag(SegmentType.STRING)]

api/core/workflow/conversation_variable_updater.py

@@ -1,7 +1,7 @@
 import abc
 from typing import Protocol
 
-from core.variables import Variable
+from core.variables import VariableBase
 
 
 class ConversationVariableUpdater(Protocol):
@@ -20,12 +20,12 @@ class ConversationVariableUpdater(Protocol):
     """
 
     @abc.abstractmethod
-    def update(self, conversation_id: str, variable: "Variable"):
+    def update(self, conversation_id: str, variable: "VariableBase"):
         """
         Updates the value of the specified conversation variable in the underlying storage.
 
         :param conversation_id: The ID of the conversation to update. Typically references `ConversationVariable.id`.
-        :param variable: The `Variable` instance containing the updated value.
+        :param variable: The `VariableBase` instance containing the updated value.
         """
         pass
 

api/core/workflow/graph_engine/entities/commands.py

@@ -11,7 +11,7 @@ from typing import Any
 
 from pydantic import BaseModel, Field
 
-from core.variables.variables import VariableUnion
+from core.variables.variables import Variable
 
 
 class CommandType(StrEnum):
@@ -46,7 +46,7 @@ class PauseCommand(GraphEngineCommand):
 class VariableUpdate(BaseModel):
     """Represents a single variable update instruction."""
 
-    value: VariableUnion = Field(description="New variable value")
+    value: Variable = Field(description="New variable value")
 
 
 class UpdateVariablesCommand(GraphEngineCommand):

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

@@ -11,7 +11,7 @@ from typing_extensions import TypeIs
 from core.model_runtime.entities.llm_entities import LLMUsage
 from core.variables import IntegerVariable, NoneSegment
 from core.variables.segments import ArrayAnySegment, ArraySegment
-from core.variables.variables import VariableUnion
+from core.variables.variables import Variable
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
 from core.workflow.enums import (
     NodeExecutionType,
@@ -240,7 +240,7 @@ class IterationNode(LLMUsageTrackingMixin, Node[IterationNodeData]):
                         datetime,
                         list[GraphNodeEventBase],
                         object | None,
-                        dict[str, VariableUnion],
+                        dict[str, Variable],
                         LLMUsage,
                     ]
                 ],
@@ -308,7 +308,7 @@ class IterationNode(LLMUsageTrackingMixin, Node[IterationNodeData]):
         item: object,
         flask_app: Flask,
         context_vars: contextvars.Context,
-    ) -> tuple[datetime, list[GraphNodeEventBase], object | None, dict[str, VariableUnion], LLMUsage]:
+    ) -> tuple[datetime, list[GraphNodeEventBase], object | None, dict[str, Variable], LLMUsage]:
         """Execute a single iteration in parallel mode and return results."""
         with preserve_flask_contexts(flask_app=flask_app, context_vars=context_vars):
             iter_start_at = datetime.now(UTC).replace(tzinfo=None)
@@ -515,11 +515,11 @@ class IterationNode(LLMUsageTrackingMixin, Node[IterationNodeData]):
 
         return variable_mapping
 
-    def _extract_conversation_variable_snapshot(self, *, variable_pool: VariablePool) -> dict[str, VariableUnion]:
+    def _extract_conversation_variable_snapshot(self, *, variable_pool: VariablePool) -> dict[str, Variable]:
         conversation_variables = variable_pool.variable_dictionary.get(CONVERSATION_VARIABLE_NODE_ID, {})
         return {name: variable.model_copy(deep=True) for name, variable in conversation_variables.items()}
 
-    def _sync_conversation_variables_from_snapshot(self, snapshot: dict[str, VariableUnion]) -> None:
+    def _sync_conversation_variables_from_snapshot(self, snapshot: dict[str, Variable]) -> None:
         parent_pool = self.graph_runtime_state.variable_pool
         parent_conversations = parent_pool.variable_dictionary.get(CONVERSATION_VARIABLE_NODE_ID, {})
 

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

@@ -1,7 +1,7 @@
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any
 
-from core.variables import SegmentType, Variable
+from core.variables import SegmentType, VariableBase
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
 from core.workflow.entities import GraphInitParams
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
@@ -73,7 +73,7 @@ class VariableAssignerNode(Node[VariableAssignerData]):
         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, Variable):
+        if not isinstance(original_variable, VariableBase):
             raise VariableOperatorNodeError("assigned variable not found")
 
         match self.node_data.write_mode:

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

@@ -2,7 +2,7 @@ import json
 from collections.abc import Mapping, MutableMapping, Sequence
 from typing import TYPE_CHECKING, Any
 
-from core.variables import SegmentType, Variable
+from core.variables import SegmentType, VariableBase
 from core.variables.consts import SELECTORS_LENGTH
 from core.workflow.constants import CONVERSATION_VARIABLE_NODE_ID
 from core.workflow.enums import NodeType, WorkflowNodeExecutionStatus
@@ -118,7 +118,7 @@ class VariableAssignerNode(Node[VariableAssignerNodeData]):
                 # ==================== Validation Part
 
                 # Check if variable exists
-                if not isinstance(variable, Variable):
+                if not isinstance(variable, VariableBase):
                     raise VariableNotFoundError(variable_selector=item.variable_selector)
 
                 # Check if operation is supported
@@ -192,7 +192,7 @@ class VariableAssignerNode(Node[VariableAssignerNodeData]):
 
         for selector in updated_variable_selectors:
             variable = self.graph_runtime_state.variable_pool.get(selector)
-            if not isinstance(variable, Variable):
+            if not isinstance(variable, VariableBase):
                 raise VariableNotFoundError(variable_selector=selector)
             process_data[variable.name] = variable.value
 
@@ -213,7 +213,7 @@ class VariableAssignerNode(Node[VariableAssignerNodeData]):
     def _handle_item(
         self,
         *,
-        variable: Variable,
+        variable: VariableBase,
         operation: Operation,
         value: Any,
     ):

api/core/workflow/runtime/variable_pool.py

@@ -9,10 +9,10 @@ from typing import Annotated, Any, Union, cast
 from pydantic import BaseModel, Field
 
 from core.file import File, FileAttribute, file_manager
-from core.variables import Segment, SegmentGroup, Variable
+from core.variables import Segment, SegmentGroup, VariableBase
 from core.variables.consts import SELECTORS_LENGTH
 from core.variables.segments import FileSegment, ObjectSegment
-from core.variables.variables import RAGPipelineVariableInput, VariableUnion
+from core.variables.variables import RAGPipelineVariableInput, Variable
 from core.workflow.constants import (
     CONVERSATION_VARIABLE_NODE_ID,
     ENVIRONMENT_VARIABLE_NODE_ID,
@@ -32,7 +32,7 @@ class VariablePool(BaseModel):
     # The first element of the selector is the node id, it's the first-level key in the dictionary.
     # Other elements of the selector are the keys in the second-level dictionary. To get the key, we hash the
     # elements of the selector except the first one.
-    variable_dictionary: defaultdict[str, Annotated[dict[str, VariableUnion], Field(default_factory=dict)]] = Field(
+    variable_dictionary: defaultdict[str, Annotated[dict[str, Variable], Field(default_factory=dict)]] = Field(
         description="Variables mapping",
         default=defaultdict(dict),
     )
@@ -46,13 +46,13 @@ class VariablePool(BaseModel):
         description="System variables",
         default_factory=SystemVariable.empty,
     )
-    environment_variables: Sequence[VariableUnion] = Field(
+    environment_variables: Sequence[Variable] = Field(
         description="Environment variables.",
-        default_factory=list[VariableUnion],
+        default_factory=list[Variable],
     )
-    conversation_variables: Sequence[VariableUnion] = Field(
+    conversation_variables: Sequence[Variable] = Field(
         description="Conversation variables.",
-        default_factory=list[VariableUnion],
+        default_factory=list[Variable],
     )
     rag_pipeline_variables: list[RAGPipelineVariableInput] = Field(
         description="RAG pipeline variables.",
@@ -105,7 +105,7 @@ class VariablePool(BaseModel):
                 f"got {len(selector)} elements"
             )
 
-        if isinstance(value, Variable):
+        if isinstance(value, VariableBase):
             variable = value
         elif isinstance(value, Segment):
             variable = variable_factory.segment_to_variable(segment=value, selector=selector)
@@ -114,9 +114,9 @@ class VariablePool(BaseModel):
             variable = variable_factory.segment_to_variable(segment=segment, selector=selector)
 
         node_id, name = self._selector_to_keys(selector)
-        # Based on the definition of `VariableUnion`,
-        # `list[Variable]` can be safely used as `list[VariableUnion]` since they are compatible.
-        self.variable_dictionary[node_id][name] = cast(VariableUnion, variable)
+        # Based on the definition of `Variable`,
+        # `VariableBase` instances can be safely used as `Variable` since they are compatible.
+        self.variable_dictionary[node_id][name] = cast(Variable, variable)
 
     @classmethod
     def _selector_to_keys(cls, selector: Sequence[str]) -> tuple[str, str]:

api/core/workflow/variable_loader.py

@@ -2,7 +2,7 @@ import abc
 from collections.abc import Mapping, Sequence
 from typing import Any, Protocol
 
-from core.variables import Variable
+from core.variables import VariableBase
 from core.variables.consts import SELECTORS_LENGTH
 from core.workflow.runtime import VariablePool
 
@@ -26,7 +26,7 @@ class VariableLoader(Protocol):
     """
 
     @abc.abstractmethod
-    def load_variables(self, selectors: list[list[str]]) -> list[Variable]:
+    def load_variables(self, selectors: list[list[str]]) -> list[VariableBase]:
         """Load variables based on the provided selectors. If the selectors are empty,
         this method should return an empty list.
 
@@ -36,7 +36,7 @@ class VariableLoader(Protocol):
         :param: selectors: a list of string list, each inner list should have at least two elements:
             - the first element is the node ID,
             - the second element is the variable name.
-        :return: a list of Variable objects that match the provided selectors.
+        :return: a list of VariableBase objects that match the provided selectors.
         """
         pass
 
@@ -46,7 +46,7 @@ class _DummyVariableLoader(VariableLoader):
     Serves as a placeholder when no variable loading is needed.
     """
 
-    def load_variables(self, selectors: list[list[str]]) -> list[Variable]:
+    def load_variables(self, selectors: list[list[str]]) -> list[VariableBase]:
         return []