test(backend): added csv and pdf upload

backend/main.py

@@ -30,7 +30,7 @@ if os.getenv("SENTRY_DSN"):
 app = FastAPI()
 
 add_cors_middleware(app)
-max_brain_size = os.getenv("MAX_BRAIN_SIZE")
+max_brain_size = os.getenv("MAX_BRAIN_SIZE", 52428800)
 max_brain_size_with_own_key = os.getenv("MAX_BRAIN_SIZE_WITH_KEY", 209715200)
 
 

backend/models/brains.py

@@ -18,7 +18,7 @@ class Brain(BaseModel):
     model: Optional[str] = "gpt-3.5-turbo-0613"
     temperature: Optional[float] = 0.0
     max_tokens: Optional[int] = 256
-    max_brain_size: Optional[int] = int(os.getenv("MAX_BRAIN_SIZE", 0))
+    max_brain_size: Optional[int] = int(os.getenv("MAX_BRAIN_SIZE",52428800))
     files: List[Any] = []
 
     class Config:

backend/routes/user_routes.py

@@ -32,7 +32,7 @@ async def get_user_endpoint(
     information about the user's API usage.
     """
 
-    max_brain_size = int(os.getenv("MAX_BRAIN_SIZE", 0))
+    max_brain_size = int(os.getenv("MAX_BRAIN_SIZE", 52428800))
     if request.headers.get("Openai-Api-Key"):
         max_brain_size = MAX_BRAIN_SIZE_WITH_OWN_KEY
 

backend/test_file/test.csv

@@ -0,0 +1,17 @@
+StanGirard/quivr,Sat May 13 2023 02:20:09 GMT+0200 (heure d’été d’Europe centrale),0
+StanGirard/quivr,Tue May 16 2023 18:03:49 GMT+0200 (heure d’été d’Europe centrale),660
+StanGirard/quivr,Thu May 18 2023 03:04:23 GMT+0200 (heure d’été d’Europe centrale),1380
+StanGirard/quivr,Thu May 18 2023 23:04:11 GMT+0200 (heure d’été d’Europe centrale),2070
+StanGirard/quivr,Sat May 20 2023 04:44:40 GMT+0200 (heure d’été d’Europe centrale),2790
+StanGirard/quivr,Sun May 21 2023 03:19:46 GMT+0200 (heure d’été d’Europe centrale),3510
+StanGirard/quivr,Mon May 22 2023 08:03:18 GMT+0200 (heure d’été d’Europe centrale),4230
+StanGirard/quivr,Tue May 23 2023 16:57:58 GMT+0200 (heure d’été d’Europe centrale),4950
+StanGirard/quivr,Sat May 27 2023 02:18:31 GMT+0200 (heure d’été d’Europe centrale),5640
+StanGirard/quivr,Thu Jun 01 2023 18:45:27 GMT+0200 (heure d’été d’Europe centrale),6360
+StanGirard/quivr,Thu Jun 08 2023 16:33:57 GMT+0200 (heure d’été d’Europe centrale),7080
+StanGirard/quivr,Mon Jun 19 2023 12:58:34 GMT+0200 (heure d’été d’Europe centrale),7800
+StanGirard/quivr,Tue Jun 27 2023 14:45:52 GMT+0200 (heure d’été d’Europe centrale),8520
+StanGirard/quivr,Fri Jun 30 2023 11:43:51 GMT+0200 (heure d’été d’Europe centrale),9210
+StanGirard/quivr,Fri Jul 07 2023 23:08:23 GMT+0200 (heure d’été d’Europe centrale),9930
+StanGirard/quivr,Mon Jul 10 2023 08:13:07 GMT+0200 (heure d’été d’Europe centrale),10650
+StanGirard/quivr,Wed Jul 12 2023 09:40:29 GMT+0200 (heure d’été d’Europe centrale),13837

backend/test_main.py

@@ -314,7 +314,7 @@ def test_upload_and_delete_file():
     default_brain_id = brain_response.json()["brain_id"]
 
     # File to upload
-    file_path = "test.txt"
+    file_path = "test_file/test.txt"
     file_name = "test.txt"  # Assuming the name of the file on the server is the same as the local file name
 
     # Set enable_summarization flag
@@ -359,7 +359,7 @@ def test_upload_explore_and_delete_file():
     default_brain_id = brain_response.json()["brain_id"]
 
     # File to upload
-    file_path = "test.txt"
+    file_path = "test_file/test.txt"
     file_name = "test.txt"  # Assuming the name of the file on the server is the same as the local file name
 
     # Set enable_summarization flag
@@ -404,6 +404,119 @@ def test_upload_explore_and_delete_file():
     assert "message" in delete_response_data
 
 
+def test_upload_explore_and_delete_file_pdf():
+    # Retrieve the default brain
+    brain_response = client.get(
+        "/brains/default", headers={"Authorization": "Bearer " + API_KEY}
+    )
+    assert brain_response.status_code == 200
+    default_brain_id = brain_response.json()["brain_id"]
+
+    # File to upload
+    file_path = "test_file/test.pdf"
+    file_name = "test.pdf"  # Assuming the name of the file on the server is the same as the local file name
+
+    # Set enable_summarization flag
+    enable_summarization = False
+
+    # Upload the file
+    with open(file_path, "rb") as file:
+        upload_response = client.post(
+            f"/upload?brain_id={default_brain_id}&enable_summarization={enable_summarization}",
+            headers={"Authorization": "Bearer " + API_KEY},
+            files={"uploadFile": file},
+        )
+
+    # Assert that the upload response status code is 200 (HTTP OK)
+    assert upload_response.status_code == 200
+    # assert it starts with File uploaded successfully:
+    assert upload_response.json()["message"].startswith("File uploaded successfully:")
+    # show message if it fails
+    if not upload_response.json()["message"].startswith("File uploaded successfully:"):
+        print(upload_response.json()["message"])
+
+    # Optionally, you can assert on specific fields in the upload response data
+    upload_response_data = upload_response.json()
+    assert "message" in upload_response_data
+
+    # Explore (Download) the file
+    explore_response = client.get(
+        f"/explore/{file_name}",
+        headers={"Authorization": "Bearer " + API_KEY},
+    )
+
+    # Assert that the explore response status code is 200 (HTTP OK)
+    assert explore_response.status_code == 200
+
+    # Delete the file
+    delete_response = client.delete(
+        f"/explore/{file_name}",
+        headers={"Authorization": "Bearer " + API_KEY},
+        params={"brain_id": default_brain_id},
+    )
+
+    # Assert that the delete response status code is 200 (HTTP OK)
+    assert delete_response.status_code == 200
+
+    # Optionally, you can assert on specific fields in the delete response data
+    delete_response_data = delete_response.json()
+    assert "message" in delete_response_data
+
+
+def test_upload_explore_and_delete_file_csv():
+    # Retrieve the default brain
+    brain_response = client.get(
+        "/brains/default", headers={"Authorization": "Bearer " + API_KEY}
+    )
+    assert brain_response.status_code == 200
+    default_brain_id = brain_response.json()["brain_id"]
+
+    # File to upload
+    file_path = "test_file/test.csv"
+    file_name = "test.csv"  # Assuming the name of the file on the server is the same as the local file name
+
+    # Set enable_summarization flag
+    enable_summarization = False
+
+    # Upload the file
+    with open(file_path, "rb") as file:
+        upload_response = client.post(
+            f"/upload?brain_id={default_brain_id}&enable_summarization={enable_summarization}",
+            headers={"Authorization": "Bearer " + API_KEY},
+            files={"uploadFile": file},
+        )
+
+    # Assert that the upload response status code is 200 (HTTP OK)
+    assert upload_response.status_code == 200
+
+    # Optionally, you can assert on specific fields in the upload response data
+    upload_response_data = upload_response.json()
+    assert "message" in upload_response_data
+
+    # Explore (Download) the file
+    explore_response = client.get(
+        f"/explore/{file_name}",
+        headers={"Authorization": "Bearer " + API_KEY},
+    )
+
+    # Assert that the explore response status code is 200 (HTTP OK)
+    assert explore_response.status_code == 200
+
+    # Delete the file
+    delete_response = client.delete(
+        f"/explore/{file_name}",
+        headers={"Authorization": "Bearer " + API_KEY},
+        params={"brain_id": default_brain_id},
+    )
+
+    # Assert that the delete response status code is 200 (HTTP OK)
+    assert delete_response.status_code == 200
+
+    # Optionally, you can assert on specific fields in the delete response data
+    delete_response_data = delete_response.json()
+    assert "message" in delete_response_data
+
+
 def test_get_user_info():
     # Send a request to get user information
     response = client.get("/user", headers={"Authorization": "Bearer " + API_KEY})