From 8c2dd4371326a2099a428bc0be00b17d6f9aa9c5 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Fri, 23 Aug 2024 10:30:42 -0300
Subject: [PATCH] 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
---
 src/backend/base/langflow/api/v1/chat.py    | 10 ++++++++--
 src/backend/base/langflow/api/v1/schemas.py | 10 +++++++++-
 src/backend/base/langflow/graph/utils.py    | 14 ++++++++++----
 3 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/src/backend/base/langflow/api/v1/chat.py b/src/backend/base/langflow/api/v1/chat.py
index 2156b62b9ba7..5e47861aa13d 100644
--- a/src/backend/base/langflow/api/v1/chat.py
+++ b/src/backend/base/langflow/api/v1/chat.py
@@ -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:
diff --git a/src/backend/base/langflow/api/v1/schemas.py b/src/backend/base/langflow/api/v1/schemas.py
index 8d31b4d71e5e..4c1e3726fab7 100644
--- a/src/backend/base/langflow/api/v1/schemas.py
+++ b/src/backend/base/langflow/api/v1/schemas.py
@@ -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
@@ -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
diff --git a/src/backend/base/langflow/graph/utils.py b/src/backend/base/langflow/graph/utils.py
index 41b20c135db0..bcc784a10f44 100644
--- a/src/backend/base/langflow/graph/utils.py
+++ b/src/backend/base/langflow/graph/utils.py
@@ -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:
@@ -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