🐛 fix(database.py): handle case where data does not contain "flows" key

✨ feat(database.py): add default argument to json.dumps to handle datetime objects
🚨 test(database.py): add tests for batch flow creation, file upload, and file download
The fix in database.py handles the case where the data dictionary does not contain the "flows" key. This is important because the code assumes that the "flows" key is present and will raise an exception if it is not. The fix adds a check to see if the "flows" key is present and if not, it creates a new FlowListCreate object with the data as a list of FlowCreate objects.

The feature in database.py adds a default argument to the json.dumps function to handle datetime objects. This is important because the default json encoder does not handle datetime objects and will raise an exception if it encounters one.

The tests in test_database.py cover the batch creation of flows, uploading a file containing flows, and downloading a file containing flows. These tests ensure that the endpoints are working as expected and that the data is being handled correctly.

src/backend/langflow/api/database.py

@@ -92,7 +92,10 @@ async def upload_file(
     """Upload flows from a file."""
     contents = await file.read()
     data = json.loads(contents)
-    flow_list = FlowListCreate(**data)
+    if "flows" in data:
+        flow_list = FlowListCreate(**data)
+    else:
+        flow_list = FlowListCreate(flows=[FlowCreate(**flow) for flow in data])
     return create_flows(session=session, flow_list=flow_list)
 
 
@@ -100,4 +103,4 @@ async def upload_file(
 async def download_file(*, session: Session = Depends(get_session)):
     """Download all flows as a file."""
     flows = read_flows(session=session)
-    return {"file": json.dumps([flow.dict() for flow in flows])}
+    return {"file": json.dumps([flow.dict() for flow in flows], default=str)}

src/backend/langflow/api/endpoints.py

@@ -35,7 +35,7 @@ async def get_load(
         if flow_obj is None:
             raise ValueError(f"Flow {flow_id} not found")
         graph_data: GraphData = json.loads(flow_obj.flow)
-        data = graph_data.dict()
+        data = graph_data.get("data")
         response = process_graph_cached(data, predict_request.message)
         return PredictResponse(
             result=response.get("result", ""),

tests/test_database.py

@@ -1,4 +1,8 @@
+from langflow.api.schemas import FlowListCreate
 from langflow.database.models.flow import FlowCreate
+import json
+from sqlalchemy.orm import Session
+from langflow.database.models.flow import Flow
 
 
 def test_create_flow(client, json_flow):
@@ -56,3 +60,79 @@ def test_delete_flow(client, json_flow):
     response = client.delete(f"/flows/{flow_id}")
     assert response.status_code == 200
     assert response.json()["message"] == "Flow deleted successfully"
+
+
+def test_create_flows(client, session: Session):
+    # Create test data
+    flow_list = FlowListCreate(
+        flows=[
+            FlowCreate(name="Flow 1", flow="Test flow 1"),
+            FlowCreate(name="Flow 2", flow="Test flow 2"),
+        ]
+    )
+    # Make request to endpoint
+    response = client.post("/flows/batch/", json=flow_list.dict())
+    # Check response status code
+    assert response.status_code == 200
+    # Check response data
+    response_data = response.json()
+    assert len(response_data) == 2
+    assert response_data[0]["name"] == "Flow 1"
+    assert response_data[0]["flow"] == "Test flow 1"
+    assert response_data[1]["name"] == "Flow 2"
+    assert response_data[1]["flow"] == "Test flow 2"
+
+
+def test_upload_file(client, session: Session):
+    # Create test data
+    flow_list = FlowListCreate(
+        flows=[
+            FlowCreate(name="Flow 1", flow="Test flow 1"),
+            FlowCreate(name="Flow 2", flow="Test flow 2"),
+        ]
+    )
+    file_contents = json.dumps(flow_list.dict())
+    # Make request to endpoint
+    #     curl -X 'POST' \
+    #   'http://127.0.0.1:7860/flows/upload/' \
+    #   -H 'accept: application/json' \
+    #   -H 'Content-Type: multipart/form-data' \
+    #   -F 'file=@examples.json;type=application/json'
+    response = client.post(
+        "/flows/upload/",
+        files={"file": ("examples.json", file_contents, "application/json")},
+    )
+    # Check response status code
+    assert response.status_code == 200
+    # Check response data
+    response_data = response.json()
+    assert len(response_data) == 2
+    assert response_data[0]["name"] == "Flow 1"
+    assert response_data[0]["flow"] == "Test flow 1"
+    assert response_data[1]["name"] == "Flow 2"
+    assert response_data[1]["flow"] == "Test flow 2"
+
+
+def test_download_file(client, session: Session):
+    # Create test data
+    flow_list = FlowListCreate(
+        flows=[
+            FlowCreate(name="Flow 1", flow="Test flow 1"),
+            FlowCreate(name="Flow 2", flow="Test flow 2"),
+        ]
+    )
+    for flow in flow_list.flows:
+        db_flow = Flow.from_orm(flow)
+        session.add(db_flow)
+    session.commit()
+    # Make request to endpoint
+    response = client.get("/flows/download/")
+    # Check response status code
+    assert response.status_code == 200
+    # Check response data
+    response_data = json.loads(response.json()["file"])
+    assert len(response_data) == 2
+    assert response_data[0]["name"] == "Flow 1"
+    assert response_data[0]["flow"] == "Test flow 1"
+    assert response_data[1]["name"] == "Flow 2"
+    assert response_data[1]["flow"] == "Test flow 2"