refactor(api): Separate SegmentType for Integer/Float to Enable Pydantic Serialization (#22025)

refactor(api): Separate SegmentType for Integer/Float to Enable Pydantic Serialization (#22025)

This PR addresses serialization issues in the VariablePool model by separating the `value_type` tags for `IntegerSegment`/`FloatSegment` and `IntegerVariable`/`FloatVariable`. Previously, both Integer and Float types shared the same `SegmentType.NUMBER` tag, causing conflicts during serialization.

Key changes:
- Introduce distinct `value_type` tags for Integer and Float segments/variables
- Add `VariableUnion` and `SegmentUnion` types for proper type discrimination
- Leverage Pydantic's discriminated union feature for seamless serialization/deserialization
- Enable accurate serialization of data structures containing these types

Closes #22024.

api/core/variables/variables.py

@@ -1,8 +1,8 @@
 from collections.abc import Sequence
-from typing import cast
+from typing import Annotated, TypeAlias, cast
 from uuid import uuid4
 
-from pydantic import Field
+from pydantic import Discriminator, Field, Tag
 
 from core.helper import encrypter
 
@@ -20,6 +20,7 @@ from .segments import (
     ObjectSegment,
     Segment,
     StringSegment,
+    get_segment_discriminator,
 )
 from .types import SegmentType
 
@@ -27,6 +28,10 @@ from .types import SegmentType
 class Variable(Segment):
     """
     A variable is a segment that has a name.
+
+    It is mainly used to store segments and their selector in VariablePool.
+
+    Note: this class is abstract, you should use subclasses of this class instead.
     """
 
     id: str = Field(
@@ -93,3 +98,28 @@ class FileVariable(FileSegment, Variable):
 
 class ArrayFileVariable(ArrayFileSegment, ArrayVariable):
     pass
+
+
+# The `VariableUnion`` type is used to enable serialization and deserialization with Pydantic.
+# Use `Variable` 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[
+    (
+        Annotated[NoneVariable, Tag(SegmentType.NONE)]
+        | Annotated[StringVariable, Tag(SegmentType.STRING)]
+        | Annotated[FloatVariable, Tag(SegmentType.FLOAT)]
+        | Annotated[IntegerVariable, Tag(SegmentType.INTEGER)]
+        | Annotated[ObjectVariable, Tag(SegmentType.OBJECT)]
+        | Annotated[FileVariable, Tag(SegmentType.FILE)]
+        | Annotated[ArrayAnyVariable, Tag(SegmentType.ARRAY_ANY)]
+        | Annotated[ArrayStringVariable, Tag(SegmentType.ARRAY_STRING)]
+        | Annotated[ArrayNumberVariable, Tag(SegmentType.ARRAY_NUMBER)]
+        | Annotated[ArrayObjectVariable, Tag(SegmentType.ARRAY_OBJECT)]
+        | Annotated[ArrayFileVariable, Tag(SegmentType.ARRAY_FILE)]
+        | Annotated[SecretVariable, Tag(SegmentType.SECRET)]
+    ),
+    Discriminator(get_segment_discriminator),
+]