mirror of
https://github.com/StanGirard/quivr.git
synced 2024-12-23 19:32:30 +03:00
fix: bug if no brain (#635)
This commit is contained in:
parent
211740b400
commit
9f2cd7b7b6
2
.github/workflows/backend-tests.yml
vendored
2
.github/workflows/backend-tests.yml
vendored
@ -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: |
|
||||
|
@ -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(
|
||||
|
@ -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
|
||||
|
@ -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."}
|
||||
|
@ -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(
|
||||
|
@ -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"):
|
||||
|
@ -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
20
backend/tests/conftest.py
Normal 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
39
backend/tests/test_api_key.py
Normal 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
164
backend/tests/test_brains.py
Normal 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
116
backend/tests/test_chats.py
Normal 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
21
backend/tests/test_explore.py
Normal 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
|
208
backend/tests/test_upload.py
Normal file
208
backend/tests/test_upload.py
Normal 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
13
backend/tests/test_user.py
Normal 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
|
@ -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,
|
||||
});
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user