QuivrHQ/quivr
1
1
Fork 0
mirror of https://github.com/QuivrHQ/quivr.git synced 2024-10-26 15:18:16 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix: bug if no brain (#635)

Browse Source
This commit is contained in:
Matt 2023-07-14 20:02:26 +01:00 committed by GitHub
parent 211740b400
commit 9f2cd7b7b6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
18 changed files with 664 additions and 607 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/backend-tests.yml vendored
View File

@ -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: |

14
backend/main.py
View File

@ -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(

20
backend/models/brains.py
View File

@ -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

69
backend/routes/brain_routes.py
View File

@ -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."}

46
backend/routes/chat_routes.py
View File

@ -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(

3
backend/routes/upload_routes.py
View File

@ -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"):

530
backend/test_main.py
View File

@ -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

20
backend/tests/conftest.py Normal file
View File

@ -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

39
backend/tests/test_api_key.py Normal file
View File

@ -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."}

164
backend/tests/test_brains.py Normal file
View File

@ -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"

116
backend/tests/test_chats.py Normal file
View File

@ -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

21
backend/tests/test_explore.py Normal file
View File

@ -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

0
backend/test_file/test.csv → backend/tests/test_files/test.csv
View File

0
backend/test_file/test.pdf → backend/tests/test_files/test.pdf
View File

0
backend/test_file/test.txt → backend/tests/test_files/test.txt
View File

208
backend/tests/test_upload.py Normal file
View File

@ -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

13
backend/tests/test_user.py Normal file
View File

@ -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

6
frontend/app/chat/[chatId]/hooks/useQuestion.ts
View File

@ -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,
});