🎨 refactor(schemas.py): add FlowListRead schema and SQLModelSerializable base model

The FlowListRead schema is added to support reading a list of flows with their styles. The SQLModelSerializable base model is added to support serialization of SQLModel objects to JSON using orjson. This improves performance and reduces memory usage.

🐛 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.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-06 16:24:56 -03:00
commit ce9b3bfbd3
3 changed files with 39 additions and 17 deletions

View file

@ -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]

View file

@ -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

View file

@ -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