Add serialize_field function to handle serialization of BaseModel and Document types

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-08 19:42:47 -03:00
commit 1d31ee753b
4 changed files with 27 additions and 30 deletions

View file

@ -3,7 +3,9 @@ from pathlib import Path
from typing import TYPE_CHECKING, List
from fastapi import HTTPException
from platformdirs import user_cache_dir
from langchain_core.documents import Document
from langflow.services.store.schema import StoreComponentCreate
from pydantic import BaseModel
from langflow.services.store.schema import StoreComponentCreate
from langflow.services.store.utils import get_lf_version_from_pypi
@ -191,3 +193,17 @@ def format_elapsed_time(elapsed_time: float) -> str:
minutes_unit = "minute" if minutes == 1 else "minutes"
seconds_unit = "second" if seconds == 1 else "seconds"
return f"{minutes} {minutes_unit}, {seconds} {seconds_unit}"
def serialize_field(value):
"""Unified serialization function for handling both BaseModel and Document types,
including handling lists of these types."""
if isinstance(value, (list, tuple)):
return [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

View file

@ -4,7 +4,7 @@ from pathlib import Path
from typing import Any, Dict, List, Optional, Union
from uuid import UUID
from langchain_core.documents import Document
from langflow.api.utils import serialize_field
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
@ -226,35 +226,11 @@ class ResultDict(BaseModel):
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
return {key: serialize_field(val) for key, val in value.items()}
return serialize_field(value)
class VertexBuildResponse(BaseModel):

View file

@ -3,6 +3,7 @@ import json
from typing import Callable, Dict, List, Optional, Union
from langchain_core.messages import AIMessage
from langflow.graph.utils import UnbuiltObject, flatten_list
from langflow.graph.vertex.base import StatefulVertex, StatelessVertex
from langflow.interface.utils import extract_input_variables_from_prompt
@ -354,5 +355,8 @@ class ChatVertex(StatelessVertex):
def dict_to_codeblock(d: dict) -> str:
json_str = json.dumps(d, indent=4)
from langflow.api.utils import serialize_field
serialized = {key: serialize_field(val) for key, val in d.items()}
json_str = json.dumps(serialized, indent=4)
return f"```json\n{json_str}\n```"

View file

@ -1,10 +1,11 @@
from typing import TYPE_CHECKING, Any, Dict, Optional, Type
import duckdb
from langflow.services.deps import get_monitor_service
from loguru import logger
from pydantic import BaseModel
from langflow.services.deps import get_monitor_service
if TYPE_CHECKING:
from langflow.api.v1.schemas import ResultDict