fix: update serialization and improve error handling (#3516)

* feat(utils): add support for V1BaseModel in serialize_field

Add support for V1BaseModel instances in the serialize_field function by
checking for a "to_json" method. If the method is not present, return the
attribute values as a dictionary.

* refactor: Update field serializer function and error handling in build_flow function
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-08-23 10:30:42 -03:00 • committed by GitHub
commit 189ef4147d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 27 additions and 7 deletions

View file

@ -327,13 +327,19 @@ async def build_flow(
build_task = asyncio.create_task(await asyncio.to_thread(_build_vertex, vertex_id, graph)) build_task = asyncio.create_task(await asyncio.to_thread(_build_vertex, vertex_id, graph))
try: try:
await build_task await build_task
except asyncio.CancelledError: except asyncio.CancelledError as exc:
logger.exception(exc)
build_task.cancel() build_task.cancel()
return return
vertex_build_response: VertexBuildResponse = build_task.result() vertex_build_response: VertexBuildResponse = build_task.result()
# send built event or error event # send built event or error event
send_event("end_vertex", {"build_data": json.loads(vertex_build_response.model_dump_json())}, queue) try:
vertex_build_response_json = vertex_build_response.model_dump_json()
build_data = json.loads(vertex_build_response_json)
except Exception as exc:
raise ValueError(f"Error serializing vertex build response: {exc}") from exc
send_event("end_vertex", {"build_data": build_data}, queue)
await client_consumed_queue.get() await client_consumed_queue.get()
if vertex_build_response.valid: if vertex_build_response.valid:
if vertex_build_response.next_vertices_ids: if vertex_build_response.next_vertices_ids:

View file

@ -4,9 +4,10 @@ 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, ConfigDict, Field, field_validator, model_serializer from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator, model_serializer
from langflow.graph.schema import RunOutputs from langflow.graph.schema import RunOutputs
from langflow.graph.utils import serialize_field
from langflow.schema import dotdict from langflow.schema import dotdict
from langflow.schema.graph import Tweaks from langflow.schema.graph import Tweaks
from langflow.schema.schema import InputType, OutputType, OutputValue from langflow.schema.schema import InputType, OutputType, OutputValue
@ -259,6 +260,13 @@ class ResultDataResponse(BaseModel):
duration: Optional[str] = None duration: Optional[str] = None
used_frozen_result: Optional[bool] = False used_frozen_result: Optional[bool] = False
@field_serializer("results")
@classmethod
def serialize_results(cls, v):
if isinstance(v, dict):
return {key: serialize_field(val) for key, val in v.items()}
return serialize_field(v)
class VertexBuildResponse(BaseModel): class VertexBuildResponse(BaseModel):
id: Optional[str] = None id: Optional[str] = None

View file

@ -1,25 +1,26 @@
import json import json
from enum import Enum from enum import Enum
from typing import TYPE_CHECKING, Any, Generator, Union, Optional from typing import TYPE_CHECKING, Any, Generator, Optional, Union
from uuid import UUID from uuid import UUID
from langchain_core.documents import Document from langchain_core.documents import Document
from loguru import logger
from pydantic import BaseModel from pydantic import BaseModel
from pydantic.v1 import BaseModel as V1BaseModel
from langflow.interface.utils import extract_input_variables_from_prompt from langflow.interface.utils import extract_input_variables_from_prompt
from langflow.schema.data import Data from langflow.schema.data import Data
from langflow.schema.message import Message from langflow.schema.message import Message
from langflow.services.database.models.transactions.model import TransactionBase
from langflow.services.database.models.transactions.crud import log_transaction as crud_log_transaction from langflow.services.database.models.transactions.crud import log_transaction as crud_log_transaction
from langflow.services.database.models.transactions.model import TransactionBase
from langflow.services.database.models.vertex_builds.crud import log_vertex_build as crud_log_vertex_build from langflow.services.database.models.vertex_builds.crud import log_vertex_build as crud_log_vertex_build
from langflow.services.database.models.vertex_builds.model import VertexBuildBase from langflow.services.database.models.vertex_builds.model import VertexBuildBase
from langflow.services.database.utils import session_getter from langflow.services.database.utils import session_getter
from langflow.services.deps import get_db_service, get_settings_service from langflow.services.deps import get_db_service, get_settings_service
from loguru import logger
if TYPE_CHECKING: if TYPE_CHECKING:
from langflow.graph.vertex.base import Vertex
from langflow.api.v1.schemas import ResultDataResponse from langflow.api.v1.schemas import ResultDataResponse
from langflow.graph.vertex.base import Vertex
class UnbuiltObject: class UnbuiltObject:
@ -74,6 +75,11 @@ def serialize_field(value):
return value.to_json() return value.to_json()
elif isinstance(value, BaseModel): elif isinstance(value, BaseModel):
return value.model_dump() return value.model_dump()
elif isinstance(value, V1BaseModel):
if hasattr(value, "to_json"):
return value.to_json()
else:
return value.dict()
elif isinstance(value, str): elif isinstance(value, str):
return {"result": value} return {"result": value}
return value return value