Feat/testing backend (#446)

* feat(pytest): added

* feat(brains): added tests

* feat(actions): pytest

.github/workflows/pytest.yml

@@ -0,0 +1,40 @@
+name: Pytest
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]    
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    defaults:
+      run:
+        working-directory: ./backend
+    environment: preview
+    strategy:
+      matrix:
+        python-version: [ "3.11"]
+
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: 'pip'
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install pytest
+          if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
+      - name: Test with pytest
+        env:
+          SUPABASE_URL: ${{vars.SUPABASE_URL}}
+          SUPABASE_SERVICE_KEY: ${{vars.SUPABASE_SERVICE_KEY}}
+          OPENAI_API_KEY: ${{vars.OPENAI_API_KEY}}
+          ANTHROPIC_API_KEY: ${{vars.ANTHROPIC_API_KEY}}
+          JWT_SECRET_KEY: ${{vars.JWT_SECRET_KEY}}
+          CI_TEST_API_KEY: ${{vars.CI_TEST_API_KEY}} 
+        run: |
+          pytest

.github/workflows/vitest.yml

@@ -12,14 +12,18 @@ jobs:
     defaults:
         run:
             working-directory: ./frontend
+    strategy:
+      matrix:
+        node-version: [18]
 
     steps:
       - uses: actions/checkout@v3
       - name: Use Node.js ${{ matrix.node-version }}
         uses: actions/setup-node@v3
         with:
-          node-version: 18
+          node-version: ${{ matrix.node-version }}
           cache: 'yarn'
           cache-dependency-path: frontend/yarn.lock
       - run: yarn
-      - run: yarn run test-unit
+      - run: yarn run test-unit
+      - run: yarn run build

backend/auth/api_key_handler.py

@@ -25,8 +25,10 @@ async def verify_api_key(
                 result.data[0]["creation_time"], "%Y-%m-%dT%H:%M:%S"
             ).date()
 
-            # Check if the API key was created today: Todo remove this check and use deleted_time instead.
-            if api_key_creation_date == current_date:
+            # Check if the API key was created in the month of the current date
+            if (api_key_creation_date.month == current_date.month) and (
+                api_key_creation_date.year == current_date.year
+            ):
                 return True
         return False
     except DateError:
@@ -62,7 +64,7 @@ async def get_user_from_api_key(
     )
 
     return (
-        {"email": user_email_data.data[0]["email"]}
+        {"email": user_email_data.data[0]["email"], "sub": user_id}
         if user_email_data.data
         else {"email": None}
     )

backend/auth/auth_bearer.py

@@ -55,4 +55,6 @@ class AuthBearer(HTTPBearer):
 
 
 def get_current_user(credentials: dict = Depends(AuthBearer())) -> User:
-    return User(email=credentials.get("email", "none"), id=credentials.get("sub", "none"))
+    return User(
+        email=credentials.get("email", "none"), id=credentials.get("sub", "none")
+    )

backend/models/brains.py

@@ -10,7 +10,7 @@ from pydantic import BaseModel
 class Brain(BaseModel):
     id: Optional[UUID] = None
     name: Optional[str] = "Default brain"
-    status: Optional[str]= "public"
+    status: Optional[str] = "public"
     model: Optional[str] = "gpt-3.5-turbo-0613"
     temperature: Optional[float] = 0.0
     max_tokens: Optional[int] = 256
@@ -31,9 +31,9 @@ class Brain(BaseModel):
     @property
     def brain_size(self):
         self.get_unique_brain_files()
-        current_brain_size = sum(float(doc['size']) for doc in self.files)
+        current_brain_size = sum(float(doc["size"]) for doc in self.files)
 
-        print('current_brain_size', current_brain_size)
+        print("current_brain_size", current_brain_size)
         return current_brain_size
 
     @property
@@ -65,30 +65,86 @@ class Brain(BaseModel):
         )
         return response.data
 
