Skip to content

Commit

Permalink
refactor: Update field serializer function and error handling in buil…
Browse files Browse the repository at this point in the history
…d_flow function
  • Loading branch information
ogabrielluiz committed Aug 23, 2024
1 parent d30da90 commit bc1d8b7
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 3 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

0 comments on commit bc1d8b7

Please sign in to comment.