Skip to content

Commit

Permalink
fix: update serialization and improve error handling (#3516)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
ogabrielluiz authored Aug 23, 2024
1 parent 7e7d5a1 commit 189ef41
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 7 deletions.
10 changes: 8 additions & 2 deletions src/backend/base/langflow/api/v1/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -327,13 +327,19 @@ async def build_vertices(
build_task = asyncio.create_task(await asyncio.to_thread(_build_vertex, vertex_id, graph))
try:
await build_task
except asyncio.CancelledError:
except asyncio.CancelledError as exc:
logger.exception(exc)
build_task.cancel()
return

vertex_build_response: VertexBuildResponse = build_task.result()
# 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()
if vertex_build_response.valid:
if vertex_build_response.next_vertices_ids:
Expand Down
10 changes: 9 additions & 1 deletion src/backend/base/langflow/api/v1/schemas.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,10 @@
from typing import Any, Dict, List, Optional, Union
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.utils import serialize_field
from langflow.schema import dotdict
from langflow.schema.graph import Tweaks
from langflow.schema.schema import InputType, OutputType, OutputValue
Expand Down Expand Up @@ -259,6 +260,13 @@ class ResultDataResponse(BaseModel):
duration: Optional[str] = None
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):
id: Optional[str] = None
Expand Down
14 changes: 10 additions & 4 deletions src/backend/base/langflow/graph/utils.py
Original file line number Diff line number Diff line change
@@ -1,25 +1,26 @@
import json
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 langchain_core.documents import Document
from loguru import logger
from pydantic import BaseModel
from pydantic.v1 import BaseModel as V1BaseModel

from langflow.interface.utils import extract_input_variables_from_prompt
from langflow.schema.data import Data
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.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.model import VertexBuildBase
from langflow.services.database.utils import session_getter
from langflow.services.deps import get_db_service, get_settings_service
from loguru import logger

if TYPE_CHECKING:
from langflow.graph.vertex.base import Vertex
from langflow.api.v1.schemas import ResultDataResponse
from langflow.graph.vertex.base import Vertex


class UnbuiltObject:
Expand Down Expand Up @@ -74,6 +75,11 @@ def serialize_field(value):
return value.to_json()
elif isinstance(value, BaseModel):
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):
return {"result": value}
return value
Expand Down

0 comments on commit 189ef41

Please sign in to comment.