-    def delete_brain(self):
-        self.commons["supabase"].table("brains").delete().match(
-            {"brain_id": self.id}
-        ).execute()
+    def delete_brain(self, user_id):
+        print("user_id", user_id)
+        print("self.id", self.id)
+        results = (
+            self.commons["supabase"]
+            .table("brains_users")
+            .select("*")
+            .match({"brain_id": self.id, "user_id": user_id, "rights": "Owner"})
+            .execute()
+        )
+        if len(results.data) == 0:
+            print("You are not the owner of this brain.")
+            return {"message": "You are not the owner of this brain."}
+        else:
+            results = (
+                self.commons["supabase"]
+                .table("brains_vectors")
+                .delete()
+                .match({"brain_id": self.id})
+                .execute()
+            )
+            print("results", results)
+
+            results = (
+                self.commons["supabase"]
+                .table("brains_users")
+                .delete()
+                .match({"brain_id": self.id})
+                .execute()
+            )
+            print("results", results)
+
+            results = (
+                self.commons["supabase"]
+                .table("brains")
+                .delete()
+                .match({"brain_id": self.id})
+                .execute()
+            )
+            print("results", results)
 
     def create_brain(self):
         commons = common_dependencies()
-        response = commons["supabase"].table("brains").insert({"name": self.name}).execute()
+        response = (
+            commons["supabase"].table("brains").insert({"name": self.name}).execute()
+        )
         # set the brainId with response.data
 
-        self.id = response.data[0]['brain_id']   
+        self.id = response.data[0]["brain_id"]
         return response.data
 
     def create_brain_user(self, user_id: UUID, rights, default_brain):
         commons = common_dependencies()
-        response = commons["supabase"].table("brains_users").insert({"brain_id": str(self.id), "user_id": str(user_id), "rights": rights, "default_brain": default_brain}).execute()
-        
+        response = (
+            commons["supabase"]
+            .table("brains_users")
+            .insert(
+                {
+                    "brain_id": str(self.id),
+                    "user_id": str(user_id),
+                    "rights": rights,
+                    "default_brain": default_brain,
+                }
+            )
+            .execute()
+        )
+
         return response.data
 
     def create_brain_vector(self, vector_id, file_sha1):
         response = (
             self.commons["supabase"]
             .table("brains_vectors")
-            .insert({"brain_id": str(self.id), "vector_id": str(vector_id), "file_sha1": file_sha1})
+            .insert(
+                {
+                    "brain_id": str(self.id),
+                    "vector_id": str(vector_id),
+                    "file_sha1": file_sha1,
+                }
+            )
             .execute()
         )
         return response.data
