Update schemas.py with serialization functions

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-07 14:07:23 -03:00
commit 58ee61d0a2

View file

@ -4,12 +4,12 @@ from pathlib import Path
from typing import Any, Dict, List, Optional, Union
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.base import orjson_dumps
from langflow.services.database.models.flow import FlowCreate, FlowRead
from langflow.services.database.models.user import UserRead
from pydantic import BaseModel, Field, field_serializer, field_validator
class BuildStatus(Enum):
@ -221,13 +221,41 @@ class VerticesOrderResponse(BaseModel):
class ResultDict(BaseModel):
"""Outputs of the vertex build process."""
results: Optional[Any] = Field(default_factory=dict)
artifacts: Optional[Any] = Field(default_factory=dict)
timedelta: Optional[float] = 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):
id: Optional[str] = None
@ -242,3 +270,4 @@ class VertexBuildResponse(BaseModel):
class VerticesBuiltResponse(BaseModel):
vertices: List[VertexBuildResponse]
vertices: List[VertexBuildResponse]