feat(api/core/app/segments): Update segment types and variables (#6734)

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

api/core/app/segments/__init__.py

@@ -1,6 +1,6 @@
 from .segment_group import SegmentGroup
 from .segments import (
-    ArraySegment,
+    ArrayAnySegment,
     FileSegment,
     FloatSegment,
     IntegerSegment,
@@ -11,7 +11,11 @@ from .segments import (
 )
 from .types import SegmentType
 from .variables import (
-    ArrayVariable,
+    ArrayAnyVariable,
+    ArrayFileVariable,
+    ArrayNumberVariable,
+    ArrayObjectVariable,
+    ArrayStringVariable,
     FileVariable,
     FloatVariable,
     IntegerVariable,
@@ -29,7 +33,7 @@ __all__ = [
     'SecretVariable',
     'FileVariable',
     'StringVariable',
-    'ArrayVariable',
+    'ArrayAnyVariable',
     'Variable',
     'SegmentType',
     'SegmentGroup',
@@ -39,7 +43,11 @@ __all__ = [
     'IntegerSegment',
     'FloatSegment',
     'ObjectSegment',
-    'ArraySegment',
+    'ArrayAnySegment',
     'FileSegment',
     'StringSegment',
+    'ArrayStringVariable',
+    'ArrayNumberVariable',
+    'ArrayObjectVariable',
+    'ArrayFileVariable',
 ]

api/core/app/segments/factory.py

@@ -4,7 +4,7 @@ from typing import Any
 from core.file.file_obj import FileVar
 
 from .segments import (
-    ArraySegment,
+    ArrayAnySegment,
     FileSegment,
     FloatSegment,
     IntegerSegment,
@@ -15,8 +15,14 @@ from .segments import (
 )
 from .types import SegmentType
 from .variables import (
+    ArrayFileVariable,
+    ArrayNumberVariable,
+    ArrayObjectVariable,
+    ArrayStringVariable,
+    FileVariable,
     FloatVariable,
     IntegerVariable,
+    ObjectVariable,
     SecretVariable,
     StringVariable,
     Variable,
@@ -33,14 +39,28 @@ def build_variable_from_mapping(m: Mapping[str, Any], /) -> Variable:
     match value_type:
         case SegmentType.STRING:
             return StringVariable.model_validate(m)
+        case SegmentType.SECRET:
+            return SecretVariable.model_validate(m)
         case SegmentType.NUMBER if isinstance(value, int):
             return IntegerVariable.model_validate(m)
         case SegmentType.NUMBER if isinstance(value, float):
             return FloatVariable.model_validate(m)
-        case SegmentType.SECRET:
-            return SecretVariable.model_validate(m)
         case SegmentType.NUMBER if not isinstance(value, float | int):
             raise ValueError(f'invalid number value {value}')
+        case SegmentType.FILE:
+            return FileVariable.model_validate(m)
+        case SegmentType.OBJECT if isinstance(value, dict):
+            return ObjectVariable.model_validate(
+                {**m, 'value': {k: build_variable_from_mapping(v) for k, v in value.items()}}
+            )
+        case SegmentType.ARRAY_STRING if isinstance(value, list):
+            return ArrayStringVariable.model_validate({**m, 'value': [build_variable_from_mapping(v) for v in value]})
+        case SegmentType.ARRAY_NUMBER if isinstance(value, list):
+            return ArrayNumberVariable.model_validate({**m, 'value': [build_variable_from_mapping(v) for v in value]})
+        case SegmentType.ARRAY_OBJECT if isinstance(value, list):
+            return ArrayObjectVariable.model_validate({**m, 'value': [build_variable_from_mapping(v) for v in value]})
+        case SegmentType.ARRAY_FILE if isinstance(value, list):
+            return ArrayFileVariable.model_validate({**m, 'value': [build_variable_from_mapping(v) for v in value]})
     raise ValueError(f'not supported value type {value_type}')
 
 
@@ -60,7 +80,7 @@ def build_segment(value: Any, /) -> Segment:
     if isinstance(value, list):
         # TODO: Limit the depth of the array
         elements = [build_segment(v) for v in value]
-        return ArraySegment(value=elements)
+        return ArrayAnySegment(value=elements)
     if isinstance(value, FileVar):
         return FileSegment(value=value)
     raise ValueError(f'not supported value {value}')

api/core/app/segments/segments.py

@@ -62,6 +62,7 @@ class StringSegment(Segment):
     value_type: SegmentType = SegmentType.STRING
     value: str
 
+
 class FloatSegment(Segment):
     value_type: SegmentType = SegmentType.NUMBER
     value: float
@@ -72,6 +73,16 @@ class IntegerSegment(Segment):
     value: int
 
 
+class FileSegment(Segment):
+    value_type: SegmentType = SegmentType.FILE
+    # TODO: embed FileVar in this model.
+    value: FileVar
+
+    @property
+    def markdown(self) -> str:
+        return self.value.to_markdown()
+
+
 class ObjectSegment(Segment):
     value_type: SegmentType = SegmentType.OBJECT
     value: Mapping[str, Segment]
@@ -96,9 +107,6 @@ class ObjectSegment(Segment):
 
 
 class ArraySegment(Segment):
-    value_type: SegmentType = SegmentType.ARRAY
-    value: Sequence[Segment]
-
     @property
     def markdown(self) -> str:
         return '\n'.join(['- ' + item.markdown for item in self.value])
@@ -107,11 +115,26 @@ class ArraySegment(Segment):
         return [v.to_object() for v in self.value]
 
 
-class FileSegment(Segment):
-    value_type: SegmentType = SegmentType.FILE
-    # TODO: embed FileVar in this model.
-    value: FileVar
+class ArrayAnySegment(ArraySegment):
+    value_type: SegmentType = SegmentType.ARRAY_ANY
+    value: Sequence[Segment]
 
-    @property
-    def markdown(self) -> str:
-        return self.value.to_markdown()
+
+class ArrayStringSegment(ArraySegment):
+    value_type: SegmentType = SegmentType.ARRAY_STRING
+    value: Sequence[StringSegment]
+
+
+class ArrayNumberSegment(ArraySegment):
+    value_type: SegmentType = SegmentType.ARRAY_NUMBER
+    value: Sequence[FloatSegment | IntegerSegment]
+
+
+class ArrayObjectSegment(ArraySegment):
+    value_type: SegmentType = SegmentType.ARRAY_OBJECT
+    value: Sequence[ObjectSegment]
+
+
+class ArrayFileSegment(ArraySegment):
+    value_type: SegmentType = SegmentType.ARRAY_FILE
+    value: Sequence[FileSegment]

api/core/app/segments/types.py

@@ -6,7 +6,11 @@ class SegmentType(str, Enum):
     NUMBER = 'number'
     STRING = 'string'
     SECRET = 'secret'
-    ARRAY = 'array'
+    ARRAY_ANY = 'array[any]'
+    ARRAY_STRING = 'array[string]'
+    ARRAY_NUMBER = 'array[number]'
+    ARRAY_OBJECT = 'array[object]'
+    ARRAY_FILE = 'array[file]'
     OBJECT = 'object'
     FILE = 'file'
 

api/core/app/segments/variables.py

@@ -1,10 +1,13 @@
-
 from pydantic import Field
 
 from core.helper import encrypter
 
 from .segments import (
-    ArraySegment,
+    ArrayAnySegment,
+    ArrayFileSegment,
+    ArrayNumberSegment,
+    ArrayObjectSegment,
+    ArrayStringSegment,
     FileSegment,
     FloatSegment,
     IntegerSegment,
@@ -41,15 +44,31 @@ class IntegerVariable(IntegerSegment, Variable):
     pass
 
 
+class FileVariable(FileSegment, Variable):
+    pass
+
+
 class ObjectVariable(ObjectSegment, Variable):
     pass
 
 
-class ArrayVariable(ArraySegment, Variable):
+class ArrayAnyVariable(ArrayAnySegment, Variable):
     pass
 
 
-class FileVariable(FileSegment, Variable):
+class ArrayStringVariable(ArrayStringSegment, Variable):
+    pass
+
+
+class ArrayNumberVariable(ArrayNumberSegment, Variable):
+    pass
+
+
+class ArrayObjectVariable(ArrayObjectSegment, Variable):
+    pass
+
+
+class ArrayFileVariable(ArrayFileSegment, Variable):
     pass