diff --git a/.github/workflows/backend-tests.yml b/.github/workflows/backend-tests.yml index e8e475385..be7771f1a 100644 --- a/.github/workflows/backend-tests.yml +++ b/.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: | diff --git a/backend/main.py b/backend/main.py index 31446f5d3..a0f1d35cd 100644 --- a/backend/main.py +++ b/backend/main.py @@ -18,25 +18,22 @@ 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(): - pypandoc.download_pandoc() + if not os.path.exists(pypandoc.get_pandoc_path()): + pypandoc.download_pandoc() app.include_router(brain_router) @@ -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( diff --git a/backend/models/brains.py b/backend/models/brains.py index e0e1f341d..f56e97eae 100644 --- a/backend/models/brains.py +++ b/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 diff --git a/backend/routes/brain_routes.py b/backend/routes/brain_routes.py index 57510f85c..32727c494 100644 --- a/backend/routes/brain_routes.py +++ b/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."} diff --git a/backend/routes/chat_routes.py b/backend/routes/chat_routes.py index 40246bf76..91ed78812 100644 --- a/backend/routes/chat_routes.py +++ b/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( diff --git a/backend/routes/upload_routes.py b/backend/routes/upload_routes.py index ab75b555d..775443fc0 100644 --- a/backend/routes/upload_routes.py +++ b/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"): diff --git a/backend/test_main.py b/backend/test_main.py deleted file mode 100644 index 37387a9a5..000000000 --- a/backend/test_main.py +++ /dev/null @@ -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 diff --git a/backend/tests/conftest.py b/backend/tests/conftest.py new file mode 100644 index 000000000..2f4e908ce --- /dev/null +++ b/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 diff --git a/backend/tests/test_api_key.py b/backend/tests/test_api_key.py new file mode 100644 index 000000000..ed2310937 --- /dev/null +++ b/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."} diff --git a/backend/tests/test_brains.py b/backend/tests/test_brains.py new file mode 100644 index 000000000..7b12022bc --- /dev/null +++ b/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" diff --git a/backend/tests/test_chats.py b/backend/tests/test_chats.py new file mode 100644 index 000000000..567a8660b --- /dev/null +++ b/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 diff --git a/backend/tests/test_explore.py b/backend/tests/test_explore.py new file mode 100644 index 000000000..67e9f1147 --- /dev/null +++ b/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 diff --git a/backend/test_file/test.csv b/backend/tests/test_files/test.csv similarity index 100% rename from backend/test_file/test.csv rename to backend/tests/test_files/test.csv diff --git a/backend/test_file/test.pdf b/backend/tests/test_files/test.pdf similarity index 100% rename from backend/test_file/test.pdf rename to backend/tests/test_files/test.pdf diff --git a/backend/test_file/test.txt b/backend/tests/test_files/test.txt similarity index 100% rename from backend/test_file/test.txt rename to backend/tests/test_files/test.txt diff --git a/backend/tests/test_upload.py b/backend/tests/test_upload.py new file mode 100644 index 000000000..eabbd4e66 --- /dev/null +++ b/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 diff --git a/backend/tests/test_user.py b/backend/tests/test_user.py new file mode 100644 index 000000000..85b82f5e4 --- /dev/null +++ b/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 diff --git a/frontend/app/chat/[chatId]/hooks/useQuestion.ts b/frontend/app/chat/[chatId]/hooks/useQuestion.ts index 76718f22d..0a23d0045 100644 --- a/frontend/app/chat/[chatId]/hooks/useQuestion.ts +++ b/frontend/app/chat/[chatId]/hooks/useQuestion.ts @@ -25,13 +25,9 @@ export const useQuestion = (): UseChatService => { chatId: string, chatQuestion: ChatQuestion ): Promise => { - if (currentBrain?.id === undefined) { - throw new Error("No current brain"); - } - const response = await addQuestion({ chatId, - brainId: currentBrain.id, + brainId: currentBrain?.id ?? "", chatQuestion, });