@@ -121,22 +177,22 @@ class Brain(BaseModel):
         """
 
         response = (
-                self.commons["supabase"]
-                .from_("brains_vectors")
-                .select("vector_id")
-                .filter("brain_id", "eq", self.id)
-                .execute()
-            )
-        
+            self.commons["supabase"]
+            .from_("brains_vectors")
+            .select("vector_id")
+            .filter("brain_id", "eq", self.id)
+            .execute()
+        )
+
         vector_ids = [item["vector_id"] for item in response.data]
 
-        print('vector_ids', vector_ids)
+        print("vector_ids", vector_ids)
 
         if len(vector_ids) == 0:
             return []
-        
+
         self.files = self.get_unique_files_from_vector_ids(vector_ids)
-        print('unique_files', self.files)
+        print("unique_files", self.files)
 
         return self.files
 
@@ -145,19 +201,31 @@ class Brain(BaseModel):
         """
         Retrieve unique user data vectors.
         """
-        print('vectors_ids', vectors_ids)
-        print('tuple(vectors_ids)', tuple(vectors_ids))
+        print("vectors_ids", vectors_ids)
+        print("tuple(vectors_ids)", tuple(vectors_ids))
         if len(vectors_ids) == 1:
-            vectors_response = self.commons['supabase'].table("vectors").select(
-                "name:metadata->>file_name, size:metadata->>file_size", count="exact") \
-                .filter("id", "eq", vectors_ids[0])\
+            vectors_response = (
+                self.commons["supabase"]
+                .table("vectors")
+                .select(
+                    "name:metadata->>file_name, size:metadata->>file_size",
+                    count="exact",
+                )
+                .filter("id", "eq", vectors_ids[0])
                 .execute()
+            )
         else:
-            vectors_response = self.commons['supabase'].table("vectors").select(
-                "name:metadata->>file_name, size:metadata->>file_size", count="exact") \
-                .filter("id", "in", tuple(vectors_ids))\
+            vectors_response = (
+                self.commons["supabase"]
+                .table("vectors")
+                .select(
+                    "name:metadata->>file_name, size:metadata->>file_size",
+                    count="exact",
+                )
+                .filter("id", "in", tuple(vectors_ids))
                 .execute()
-            
+            )
+
         documents = vectors_response.data  # Access the data from the response
         # Convert each dictionary to a tuple of items, then to a set to remove duplicates, and then back to a dictionary
         unique_files = [dict(t) for t in set(tuple(d.items()) for d in documents)]
@@ -165,32 +233,52 @@ class Brain(BaseModel):
 
     def delete_file_from_brain(self, file_name: str):
         # First, get the vector_ids associated with the file_name
-        vector_response = self.commons["supabase"].table("vectors").select("id").filter("metadata->>file_name", "eq", file_name).execute()
+        vector_response = (
+            self.commons["supabase"]
+            .table("vectors")
+            .select("id")
+            .filter("metadata->>file_name", "eq", file_name)
+            .execute()
+        )
         vector_ids = [item["id"] for item in vector_response.data]
 
         # For each vector_id, delete the corresponding entry from the 'brains_vectors' table
         for vector_id in vector_ids:
-            self.commons["supabase"].table("brains_vectors").delete().filter("vector_id", "eq", vector_id).filter("brain_id", "eq", self.id).execute()
+            self.commons["supabase"].table("brains_vectors").delete().filter(
+                "vector_id", "eq", vector_id
+            ).filter("brain_id", "eq", self.id).execute()
 
             # Check if the vector is still associated with any other brains
-            associated_brains_response = self.commons["supabase"].table("brains_vectors").select("brain_id").filter("vector_id", "eq", vector_id).execute()
-            associated_brains = [item["brain_id"] for item in associated_brains_response.data]
+            associated_brains_response = (
+                self.commons["supabase"]
+                .table("brains_vectors")
+                .select("brain_id")
+                .filter("vector_id", "eq", vector_id)
+                .execute()
+            )
+            associated_brains = [
+                item["brain_id"] for item in associated_brains_response.data
+            ]
 
             # If the vector is not associated with any other brains, delete it from 'vectors' table
             if not associated_brains:
-                self.commons["supabase"].table("vectors").delete().filter("id", "eq", vector_id).execute()
+                self.commons["supabase"].table("vectors").delete().filter(
+                    "id", "eq", vector_id
+                ).execute()
 
         return {"message": f"File {file_name} in brain {self.id} has been deleted."}
 
-                        
+
 def get_default_user_brain(user: User):
     commons = common_dependencies()
     response = (
         commons["supabase"]
-        .from_("brains_users") # I'm assuming this is the correct table
-        .select("brain_id") 
+        .from_("brains_users")  # I'm assuming this is the correct table
+        .select("brain_id")
         .filter("user_id", "eq", user.id)
-        .filter("default_brain", "eq", True) # Assuming 'default' is the correct column name
+        .filter(
+            "default_brain", "eq", True
+        )  # Assuming 'default' is the correct column name
         .execute()
     )
 
@@ -207,8 +295,7 @@ def get_default_user_brain(user: User):
             .filter("brain_id", "eq", default_brain_id)
             .execute()
         )
-        
+
         return brain_response.data[0] if brain_response.data else None
 
     return None
-

backend/routes/api_key_routes.py

@@ -21,6 +21,7 @@ class ApiKeyInfo(BaseModel):
 
 class ApiKey(BaseModel):
     api_key: str
+    key_id: str
 
 
 api_key_router = APIRouter()
@@ -76,7 +77,7 @@ async def create_api_key(
             return {"api_key": "Error creating new API key."}
     logger.info(f"Created new API key for user {current_user.email}.")
 
-    return {"api_key": new_api_key}
+    return {"api_key": new_api_key, "key_id": str(new_key_id)}
 
 
 @api_key_router.delete(

backend/routes/brain_routes.py

@@ -41,7 +41,9 @@ async def brain_endpoint(current_user: User = Depends(get_current_user)):
     return {"brains": brains}
 
 
-@brain_router.get("/brains/default/", dependencies=[Depends(AuthBearer())], tags=["Brain"])
+@brain_router.get(
+    "/brains/default/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
+)
 async def get_default_brain_endpoint(current_user: User = Depends(get_current_user)):
     """
     Retrieve the default brain for the current user.
