🐛 fix(flow.py): change flow field type from str to Dict to allow for JSON data

✨ feat(flow.py): add validator to ensure flow field is a valid JSON object with required fields
The flow field in the FlowBase model has been changed from a string to a dictionary to allow for JSON data. A validator has been added to ensure that the flow field is a valid JSON object with the required fields. The tests have been updated to reflect these changes.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-05 13:40:14 -03:00
commit ec585b8acc
3 changed files with 126 additions and 38 deletions

View file

@ -1,6 +1,7 @@
from sqlmodel import Field, SQLModel, Relationship
from pydantic import validator
from sqlmodel import Field, SQLModel, Relationship, JSON, Column
from uuid import UUID, uuid4
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Dict
if TYPE_CHECKING:
from langflow.database.models.flowstyle import FlowStyle
@ -8,9 +9,36 @@ if TYPE_CHECKING:
class FlowBase(SQLModel):
name: str = Field(index=True)
flow: str = Field(index=False)
flow: Dict = Field(default_factory=dict, sa_column=Column(JSON))
style: "FlowStyle" = Relationship(back_populates="flow")
@validator("flow")
def validate_json(v):
# dict_keys(['description', 'name', 'id', 'data'])
if not isinstance(v, dict):
raise ValueError("Flow must be a valid JSON")
if "description" not in v.keys():
raise ValueError("Flow must have a description")
if "data" not in v.keys():
raise ValueError("Flow must have data")
# data must contain nodes and edges
if "nodes" not in v["data"].keys():
raise ValueError("Flow must have nodes")
if "edges" not in v["data"].keys():
raise ValueError("Flow must have edges")
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)