diff --git a/src/backend/langflow/api/v1/flows.py b/src/backend/langflow/api/v1/flows.py index d047cbcd5..f560b1105 100644 --- a/src/backend/langflow/api/v1/flows.py +++ b/src/backend/langflow/api/v1/flows.py @@ -43,13 +43,13 @@ def read_flows(*, session: Session = Depends(get_session)): @router.get("/{flow_id}", response_model=FlowReadWithStyle) def read_flow(*, session: Session = Depends(get_session), flow_id: UUID): """Read a flow.""" - flow = session.get(Flow, flow_id) - if not flow: + if flow := session.get(Flow, flow_id): + return flow + else: raise HTTPException(status_code=404, detail="Flow not found") - return flow -@router.patch("/{flow_id}", response_model=FlowRead) +@router.patch("/{flow_id}", response_model=FlowReadWithStyle) def update_flow( *, session: Session = Depends(get_session), flow_id: UUID, flow: FlowUpdate ): diff --git a/src/backend/langflow/database/models/flow.py b/src/backend/langflow/database/models/flow.py index 6305de3cc..f9e3aa249 100644 --- a/src/backend/langflow/database/models/flow.py +++ b/src/backend/langflow/database/models/flow.py @@ -12,24 +12,21 @@ from langflow.database.models.flow_style import FlowStyle, FlowStyleRead class FlowBase(SQLModelSerializable): name: str = Field(index=True) - flow: Optional[Dict] = Field(default_factory=dict, sa_column=Column(JSON)) + description: Optional[str] = Field(index=True) + data: Optional[Dict] = Field(default=None) - @validator("flow") + @validator("data") def validate_json(v): # dict_keys(['description', 'name', 'id', 'data']) if not v: return v 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(): + if "nodes" not in v.keys(): raise ValueError("Flow must have nodes") - if "edges" not in v["data"].keys(): + if "edges" not in v.keys(): raise ValueError("Flow must have edges") return v @@ -37,8 +34,11 @@ class FlowBase(SQLModelSerializable): class Flow(FlowBase, table=True): id: UUID = Field(default_factory=uuid4, primary_key=True, unique=True) + data: Optional[Dict] = Field(default=None, sa_column=Column(JSON)) style: Optional["FlowStyle"] = Relationship( - back_populates="flow", sa_relationship_kwargs={"uselist": False} + back_populates="flow", + # use "uselist=False" to make it a one-to-one relationship + sa_relationship_kwargs={"uselist": False}, ) @@ -56,4 +56,5 @@ class FlowReadWithStyle(FlowRead): class FlowUpdate(SQLModelSerializable): name: Optional[str] = None - flow: Optional[Dict] = None + description: Optional[str] = None + data: Optional[Dict] = None diff --git a/tests/test_database.py b/tests/test_database.py index 516c04ade..1e449f373 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1,6 +1,6 @@ from uuid import UUID, uuid4 from langflow.api.v1.schemas import FlowListCreate -from langflow.database.models.flow import FlowCreate +from langflow.database.models.flow import FlowCreate, FlowUpdate import json from sqlalchemy.orm import Session from langflow.database.models.flow import Flow @@ -33,25 +33,29 @@ def json_style(): def test_create_flow(client: TestClient, json_flow: str): - flow = FlowCreate(name="Test Flow", flow=json.loads(json_flow)) + flow = json.loads(json_flow) + data = flow["data"] + flow = FlowCreate(name="Test Flow", description="description", data=data) response = client.post("api/v1/flows/", json=flow.dict()) assert response.status_code == 200 assert response.json()["name"] == flow.name - assert response.json()["flow"] == flow.flow + assert response.json()["data"] == flow.data # flow is optional so we can create a flow without a flow flow = FlowCreate(name="Test Flow") response = client.post("api/v1/flows/", json=flow.dict(exclude_unset=True)) assert response.status_code == 200 assert response.json()["name"] == flow.name - assert response.json()["flow"] == flow.flow + assert response.json()["data"] == flow.data def test_read_flows(client: TestClient, json_flow: str): - flow = FlowCreate(name="Test Flow", flow=json.loads(json_flow)) + flow_data = json.loads(json_flow) + data = flow_data["data"] + flow = FlowCreate(name="Test Flow", description="description", data=data) response = client.post("api/v1/flows/", json=flow.dict()) assert response.status_code == 200 assert response.json()["name"] == flow.name - assert response.json()["flow"] == flow.flow + assert response.json()["data"] == flow.data flow_style = FlowStyleCreate(color="red", emoji="👍", flow_id=response.json()["id"]) response = client.post( @@ -62,11 +66,11 @@ def test_read_flows(client: TestClient, json_flow: str): assert response.json()["emoji"] == flow_style.emoji assert response.json()["flow_id"] == str(flow_style.flow_id) - flow = FlowCreate(name="Test Flow", flow=json.loads(json_flow)) + flow = FlowCreate(name="Test Flow", description="description", data=data) response = client.post("api/v1/flows/", json=flow.dict()) assert response.status_code == 200 assert response.json()["name"] == flow.name - assert response.json()["flow"] == flow.flow + assert response.json()["data"] == flow.data # Now we need to create FlowStyle objects for each Flow flow_style = FlowStyleCreate( @@ -86,7 +90,9 @@ def test_read_flows(client: TestClient, json_flow: str): def test_read_flow(client: TestClient, json_flow: str): - flow = FlowCreate(name="Test Flow", flow=json.loads(json_flow)) + flow = json.loads(json_flow) + data = flow["data"] + flow = FlowCreate(name="Test Flow", description="description", data=data) response = client.post("api/v1/flows/", json=flow.dict()) flow_id = response.json()["id"] # flow_id should be a UUID but is a string # turn it into a UUID @@ -105,26 +111,32 @@ def test_read_flow(client: TestClient, json_flow: str): response = client.get(f"api/v1/flows/{flow_id}") assert response.status_code == 200 assert response.json()["name"] == flow.name - assert response.json()["flow"] == flow.flow + assert response.json()["data"] == flow.data assert response.json()["style"]["color"] == flow_style.color def test_update_flow(client: TestClient, json_flow: str): - flow = FlowCreate(name="Test Flow", flow=json.loads(json_flow)) + flow = json.loads(json_flow) + data = flow["data"] + flow = FlowCreate(name="Test Flow", description="description", data=data) response = client.post("api/v1/flows/", json=flow.dict()) flow_id = response.json()["id"] - updated_flow = FlowCreate( + updated_flow = FlowUpdate( name="Updated Flow", - flow=json.loads(json_flow.replace("BasicExample", "Updated Flow")), + description="updated description", + data=data, ) response = client.patch(f"api/v1/flows/{flow_id}", json=updated_flow.dict()) assert response.status_code == 200 assert response.json()["name"] == updated_flow.name - assert response.json()["flow"] == updated_flow.flow + assert response.json()["description"] == updated_flow.description + assert response.json()["data"] == updated_flow.data def test_delete_flow(client: TestClient, json_flow: str): - flow = FlowCreate(name="Test Flow", flow=json.loads(json_flow)) + flow = json.loads(json_flow) + data = flow["data"] + flow = FlowCreate(name="Test Flow", description="description", data=data) response = client.post("api/v1/flows/", json=flow.dict()) flow_id = response.json()["id"] response = client.delete(f"api/v1/flows/{flow_id}") @@ -133,11 +145,13 @@ def test_delete_flow(client: TestClient, json_flow: str): def test_create_flows(client: TestClient, session: Session, json_flow: str): + flow = json.loads(json_flow) + data = flow["data"] # Create test data flow_list = FlowListCreate( flows=[ - FlowCreate(name="Flow 1", flow=json.loads(json_flow)), - FlowCreate(name="Flow 2", flow=json.loads(json_flow)), + FlowCreate(name="Flow 1", description="description", data=data), + FlowCreate(name="Flow 2", description="description", data=data), ] ) # Make request to endpoint @@ -148,17 +162,21 @@ def test_create_flows(client: TestClient, session: Session, json_flow: str): response_data = response.json() assert len(response_data) == 2 assert response_data[0]["name"] == "Flow 1" - assert response_data[0]["flow"] == json.loads(json_flow) + assert response_data[0]["description"] == "description" + assert response_data[0]["data"] == data assert response_data[1]["name"] == "Flow 2" - assert response_data[1]["flow"] == json.loads(json_flow) + assert response_data[1]["description"] == "description" + assert response_data[1]["data"] == data def test_upload_file(client: TestClient, session: Session, json_flow: str): + flow = json.loads(json_flow) + data = flow["data"] # Create test data flow_list = FlowListCreate( flows=[ - FlowCreate(name="Flow 1", flow=json.loads(json_flow)), - FlowCreate(name="Flow 2", flow=json.loads(json_flow)), + FlowCreate(name="Flow 1", description="description", data=data), + FlowCreate(name="Flow 2", description="description", data=data), ] ) file_contents = json.dumps(flow_list.dict()) @@ -172,17 +190,21 @@ def test_upload_file(client: TestClient, session: Session, json_flow: str): response_data = response.json() assert len(response_data) == 2 assert response_data[0]["name"] == "Flow 1" - assert response_data[0]["flow"] == json.loads(json_flow) + assert response_data[0]["description"] == "description" + assert response_data[0]["data"] == data assert response_data[1]["name"] == "Flow 2" - assert response_data[1]["flow"] == json.loads(json_flow) + assert response_data[1]["description"] == "description" + assert response_data[1]["data"] == data def test_download_file(client: TestClient, session: Session, json_flow): + flow = json.loads(json_flow) + data = flow["data"] # Create test data flow_list = FlowListCreate( flows=[ - FlowCreate(name="Flow 1", flow=json.loads(json_flow)), - FlowCreate(name="Flow 2", flow=json.loads(json_flow)), + FlowCreate(name="Flow 1", description="description", data=data), + FlowCreate(name="Flow 2", description="description", data=data), ] ) for flow in flow_list.flows: @@ -197,13 +219,15 @@ def test_download_file(client: TestClient, session: Session, json_flow): response_data = response.json()["flows"] assert len(response_data) == 2 assert response_data[0]["name"] == "Flow 1" - assert response_data[0]["flow"] == json.loads(json_flow) + assert response_data[0]["description"] == "description" + assert response_data[0]["data"] == data assert response_data[1]["name"] == "Flow 2" - assert response_data[1]["flow"] == json.loads(json_flow) + assert response_data[1]["description"] == "description" + assert response_data[1]["data"] == data def test_create_flow_with_invalid_data(client: TestClient): - flow = {"name": "a" * 256, "flow": "Invalid flow data"} + flow = {"name": "a" * 256, "data": "Invalid flow data"} response = client.post("api/v1/flows/", json=flow) assert response.status_code == 422 @@ -216,20 +240,25 @@ def test_get_nonexistent_flow(client: TestClient): def test_update_flow_idempotency(client: TestClient, json_flow: str): - flow = FlowCreate(name="Test Flow", flow=json.loads(json_flow)) - response = client.post("api/v1/flows/", json=flow.dict()) + flow_data = json.loads(json_flow) + data = flow_data["data"] + flow_data = FlowCreate(name="Test Flow", description="description", data=data) + response = client.post("api/v1/flows/", json=flow_data.dict()) flow_id = response.json()["id"] - updated_flow = FlowCreate(name="Updated Flow", flow=json.loads(json_flow)) + updated_flow = FlowCreate(name="Updated Flow", description="description", data=data) response1 = client.put(f"api/v1/flows/{flow_id}", json=updated_flow.dict()) response2 = client.put(f"api/v1/flows/{flow_id}", json=updated_flow.dict()) assert response1.json() == response2.json() def test_update_nonexistent_flow(client: TestClient, json_flow: str): + flow_data = json.loads(json_flow) + data = flow_data["data"] uuid = uuid4() updated_flow = FlowCreate( name="Updated Flow", - flow=json.loads(json_flow.replace("BasicExample", "Updated Flow")), + description="description", + data=data, ) response = client.patch(f"api/v1/flows/{uuid}", json=updated_flow.dict()) assert response.status_code == 404 @@ -248,7 +277,9 @@ def test_read_empty_flows(client: TestClient): def test_stress_create_flow(client: TestClient, json_flow: str): - flow = FlowCreate(name="Test Flow", flow=json.loads(json_flow)) + flow_data = json.loads(json_flow) + data = flow_data["data"] + flow = FlowCreate(name="Test Flow", description="description", data=data) def create_flow(): response = client.post("api/v1/flows/", json=flow.dict())