@@ -49,7 +51,7 @@ async def get_default_brain_endpoint(current_user: User = Depends(get_current_us
     - `current_user`: The current authenticated user.
     - Returns the default brain for the user.
 
-    This endpoint retrieves the default brain associated with the current authenticated user. 
+    This endpoint retrieves the default brain associated with the current authenticated user.
     The default brain is defined as the brain marked as default in the brains_users table.
     """
 
@@ -87,15 +89,17 @@ async def get_brain_endpoint(brain_id: UUID):
 @brain_router.delete(
     "/brains/{brain_id}/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
 )
-async def delete_brain_endpoint(brain_id: UUID, current_user: User = Depends(get_current_user),):
+async def delete_brain_endpoint(
+    brain_id: UUID,
+    current_user: User = Depends(get_current_user),
+):
     """
     Delete a specific brain by brain ID.
     """
     # [TODO] check if the user is the owner of the brain
-    
-    current_user.id, 
+
     brain = Brain(id=brain_id)
-    brain.delete_brain()
+    brain.delete_brain(current_user.id)
 
     return {"message": f"{brain_id}  has been deleted."}
 
@@ -127,15 +131,21 @@ async def create_brain_endpoint(
     """
 
     brain = Brain(name=brain.name)
-    
+
     brain.create_brain()
     default_brain = get_default_user_brain(current_user)
     if default_brain:
         logger.info(f"Default brain already exists for user {current_user.id}")
-        brain.create_brain_user(user_id=current_user.id, rights="Owner", default_brain=False)
+        brain.create_brain_user(
+            user_id=current_user.id, rights="Owner", default_brain=False
+        )
     else:
-        logger.info(f"Default brain does not exist for user {current_user.id}. It will be created.")
-        brain.create_brain_user(user_id=current_user.id, rights="Owner", default_brain=True)
+        logger.info(
+            f"Default brain does not exist for user {current_user.id}. It will be created."
+        )
+        brain.create_brain_user(
+            user_id=current_user.id, rights="Owner", default_brain=True
+        )
 
     return {"id": brain.id, "name": brain.name}
 

backend/routes/chat_routes.py

@@ -76,10 +76,10 @@ def check_user_limit(
 ):
     if user.user_openai_api_key is None:
         date = time.strftime("%Y%m%d")
-        max_requests_number = os.getenv("MAX_REQUESTS_NUMBER")
+        max_requests_number = os.getenv("MAX_REQUESTS_NUMBER", 1000)
 
         user.increment_user_request_count(date)
-        if user.requests_count >= float(max_requests_number):
+        if user.requests_count >= max_requests_number:
             raise HTTPException(
                 status_code=429,
                 detail="You have reached the maximum number of requests for today.",

backend/test_main.py

@@ -0,0 +1,216 @@
+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")
+
+
+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
+    # Assuming the key_id is part of the api_key_info response. If not, adjust this.
+    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 UUID for brain_id
+    random_brain_id = str(uuid.uuid4())
+    random_brain_id = str(uuid.uuid4())
+
+    # 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 = {
+        "brain_id": random_brain_id,
+        "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_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