From ce9b3bfbd33a6fa92a26024b85b334cb172c6cf4 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Tue, 6 Jun 2023 16:24:56 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=8E=A8=20refactor(schemas.py):=20add=20Fl?=
 =?UTF-8?q?owListRead=20schema=20and=20SQLModelSerializable=20base=20model?=
 =?UTF-8?q?=20The=20FlowListRead=20schema=20is=20added=20to=20support=20re?=
 =?UTF-8?q?ading=20a=20list=20of=20flows=20with=20their=20styles.=20The=20?=
 =?UTF-8?q?SQLModelSerializable=20base=20model=20is=20added=20to=20support?=
 =?UTF-8?q?=20serialization=20of=20SQLModel=20objects=20to=20JSON=20using?=
 =?UTF-8?q?=20orjson.=20This=20improves=20performance=20and=20reduces=20me?=
 =?UTF-8?q?mory=20usage.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

🐛 fix(flow.py): add optional style relationship to Flow model
The style relationship is now optional to allow for flows without styles. This is achieved by setting the uselist parameter of the sa_relationship_kwargs to False.

✨ feat(flow.py): add FlowReadWithStyle and FlowUpdate models
The FlowReadWithStyle model is added to support reading a flow with its style. The FlowUpdate model is added to support updating a flow.
---
 src/backend/langflow/api/v1/schemas.py       |  7 ++--
 src/backend/langflow/database/models/base.py | 14 ++++++++
 src/backend/langflow/database/models/flow.py | 35 +++++++++++---------
 3 files changed, 39 insertions(+), 17 deletions(-)
 create mode 100644 src/backend/langflow/database/models/base.py

diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py
index 7d5de2b63..d573a2ae2 100644
--- a/src/backend/langflow/api/v1/schemas.py
+++ b/src/backend/langflow/api/v1/schemas.py
@@ -1,6 +1,5 @@
 from typing import Any, Dict, List, Union
-from langflow.database.models.flow import FlowCreate
-
+from langflow.database.models.flow import FlowCreate, FlowRead
 from pydantic import BaseModel, validator
 
 
@@ -73,3 +72,7 @@ class FileResponse(ChatMessage):
 
 class FlowListCreate(BaseModel):
     flows: List[FlowCreate]
+
+
+class FlowListRead(BaseModel):
+    flows: List[FlowRead]
diff --git a/src/backend/langflow/database/models/base.py b/src/backend/langflow/database/models/base.py
new file mode 100644
index 000000000..e20895b93
--- /dev/null
+++ b/src/backend/langflow/database/models/base.py
@@ -0,0 +1,14 @@
+from sqlmodel import SQLModel
+import orjson
+
+
+def orjson_dumps(v, *, default):
+    # orjson.dumps returns bytes, to match standard json.dumps we need to decode
+    return orjson.dumps(v, default=default).decode()
+
+
+class SQLModelSerializable(SQLModel):
+    class Config:
+        orm_mode = True
+        json_loads = orjson.loads
+        json_dumps = orjson_dumps
diff --git a/src/backend/langflow/database/models/flow.py b/src/backend/langflow/database/models/flow.py
index 5ad49c668..0b71463a0 100644
--- a/src/backend/langflow/database/models/flow.py
+++ b/src/backend/langflow/database/models/flow.py
@@ -1,16 +1,18 @@
+# Path: src/backend/langflow/database/models/flow.py
+
+from langflow.database.models.base import SQLModelSerializable
 from pydantic import validator
-from sqlmodel import Field, SQLModel, Relationship, JSON, Column
+from sqlmodel import Field, Relationship, JSON, Column
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Dict
+from typing import Dict, Optional
 
-if TYPE_CHECKING:
-    from langflow.database.models.flowstyle import FlowStyle
+# if TYPE_CHECKING:
+from langflow.database.models.flow_style import FlowStyle, FlowStyleRead
 
 
-class FlowBase(SQLModel):
+class FlowBase(SQLModelSerializable):
     name: str = Field(index=True)
     flow: Dict = Field(default_factory=dict, sa_column=Column(JSON))
-    style: "FlowStyle" = Relationship(back_populates="flow")
 
     @validator("flow")
     def validate_json(v):
@@ -30,18 +32,12 @@ class FlowBase(SQLModel):
 
         return v
 
-    # @validator("flow")
-    # def flow_must_be_json(cls, v):
-    #     try:
-    #         valid_json = json.loads(v)
-
-    #     except Exception as e:
-    #         raise ValueError(f"Flow must be a valid JSON: {e}") from e
-    #     return v
-
 
 class Flow(FlowBase, table=True):
     id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True)
+    style: Optional["FlowStyle"] = Relationship(
+        back_populates="flow", sa_relationship_kwargs={"uselist": False}
+    )
 
 
 class FlowCreate(FlowBase):
@@ -50,3 +46,12 @@ class FlowCreate(FlowBase):
 
 class FlowRead(FlowBase):
     id: UUID
+
+
+class FlowReadWithStyle(FlowRead):
+    style: Optional["FlowStyleRead"] = None
+
+
+class FlowUpdate(SQLModelSerializable):
+    name: Optional[str] = None
+    flow: Optional[Dict] = None