Update schemas.py with serialization functions
This commit is contained in:
parent
0abeda51e9
commit
58ee61d0a2
1 changed files with 33 additions and 4 deletions
|
|
@ -4,12 +4,12 @@ from pathlib import Path
|
||||||
from typing import Any, Dict, List, Optional, Union
|
from typing import Any, Dict, List, Optional, Union
|
||||||
from uuid import UUID
|
from uuid import UUID
|
||||||
|
|
||||||
from pydantic import BaseModel, Field, field_validator
|
from langchain_core.documents import Document
|
||||||
|
|
||||||
from langflow.services.database.models.api_key.model import ApiKeyRead
|
from langflow.services.database.models.api_key.model import ApiKeyRead
|
||||||
from langflow.services.database.models.base import orjson_dumps
|
from langflow.services.database.models.base import orjson_dumps
|
||||||
from langflow.services.database.models.flow import FlowCreate, FlowRead
|
from langflow.services.database.models.flow import FlowCreate, FlowRead
|
||||||
from langflow.services.database.models.user import UserRead
|
from langflow.services.database.models.user import UserRead
|
||||||
|
from pydantic import BaseModel, Field, field_serializer, field_validator
|
||||||
|
|
||||||
|
|
||||||
class BuildStatus(Enum):
|
class BuildStatus(Enum):
|
||||||
|
|
@ -221,13 +221,41 @@ class VerticesOrderResponse(BaseModel):
|
||||||
|
|
||||||
|
|
||||||
class ResultDict(BaseModel):
|
class ResultDict(BaseModel):
|
||||||
"""Outputs of the vertex build process."""
|
|
||||||
|
|
||||||
results: Optional[Any] = Field(default_factory=dict)
|
results: Optional[Any] = Field(default_factory=dict)
|
||||||
artifacts: Optional[Any] = Field(default_factory=dict)
|
artifacts: Optional[Any] = Field(default_factory=dict)
|
||||||
timedelta: Optional[float] = None
|
timedelta: Optional[float] = None
|
||||||
duration: Optional[str] = None
|
duration: Optional[str] = None
|
||||||
|
|
||||||
|
def serialize_field(self, value):
|
||||||
|
"""Unified serialization function for handling both BaseModel and Document types,
|
||||||
|
including handling lists of these types."""
|
||||||
|
if isinstance(value, (list, tuple)):
|
||||||
|
return [self.serialize_field(v) for v in value]
|
||||||
|
elif isinstance(value, Document):
|
||||||
|
return value.to_json()
|
||||||
|
elif isinstance(value, BaseModel):
|
||||||
|
return value.model_dump()
|
||||||
|
elif isinstance(value, str):
|
||||||
|
return {"result": value}
|
||||||
|
return value
|
||||||
|
|
||||||
|
@field_serializer("results")
|
||||||
|
def serialize_results(self, value):
|
||||||
|
if isinstance(value, dict):
|
||||||
|
return {key: self.serialize_field(val) for key, val in value.items()}
|
||||||
|
return self.serialize_field(value)
|
||||||
|
|
||||||
|
|
||||||
|
def serialize_list_of_documents_or_base_models(value):
|
||||||
|
if isinstance(value, list):
|
||||||
|
for i, val in enumerate(value):
|
||||||
|
if isinstance(val, Document):
|
||||||
|
value[i] = val.to_json()
|
||||||
|
elif isinstance(val, BaseModel):
|
||||||
|
value[i] = val.model_dump()
|
||||||
|
|
||||||
|
return value
|
||||||
|
|
||||||
|
|
||||||
class VertexBuildResponse(BaseModel):
|
class VertexBuildResponse(BaseModel):
|
||||||
id: Optional[str] = None
|
id: Optional[str] = None
|
||||||
|
|
@ -242,3 +270,4 @@ class VertexBuildResponse(BaseModel):
|
||||||
|
|
||||||
class VerticesBuiltResponse(BaseModel):
|
class VerticesBuiltResponse(BaseModel):
|
||||||
vertices: List[VertexBuildResponse]
|
vertices: List[VertexBuildResponse]
|
||||||
|
vertices: List[VertexBuildResponse]
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue