From 1d31ee753bc2ebf4f1d15dfeaa50231d6c4c76d7 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Thu, 8 Feb 2024 19:42:47 -0300
Subject: [PATCH] Add serialize_field function to handle serialization of
 BaseModel and Document types

---
 src/backend/langflow/api/utils.py             | 18 ++++++++++-
 src/backend/langflow/api/v1/schemas.py        | 30 ++-----------------
 src/backend/langflow/graph/vertex/types.py    |  6 +++-
 .../langflow/services/monitor/utils.py        |  3 +-
 4 files changed, 27 insertions(+), 30 deletions(-)

diff --git a/src/backend/langflow/api/utils.py b/src/backend/langflow/api/utils.py
index 2b47eee86..a384ccbbf 100644
--- a/src/backend/langflow/api/utils.py
+++ b/src/backend/langflow/api/utils.py
@@ -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
diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py
index 1e2e9bffa..066934024 100644
--- a/src/backend/langflow/api/v1/schemas.py
+++ b/src/backend/langflow/api/v1/schemas.py
@@ -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):
diff --git a/src/backend/langflow/graph/vertex/types.py b/src/backend/langflow/graph/vertex/types.py
index 65da2bff9..b350770f1 100644
--- a/src/backend/langflow/graph/vertex/types.py
+++ b/src/backend/langflow/graph/vertex/types.py
@@ -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```"
diff --git a/src/backend/langflow/services/monitor/utils.py b/src/backend/langflow/services/monitor/utils.py
index 9bd32e243..5324700a7 100644
--- a/src/backend/langflow/services/monitor/utils.py
+++ b/src/backend/langflow/services/monitor/utils.py
@@ -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