fix: bug if no brain (#635)

.github/workflows/backend-tests.yml

@@ -42,7 +42,7 @@ jobs:
           JWT_SECRET_KEY: ${{secrets.JWT_SECRET_KEY}}
           CI_TEST_API_KEY: ${{secrets.CI_TEST_API_KEY}} 
         run: |
-          pytest
+          python -m pytest tests/
         
       - name: Static type checking with pyright
         run: |

backend/main.py

@@ -18,24 +18,21 @@ from routes.user_routes import user_router
 
 logger = get_logger(__name__)
 
-if os.getenv("SENTRY_DSN"):
+sentry_dsn = os.getenv("SENTRY_DSN")
+if sentry_dsn:
     sentry_sdk.init(
-        dsn=os.getenv("SENTRY_DSN"),
-        # Set traces_sample_rate to 1.0 to capture 100%
-        # of transactions for performance monitoring.
-        # We recommend adjusting this value in production,
+        dsn=sentry_dsn,
         traces_sample_rate=1.0,
     )
 
 app = FastAPI()
 
 add_cors_middleware(app)
-max_brain_size = os.getenv("MAX_BRAIN_SIZE", 52428800)
-max_brain_size_with_own_key = os.getenv("MAX_BRAIN_SIZE_WITH_KEY", 209715200)
 
 
 @app.on_event("startup")
 async def startup_event():
+    if not os.path.exists(pypandoc.get_pandoc_path()):
         pypandoc.download_pandoc()
 
 
@@ -49,6 +46,7 @@ app.include_router(user_router)
 app.include_router(api_key_router)
 app.include_router(subscription_router)
 
+
 @app.exception_handler(HTTPException)
 async def http_exception_handler(_, exc):
     return JSONResponse(

backend/models/brains.py

@@ -3,11 +3,10 @@ from typing import Any, List, Optional
 from uuid import UUID
 
 from logger import get_logger
-from pydantic import BaseModel
-from utils.vectors import get_unique_files_from_vector_ids
-
 from models.settings import CommonsDep, common_dependencies
 from models.users import User
+from pydantic import BaseModel
+from utils.vectors import get_unique_files_from_vector_ids
 
 logger = get_logger(__name__)
 
@@ -153,7 +152,7 @@ class Brain(BaseModel):
         self.id = response.data[0]["brain_id"]
         return response.data
 
-    def create_brain_user(self, user_id: UUID, rights, default_brain):
+    def create_brain_user(self, user_id: UUID, rights, default_brain: bool):
         commons = common_dependencies()
         response = (
             commons["supabase"]
@@ -279,7 +278,6 @@ def get_default_user_brain(user: User):
         .execute()
     )
 
-    logger.info("Default brain response:", response.data)
     default_brain_id = response.data[0]["brain_id"] if response.data else None
 
     logger.info(f"Default brain id: {default_brain_id}")
@@ -295,4 +293,14 @@ def get_default_user_brain(user: User):
 
         return brain_response.data[0] if brain_response.data else None
 
-    return None
+
+def get_default_user_brain_or_create_new(user: User) -> Brain:
+    default_brain = get_default_user_brain(user)
+
+    if default_brain:
+        return Brain.create(**default_brain)
+    else:
+        brain = Brain.create()
+        brain.create_brain()
+        brain.create_brain_user(user.id, "Owner", True)
+        return brain

backend/routes/brain_routes.py

@@ -1,33 +1,22 @@
-from typing import Optional
 from uuid import UUID
 
 from auth import AuthBearer, get_current_user
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, HTTPException
 from logger import get_logger
-from models.brains import Brain, get_default_user_brain
+from models.brains import (
+    Brain,
+    get_default_user_brain,
+    get_default_user_brain_or_create_new,
+)
 from models.settings import common_dependencies
 from models.users import User
-from pydantic import BaseModel
-
-from routes.authorizations.brain_authorization import (
-    has_brain_authorization,
-)
+from routes.authorizations.brain_authorization import has_brain_authorization
 
 logger = get_logger(__name__)
 
 brain_router = APIRouter()
 
 
-class BrainToUpdate(BaseModel):
-    brain_id: UUID
-    name: Optional[str] = "New Brain"
-    status: Optional[str] = "public"
-    model: Optional[str] = "gpt-3.5-turbo-0613"
-    temperature: Optional[float] = 0.0
-    max_tokens: Optional[int] = 256
-    file_sha1: Optional[str] = ""
-
-
 # get all brains
 @brain_router.get("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"])
 async def brain_endpoint(current_user: User = Depends(get_current_user)):
@@ -45,6 +34,7 @@ async def brain_endpoint(current_user: User = Depends(get_current_user)):
     return {"brains": brains}
 
 
+# get default brain
 @brain_router.get(
     "/brains/default/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
 )
@@ -59,23 +49,14 @@ async def get_default_brain_endpoint(current_user: User = Depends(get_current_us
     The default brain is defined as the brain marked as default in the brains_users table.
     """
 
-    default_brain = get_default_user_brain(current_user)
-
-    if default_brain is None:
-        logger.info(f"No default brain found for user {current_user.id}. Creating one.")
-
-        brain = Brain(name="Default brain")
-        brain.create_brain()
-        brain.create_brain_user(
-            user_id=current_user.id, rights="Owner", default_brain=True
-        )
-
-        default_brain = get_default_user_brain(current_user)
-
-    return default_brain
+    brain = get_default_user_brain_or_create_new(current_user)
+    return {
+        "id": brain.id,
+        "name": brain.name,
+    }
 
 
-# get one brain
+# get one brain - Currently not used in FE
 @brain_router.get(
     "/brains/{brain_id}/",
     dependencies=[Depends(AuthBearer()), Depends(has_brain_authorization())],
@@ -97,12 +78,15 @@ async def get_brain_endpoint(
     brains = brain.get_brain_details()
     if len(brains) > 0:
         return {
-            "brainId": brain_id,
-            "brainName": brains[0]["name"],
+            "id": brain_id,
+            "name": brains[0]["name"],
             "status": brains[0]["status"],
         }
     else:
-        return {"error": f"No brain found with brain_id {brain_id}"}
+        return HTTPException(
+            status_code=404,
+            detail="Brain not found",
+        )
 
 
 # delete one brain
@@ -124,20 +108,10 @@ async def delete_brain_endpoint(
     return {"message": f"{brain_id}  has been deleted."}
 
 
-class BrainObject(BaseModel):
-    brain_id: Optional[UUID]
-    name: Optional[str] = "New Brain"
-    status: Optional[str] = "public"
-    model: Optional[str] = "gpt-3.5-turbo-0613"
-    temperature: Optional[float] = 0.0
-    max_tokens: Optional[int] = 256
-    file_sha1: Optional[str] = ""
-
-
 # create new brain
 @brain_router.post("/brains/", dependencies=[Depends(AuthBearer())], tags=["Brain"])
 async def create_brain_endpoint(
-    brain: BrainObject,
+    brain: Brain,
     current_user: User = Depends(get_current_user),
 ):
     """
@@ -205,7 +179,6 @@ async def update_brain_endpoint(
         brain.update_brain_with_file(
             file_sha1=input_brain.file_sha1  # pyright: ignore reportPrivateUsage=none
         )
-        print("brain:", brain)
 
     brain.update_brain_fields(commons, brain)  # pyright: ignore reportPrivateUsage=none
     return {"message": f"Brain {brain_id} has been updated."}

backend/routes/chat_routes.py

@@ -10,6 +10,7 @@ from fastapi.responses import StreamingResponse
 from llm.openai import OpenAIBrainPicking
 from llm.openai_functions import OpenAIFunctionsBrainPicking
 from llm.private_gpt4all import PrivateGPT4AllBrainPicking
+from models.brains import get_default_user_brain_or_create_new
 from models.chat import Chat, ChatHistory
 from models.chats import ChatQuestion
 from models.settings import LLMSettings, common_dependencies
@@ -27,6 +28,21 @@ from utils.constants import (
 chat_router = APIRouter()
 
 
+class NullableUUID:
+    @classmethod
+    def __get_validators__(cls):
+        yield cls.validate
+
+    @classmethod
+    def validate(cls, v):
+        if v == "":
+            return None
+        try:
+            return UUID(v)
+        except ValueError:
+            return None
+
+
 def get_chat_details(commons, chat_id):
     response = (
         commons["supabase"]
@@ -151,22 +167,30 @@ async def create_chat_handler(
 
 # add new question to chat
 @chat_router.post(
-    "/chat/{chat_id}/question", dependencies=[Depends(AuthBearer())], tags=["Chat"]
+    "/chat/{chat_id}/question",
+    dependencies=[
+        Depends(
+            AuthBearer(),
+        ),
+    ],
+    tags=["Chat"],
 )
 async def create_question_handler(
     request: Request,
     chat_question: ChatQuestion,
     chat_id: UUID,
-    brain_id: UUID = Query(..., description="The ID of the brain"),
+    brain_id: NullableUUID
+    | UUID
+    | None = Query(..., description="The ID of the brain"),
     current_user: User = Depends(get_current_user),
 ) -> ChatHistory:
     current_user.user_openai_api_key = request.headers.get("Openai-Api-Key")
-    print("current_user", current_user)
     try:
         check_user_limit(current_user)
         llm_settings = LLMSettings()
 
-        # TODO: RBAC with current_user
+        if not brain_id:
+            brain_id = get_default_user_brain_or_create_new(current_user).id
 
         if llm_settings.private:
             gpt_answer_generator = PrivateGPT4AllBrainPicking(
@@ -209,16 +233,26 @@ async def create_question_handler(
 # stream new question response from chat
 @chat_router.post(
     "/chat/{chat_id}/question/stream",
-    dependencies=[Depends(AuthBearer())],
+    dependencies=[
+        Depends(
+            AuthBearer(),
+        ),
+    ],
     tags=["Chat"],
 )
 async def create_stream_question_handler(
     request: Request,
     chat_question: ChatQuestion,
     chat_id: UUID,
-    brain_id: UUID = Query(..., description="The ID of the brain"),
+    brain_id: NullableUUID
+    | UUID
+    | None = Query(..., description="The ID of the brain"),
     current_user: User = Depends(get_current_user),
 ) -> StreamingResponse:
+    # TODO: check if the user has access to the brain
+    if not brain_id:
+        brain_id = get_default_user_brain_or_create_new(current_user).id
+
     if chat_question.model not in streaming_compatible_models:
         # Forward the request to the none streaming endpoint
         return await create_question_handler(

backend/routes/upload_routes.py

@@ -34,11 +34,8 @@ async def upload_file(
     it can optionally apply summarization to the file's content. The response message will indicate the status of the upload.
     """
 
-    print(brain_id, "brain_id")
-
     # [TODO] check if the user is the owner/editor of the brain
     brain = Brain(id=brain_id)
-    print("brain", brain)
     commons = common_dependencies()
 
     if request.headers.get("Openai-Api-Key"):

backend/test_main.py

@@ -1,530 +0,0 @@
-import os
-import random
-import string
-import uuid
-
-from fastapi.testclient import TestClient
-from main import app
-
-client = TestClient(app)
-
-API_KEY = os.getenv("CI_TEST_API_KEY")
-
-if not API_KEY:
-    raise ValueError("CI_TEST_API_KEY environment variable not set. Cannot run tests.")
-
-
-def test_read_main():
-    response = client.get("/")
-    assert response.status_code == 200
-    assert response.json() == {"status": "OK"}
-
-
-def test_create_and_delete_api_key():
-    # First, let's create an API key
-    response = client.post(
-        "/api-key",
-        headers={
-            "Authorization": "Bearer " + API_KEY,
-        },
-    )
-    assert response.status_code == 200
-    api_key_info = response.json()
-    assert "api_key" in api_key_info
-
-    # Extract the created api_key from the response
-    api_key = api_key_info["api_key"]
-
-    # Now, let's verify the API key
-    verify_response = client.get(
-        "/user",
-        headers={
-            "Authorization": f"Bearer {api_key}",
-        },
-    )
-    assert verify_response.status_code == 200
-
-    # Now, let's delete the API key
-    assert "key_id" in api_key_info
-    key_id = api_key_info["key_id"]
-
-    delete_response = client.delete(
-        f"/api-key/{key_id}", headers={"Authorization": f"Bearer {API_KEY}"}
-    )
-    assert delete_response.status_code == 200
-    assert delete_response.json() == {"message": "API key deleted."}
-
-
-def test_retrieve_default_brain():
-    # Making a GET request to the /brains/default/ endpoint
-    response = client.get(
-        "/brains/default/",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-
-    # Optionally, assert on specific fields in the response
-    response_data = response.json()
-    # e.g., assert that the response contains a 'brain_id' field
-    assert "brain_id" in response_data
-
-
-def test_create_brain():
-
-    # Generate a random name for the brain
-    random_brain_name = "".join(
-        random.choices(string.ascii_letters + string.digits, k=10)
-    )
-
-    # Set up the request payload
-    payload = {
-        "name": random_brain_name,
-        "status": "public",
-        "model": "gpt-3.5-turbo-0613",
-        "temperature": 0,
-        "max_tokens": 256,
-        "file_sha1": "",
-    }
-
-    # Making a POST request to the /brains/ endpoint
-    response = client.post(
-        "/brains/",
-        json=payload,
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-
-    # Optionally, assert on specific fields in the response
-    response_data = response.json()
-    # e.g., assert that the response contains a 'brain_id' field
-    assert "id" in response_data
-    assert "name" in response_data
-
-    # Optionally, assert that the returned 'name' matches the one sent in the request
-    assert response_data["name"] == payload["name"]
-
-
-def test_retrieve_all_brains():
-    # Making a GET request to the /brains/ endpoint to retrieve all brains for the current user
-    response = client.get(
-        "/brains/",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-
-    response_data = response.json()
-    # Optionally, you can loop through the brains and assert on specific fields in each brain
-    for brain in response_data["brains"]:
-        assert "id" in brain
-        assert "name" in brain
-
-
-def test_delete_all_brains():
-    # First, retrieve all brains for the current user
-    response = client.get(
-        "/brains/",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-
-    response_data = response.json()
-
-    # Loop through each brain and send a DELETE request
-    for brain in response_data["brains"]:
-        brain_id = brain["id"]
-
-        # Send a DELETE request to delete the specific brain
-        delete_response = client.delete(
-            f"/brains/{brain_id}/",
-            headers={"Authorization": "Bearer " + API_KEY},
-        )
-
-        # Assert that the DELETE response status code is 200 (HTTP OK)
-        assert delete_response.status_code == 200
-
-
-def test_delete_all_brains_and_get_default_brain():
-    # First create a new brain
-    test_create_brain()
-
-    # Now, retrieve all brains for the current user
-    response = client.get(
-        "/brains/",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-    assert len(response.json()["brains"]) > 0
-
-    test_delete_all_brains()
-
-    # Now, retrieve all brains for the current user
-    response = client.get(
-        "/brains/",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-    assert len(response.json()["brains"]) == 0
-
-    # Get the default brain, it should create one if it doesn't exist
-    response = client.get(
-        "/brains/default/",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-    assert response.json()["name"] == "Default brain"
-
-    # Now, retrieve all brains for the current user
-    response = client.get(
-        "/brains/",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that there is only one brain
-    response_data = response.json()
-    assert len(response_data) == 1
-    for brain in response_data["brains"]:
-        assert "id" in brain
-        assert "name" in brain
-
-    # Assert that the brain is the default brain
-    assert response_data["brains"][0]["name"] == "Default brain"
-
-
-def test_get_all_chats():
-    # Making a GET request to the /chat endpoint to retrieve all chats
-    response = client.get(
-        "/chat",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-
-    # Assert that the response data is a list
-    response_data = response.json()
-
-    # Optionally, you can loop through the chats and assert on specific fields
-    for chat in response_data["chats"]:
-        # e.g., assert that each chat object contains 'chat_id' and 'chat_name'
-        assert "chat_id" in chat
-        assert "chat_name" in chat
-
-
-def test_create_chat_and_talk():
-    # Make a POST request to chat with the default brain and a random chat name
-    random_chat_name = "".join(
-        random.choices(string.ascii_letters + string.digits, k=10)
-    )
-
-    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"]
-    print("Default brain id: " + default_brain_id)
-
-    # Create a chat
-    response = client.post(
-        "/chat",
-        json={"name": random_chat_name},
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-    assert response.status_code == 200
-
-    # now talk to the chat with a question
-    response_data = response.json()
-    print(response_data)
-    chat_id = response_data["chat_id"]
-    response = client.post(
-        f"/chat/{chat_id}/question?brain_id={default_brain_id}",
-        json={
-            "model": "gpt-3.5-turbo-0613",
-            "question": "Hello, how are you?",
-            "temperature": "0",
-            "max_tokens": "256",
-        },
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-    assert response.status_code == 200
-
-    response = client.post(
-        f"/chat/{chat_id}/question?brain_id={default_brain_id}",
-        json={
-            "model": "gpt-4",
-            "question": "Hello, how are you?",
-            "temperature": "0",
-            "max_tokens": "256",
-        },
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-    print(response)
-    assert response.status_code == 200
-
-    # Now, let's delete the chat
-    # Assuming the chat_id is part of the chat_info response. If not, adjust this.
-    delete_response = client.delete(
-        "/chat/" + chat_id, headers={"Authorization": "Bearer " + API_KEY}
-    )
-    assert delete_response.status_code == 200
-
-
-def test_explore_with_default_brain():
-    # 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"]
-
-    # Now use the default brain_id as parameter in the /explore/ endpoint
-    response = client.get(
-        f"/explore/{default_brain_id}",
-        headers={"Authorization": "Bearer " + API_KEY},
-    )
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-
-    # Optionally, you can assert on specific fields in the response data
-    response_data = response.json()
-    # e.g., assert that the response contains a 'results' field
-    assert "documents" in response_data
-
-
-def test_upload_and_delete_file():
-    # 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.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
-    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
-
-    # 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_txt():
-    # 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.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
-    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_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:
-
-    # Optionally, you can assert on specific fields in the upload response data
-    upload_response_data = upload_response.json()
-    assert "message" in upload_response_data
-    assert "type" in upload_response_data
-    assert upload_response_data["type"] == "success"
-
-    # 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})
-
-    # Assert that the response status code is 200 (HTTP OK)
-    assert response.status_code == 200
-
-    # Assert that the response contains the expected fields
-    user_info = response.json()
-    assert "email" in user_info
-    assert "max_brain_size" in user_info
-    assert "current_brain_size" in user_info
-    assert "date" in user_info

backend/tests/conftest.py

@@ -0,0 +1,20 @@
+import os
+
+import pytest
+from fastapi.testclient import TestClient
+from main import app
+
+
+@pytest.fixture(scope="module")
+def client():
+    return TestClient(app)
+
+
+@pytest.fixture(scope="module")
+def api_key():
+    API_KEY = os.getenv("CI_TEST_API_KEY")
+    if not API_KEY:
+        raise ValueError(
+            "CI_TEST_API_KEY environment variable not set. Cannot run tests."
+        )
+    return API_KEY

backend/tests/test_api_key.py

@@ -0,0 +1,39 @@
+def test_read_main(client, api_key):
+    response = client.get("/")
+    assert response.status_code == 200
+    assert response.json() == {"status": "OK"}
+
+
+def test_create_and_delete_api_key(client, api_key):
+    # First, let's create an API key
+    response = client.post(
+        "/api-key",
+        headers={
+            "Authorization": "Bearer " + api_key,
+        },
+    )
+    assert response.status_code == 200
+    api_key_info = response.json()
+    assert "api_key" in api_key_info
+
+    # Extract the created api_key from the response
+    api_key = api_key_info["api_key"]
+
+    # Now, let's verify the API key
+    verify_response = client.get(
+        "/user",
+        headers={
+            "Authorization": f"Bearer {api_key}",
+        },
+    )
+    assert verify_response.status_code == 200
+
+    # Now, let's delete the API key
+    assert "key_id" in api_key_info
+    key_id = api_key_info["key_id"]
+
+    delete_response = client.delete(
+        f"/api-key/{key_id}", headers={"Authorization": f"Bearer {api_key}"}
+    )
+    assert delete_response.status_code == 200
+    assert delete_response.json() == {"message": "API key deleted."}

backend/tests/test_brains.py

@@ -0,0 +1,164 @@
+import random
+import string
+
+
+def test_retrieve_default_brain(client, api_key):
+    # Making a GET request to the /brains/default/ endpoint
+    response = client.get(
+        "/brains/default/",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    default_brain = response.json()
+    assert "id" in default_brain
+    assert "name" in default_brain
+
+
+def test_create_brain(client, api_key):
+    # Generate a random name for the brain
+    random_brain_name = "".join(
+        random.choices(string.ascii_letters + string.digits, k=10)
+    )
+
+    # Set up the request payload
+    payload = {
+        "name": random_brain_name,
+        "status": "public",
+        "model": "gpt-3.5-turbo-0613",
+        "temperature": 0,
+        "max_tokens": 256,
+        "file_sha1": "",
+    }
+
+    # Making a POST request to the /brains/ endpoint
+    response = client.post(
+        "/brains/",
+        json=payload,
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    # Optionally, assert on specific fields in the response
+    response_data = response.json()
+    # e.g., assert that the response contains a 'brain_id' field
+    assert "id" in response_data
+    assert "name" in response_data
+
+    # Optionally, assert that the returned 'name' matches the one sent in the request
+    assert response_data["name"] == payload["name"]
+
+
+def test_retrieve_all_brains(client, api_key):
+    # Making a GET request to the /brains/ endpoint to retrieve all brains for the current user
+    response = client.get(
+        "/brains/",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    response_data = response.json()
+    # Optionally, you can loop through the brains and assert on specific fields in each brain
+    for brain in response_data["brains"]:
+        assert "id" in brain
+        assert "name" in brain
+
+
+def test_retrieve_one_brain(client, api_key):
+    # Making a GET request to the /brains/default/ endpoint
+    response = client.get(
+        "/brains/default/",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    response_data = response.json()
+
+    # Extract the brain_id from the response
+    brain_id = response_data["id"]
+
+    # Making a GET request to the /brains/{brain_id}/ endpoint
+    response = client.get(
+        f"/brains/{brain_id}/",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    brain = response.json()
+    assert "id" in brain
+    assert "name" in brain
+    assert "status" in brain
+
+
+def test_delete_all_brains(client, api_key):
+    # First, retrieve all brains for the current user
+    response = client.get(
+        "/brains/",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    response_data = response.json()
+
+    # Loop through each brain and send a DELETE request
+    for brain in response_data["brains"]:
+        brain_id = brain["id"]
+
+        # Send a DELETE request to delete the specific brain
+        delete_response = client.delete(
+            f"/brains/{brain_id}/",
+            headers={"Authorization": "Bearer " + api_key},
+        )
+
+        # Assert that the DELETE response status code is 200 (HTTP OK)
+        assert delete_response.status_code == 200
+
+    # Finally, retrieve all brains for the current user
+    response = client.get(
+        "/brains/",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+    response_data = response.json()
+    assert len(response_data["brains"]) == 0
+
+
+def test_delete_all_brains_and_get_default_brain(client, api_key):
+    # First create a new brain
+    test_create_brain(client, api_key)
+
+    # Now, retrieve all brains for the current user
+    response = client.get(
+        "/brains/",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+    assert len(response.json()["brains"]) > 0
+
+    test_delete_all_brains(client, api_key)
+
+    # Get the default brain, it should create one if it doesn't exist
+    response = client.get(
+        "/brains/default/",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+    assert response.json()["name"] == "Default brain"

backend/tests/test_chats.py

@@ -0,0 +1,116 @@
+import random
+import string
+
+
+def test_get_all_chats(client, api_key):
+    # Making a GET request to the /chat endpoint to retrieve all chats
+    response = client.get(
+        "/chat",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    # Assert that the response data is a list
+    response_data = response.json()
+
+    # Optionally, you can loop through the chats and assert on specific fields
+    for chat in response_data["chats"]:
+        # e.g., assert that each chat object contains 'chat_id' and 'chat_name'
+        assert "chat_id" in chat
+        assert "chat_name" in chat
+
+
+def test_create_chat_and_talk(client, api_key):
+    # Make a POST request to chat with the default brain and a random chat name
+    random_chat_name = "".join(
+        random.choices(string.ascii_letters + string.digits, k=10)
+    )
+
+    brain_response = client.get(
+        "/brains/default", headers={"Authorization": "Bearer " + api_key}
+    )
+    assert brain_response.status_code == 200
+    default_brain_id = brain_response.json()["id"]
+    print("Default brain id: " + default_brain_id)
+
+    # Create a chat
+    response = client.post(
+        "/chat",
+        json={"name": random_chat_name},
+        headers={"Authorization": "Bearer " + api_key},
+    )
+    assert response.status_code == 200
+
+    # now talk to the chat with a question
+    response_data = response.json()
+    print(response_data)
+    chat_id = response_data["chat_id"]
+    response = client.post(
+        f"/chat/{chat_id}/question?brain_id={default_brain_id}",
+        json={
+            "model": "gpt-3.5-turbo-0613",
+            "question": "Hello, how are you?",
+            "temperature": "0",
+            "max_tokens": "256",
+        },
+        headers={"Authorization": "Bearer " + api_key},
+    )
+    assert response.status_code == 200
+
+    response = client.post(
+        f"/chat/{chat_id}/question?brain_id={default_brain_id}",
+        json={
+            "model": "gpt-4",
+            "question": "Hello, how are you?",
+            "temperature": "0",
+            "max_tokens": "256",
+        },
+        headers={"Authorization": "Bearer " + api_key},
+    )
+    print(response)
+    assert response.status_code == 200
+
+    # Now, let's delete the chat
+    delete_response = client.delete(
+        "/chat/" + chat_id, headers={"Authorization": "Bearer " + api_key}
+    )
+    assert delete_response.status_code == 200
+
+
+def test_create_chat_and_talk_with_no_brain(client, api_key):
+    # Make a POST request to chat with no brain id and a random chat name
+    random_chat_name = "".join(
+        random.choices(string.ascii_letters + string.digits, k=10)
+    )
+
+    # Create a chat
+    response = client.post(
+        "/chat",
+        json={"name": random_chat_name},
+        headers={"Authorization": "Bearer " + api_key},
+    )
+    assert response.status_code == 200
+
+    # now talk to the chat with a question
+    response_data = response.json()
+    print(response_data)
+    chat_id = response_data["chat_id"]
+    response = client.post(
+        f"/chat/{chat_id}/question?brain_id=",
+        json={
+            "model": "gpt-3.5-turbo-0613",
+            "question": "Hello, how are you?",
+            "temperature": "0",
+            "max_tokens": "256",
+        },
+        headers={"Authorization": "Bearer " + api_key},
+    )
+    assert response.status_code == 200
+
+    # Now, let's delete the chat
+    delete_response = client.delete(
+        "/chat/" + chat_id, headers={"Authorization": "Bearer " + api_key}
+    )
+    assert delete_response.status_code == 200

backend/tests/test_explore.py

@@ -0,0 +1,21 @@
+def test_explore_with_default_brain(client, api_key):
+    # 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()["id"]
+
+    # Now use the default brain_id as parameter in the /explore/ endpoint
+    response = client.get(
+        f"/explore/{default_brain_id}",
+        headers={"Authorization": "Bearer " + api_key},
+    )
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    # Optionally, you can assert on specific fields in the response data
+    response_data = response.json()
+    # e.g., assert that the response contains a 'results' field
+    assert "documents" in response_data

backend/tests/test_upload.py

@@ -0,0 +1,208 @@
+def test_upload_and_delete_file(client, api_key):
+    # 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()["id"]
+
+    # File to upload
+    file_path = "tests/test_files/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
+    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
+
+    # 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_txt(client, api_key):
+    # 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()["id"]
+
+    # File to upload
+    file_path = "tests/test_files/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
+    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_upload_explore_and_delete_file_pdf(client, api_key):
+    # 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()["id"]
+
+    # File to upload
+    file_path = "tests/test_files/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:
+
+    # Optionally, you can assert on specific fields in the upload response data
+    upload_response_data = upload_response.json()
+    assert "message" in upload_response_data
+    assert "type" in upload_response_data
+    assert upload_response_data["type"] == "success"
+
+    # 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(client, api_key):
+    # 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()["id"]
+
+    # File to upload
+    file_path = "tests/test_files/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

backend/tests/test_user.py

@@ -0,0 +1,13 @@
+def test_get_user_info(client, api_key):
+    # Send a request to get user information
+    response = client.get("/user", headers={"Authorization": "Bearer " + api_key})
+
+    # Assert that the response status code is 200 (HTTP OK)
+    assert response.status_code == 200
+
+    # Assert that the response contains the expected fields
+    user_info = response.json()
+    assert "email" in user_info
+    assert "max_brain_size" in user_info
+    assert "current_brain_size" in user_info
+    assert "date" in user_info

frontend/app/chat/[chatId]/hooks/useQuestion.ts

@@ -25,13 +25,9 @@ export const useQuestion = (): UseChatService => {
     chatId: string,
     chatQuestion: ChatQuestion
   ): Promise<void> => {
-    if (currentBrain?.id === undefined) {
-      throw new Error("No current brain");
-    }
-
     const response = await addQuestion({
       chatId,
-      brainId: currentBrain.id,
+      brainId: currentBrain?.id ?? "",
       chatQuestion,
     });