mirror of
https://github.com/QuivrHQ/quivr.git
synced 2024-09-11 14:36:35 +03:00
Feat/testing backend (#446)
* feat(pytest): added * feat(brains): added tests * feat(actions): pytest
This commit is contained in:
parent
bab76ba7e5
commit
e076bbe79f
40
.github/workflows/pytest.yml
vendored
Normal file
40
.github/workflows/pytest.yml
vendored
Normal file
@ -0,0 +1,40 @@
|
||||
name: Pytest
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ main ]
|
||||
pull_request:
|
||||
branches: [ main ]
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
defaults:
|
||||
run:
|
||||
working-directory: ./backend
|
||||
environment: preview
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: [ "3.11"]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
cache: 'pip'
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install pytest
|
||||
if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
|
||||
- name: Test with pytest
|
||||
env:
|
||||
SUPABASE_URL: ${{vars.SUPABASE_URL}}
|
||||
SUPABASE_SERVICE_KEY: ${{vars.SUPABASE_SERVICE_KEY}}
|
||||
OPENAI_API_KEY: ${{vars.OPENAI_API_KEY}}
|
||||
ANTHROPIC_API_KEY: ${{vars.ANTHROPIC_API_KEY}}
|
||||
JWT_SECRET_KEY: ${{vars.JWT_SECRET_KEY}}
|
||||
CI_TEST_API_KEY: ${{vars.CI_TEST_API_KEY}}
|
||||
run: |
|
||||
pytest
|
8
.github/workflows/vitest.yml
vendored
8
.github/workflows/vitest.yml
vendored
@ -12,14 +12,18 @@ jobs:
|
||||
defaults:
|
||||
run:
|
||||
working-directory: ./frontend
|
||||
strategy:
|
||||
matrix:
|
||||
node-version: [18]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v3
|
||||
- name: Use Node.js ${{ matrix.node-version }}
|
||||
uses: actions/setup-node@v3
|
||||
with:
|
||||
node-version: 18
|
||||
node-version: ${{ matrix.node-version }}
|
||||
cache: 'yarn'
|
||||
cache-dependency-path: frontend/yarn.lock
|
||||
- run: yarn
|
||||
- run: yarn run test-unit
|
||||
- run: yarn run test-unit
|
||||
- run: yarn run build
|
0
backend/auth/__init__.py
Normal file
0
backend/auth/__init__.py
Normal file
@ -25,8 +25,10 @@ async def verify_api_key(
|
||||
result.data[0]["creation_time"], "%Y-%m-%dT%H:%M:%S"
|
||||
).date()
|
||||
|
||||
# Check if the API key was created today: Todo remove this check and use deleted_time instead.
|
||||
if api_key_creation_date == current_date:
|
||||
# Check if the API key was created in the month of the current date
|
||||
if (api_key_creation_date.month == current_date.month) and (
|
||||
api_key_creation_date.year == current_date.year
|
||||
):
|
||||
return True
|
||||
return False
|
||||
except DateError:
|
||||
@ -62,7 +64,7 @@ async def get_user_from_api_key(
|
||||
)
|
||||
|
||||
return (
|
||||
{"email": user_email_data.data[0]["email"]}
|
||||
{"email": user_email_data.data[0]["email"], "sub": user_id}
|
||||
if user_email_data.data
|
||||
else {"email": None}
|
||||
)
|
||||
|
@ -55,4 +55,6 @@ class AuthBearer(HTTPBearer):
|
||||
|
||||
|
||||
def get_current_user(credentials: dict = Depends(AuthBearer())) -> User:
|
||||
return User(email=credentials.get("email", "none"), id=credentials.get("sub", "none"))
|
||||
return User(
|
||||
email=credentials.get("email", "none"), id=credentials.get("sub", "none")
|
||||
)
|
||||
|
0
backend/crawl/__init__.py
Normal file
0
backend/crawl/__init__.py
Normal file
@ -10,7 +10,7 @@ from pydantic import BaseModel
|
||||
class Brain(BaseModel):
|
||||
id: Optional[UUID] = None
|
||||
name: Optional[str] = "Default brain"
|
||||
status: Optional[str]= "public"
|
||||
status: Optional[str] = "public"
|
||||
model: Optional[str] = "gpt-3.5-turbo-0613"
|
||||
temperature: Optional[float] = 0.0
|
||||
max_tokens: Optional[int] = 256
|
||||
@ -31,9 +31,9 @@ class Brain(BaseModel):
|
||||
@property
|
||||
def brain_size(self):
|
||||
self.get_unique_brain_files()
|
||||
current_brain_size = sum(float(doc['size']) for doc in self.files)
|
||||
current_brain_size = sum(float(doc["size"]) for doc in self.files)
|
||||
|
||||
print('current_brain_size', current_brain_size)
|
||||
print("current_brain_size", current_brain_size)
|
||||
return current_brain_size
|
||||
|
||||
@property
|
||||
@ -65,30 +65,86 @@ class Brain(BaseModel):
|
||||
)
|
||||
return response.data
|
||||
|
||||
def delete_brain(self):
|
||||
self.commons["supabase"].table("brains").delete().match(
|
||||
{"brain_id": self.id}
|
||||
).execute()
|
||||
def delete_brain(self, user_id):
|
||||
print("user_id", user_id)
|
||||
print("self.id", self.id)
|
||||
results = (
|
||||
self.commons["supabase"]
|
||||
.table("brains_users")
|
||||
.select("*")
|
||||
.match({"brain_id": self.id, "user_id": user_id, "rights": "Owner"})
|
||||
.execute()
|
||||
)
|
||||
if len(results.data) == 0:
|
||||
print("You are not the owner of this brain.")
|
||||
return {"message": "You are not the owner of this brain."}
|
||||
else:
|
||||
results = (
|
||||
self.commons["supabase"]
|
||||
.table("brains_vectors")
|
||||
.delete()
|
||||
.match({"brain_id": self.id})
|
||||
.execute()
|
||||
)
|
||||
print("results", results)
|
||||
|
||||
results = (
|
||||
self.commons["supabase"]
|
||||
.table("brains_users")
|
||||
.delete()
|
||||
.match({"brain_id": self.id})
|
||||
.execute()
|
||||
)
|
||||
print("results", results)
|
||||
|
||||
results = (
|
||||
self.commons["supabase"]
|
||||
.table("brains")
|
||||
.delete()
|
||||
.match({"brain_id": self.id})
|
||||
.execute()
|
||||
)
|
||||
print("results", results)
|
||||
|
||||
def create_brain(self):
|
||||
commons = common_dependencies()
|
||||
response = commons["supabase"].table("brains").insert({"name": self.name}).execute()
|
||||
response = (
|
||||
commons["supabase"].table("brains").insert({"name": self.name}).execute()
|
||||
)
|
||||
# set the brainId with response.data
|
||||
|
||||
self.id = response.data[0]['brain_id']
|
||||
self.id = response.data[0]["brain_id"]
|
||||
return response.data
|
||||
|
||||
def create_brain_user(self, user_id: UUID, rights, default_brain):
|
||||
commons = common_dependencies()
|
||||
response = commons["supabase"].table("brains_users").insert({"brain_id": str(self.id), "user_id": str(user_id), "rights": rights, "default_brain": default_brain}).execute()
|
||||
|
||||
response = (
|
||||
commons["supabase"]
|
||||
.table("brains_users")
|
||||
.insert(
|
||||
{
|
||||
"brain_id": str(self.id),
|
||||
"user_id": str(user_id),
|
||||
"rights": rights,
|
||||
"default_brain": default_brain,
|
||||
}
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
|
||||
return response.data
|
||||
|
||||
def create_brain_vector(self, vector_id, file_sha1):
|
||||
response = (
|
||||
self.commons["supabase"]
|
||||
.table("brains_vectors")
|
||||
.insert({"brain_id": str(self.id), "vector_id": str(vector_id), "file_sha1": file_sha1})
|
||||
.insert(
|
||||
{
|
||||
"brain_id": str(self.id),
|
||||
"vector_id": str(vector_id),
|
||||
"file_sha1": file_sha1,
|
||||
}
|
||||
)
|
||||
.execute()
|
||||
)
|
||||
return response.data
|
||||
@ -121,22 +177,22 @@ class Brain(BaseModel):
|
||||
"""
|
||||
|
||||
response = (
|
||||
self.commons["supabase"]
|
||||
.from_("brains_vectors")
|
||||
.select("vector_id")
|
||||
.filter("brain_id", "eq", self.id)
|
||||
.execute()
|
||||
)
|
||||
|
||||
self.commons["supabase"]
|
||||
.from_("brains_vectors")
|
||||
.select("vector_id")
|
||||
.filter("brain_id", "eq", self.id)
|
||||
.execute()
|
||||
)
|
||||
|
||||
vector_ids = [item["vector_id"] for item in response.data]
|
||||
|
||||
print('vector_ids', vector_ids)
|
||||
print("vector_ids", vector_ids)
|
||||
|
||||
if len(vector_ids) == 0:
|
||||
return []
|
||||
|
||||
|
||||
self.files = self.get_unique_files_from_vector_ids(vector_ids)
|
||||
print('unique_files', self.files)
|
||||
print("unique_files", self.files)
|
||||
|
||||
return self.files
|
||||
|
||||
@ -145,19 +201,31 @@ class Brain(BaseModel):
|
||||
"""
|
||||
Retrieve unique user data vectors.
|
||||
"""
|
||||
print('vectors_ids', vectors_ids)
|
||||
print('tuple(vectors_ids)', tuple(vectors_ids))
|
||||
print("vectors_ids", vectors_ids)
|
||||
print("tuple(vectors_ids)", tuple(vectors_ids))
|
||||
if len(vectors_ids) == 1:
|
||||
vectors_response = self.commons['supabase'].table("vectors").select(
|
||||
"name:metadata->>file_name, size:metadata->>file_size", count="exact") \
|
||||
.filter("id", "eq", vectors_ids[0])\
|
||||
vectors_response = (
|
||||
self.commons["supabase"]
|
||||
.table("vectors")
|
||||
.select(
|
||||
"name:metadata->>file_name, size:metadata->>file_size",
|
||||
count="exact",
|
||||
)
|
||||
.filter("id", "eq", vectors_ids[0])
|
||||
.execute()
|
||||
)
|
||||
else:
|
||||
vectors_response = self.commons['supabase'].table("vectors").select(
|
||||
"name:metadata->>file_name, size:metadata->>file_size", count="exact") \
|
||||
.filter("id", "in", tuple(vectors_ids))\
|
||||
vectors_response = (
|
||||
self.commons["supabase"]
|
||||
.table("vectors")
|
||||
.select(
|
||||
"name:metadata->>file_name, size:metadata->>file_size",
|
||||
count="exact",
|
||||
)
|
||||
.filter("id", "in", tuple(vectors_ids))
|
||||
.execute()
|
||||
|
||||
)
|
||||
|
||||
documents = vectors_response.data # Access the data from the response
|
||||
# Convert each dictionary to a tuple of items, then to a set to remove duplicates, and then back to a dictionary
|
||||
unique_files = [dict(t) for t in set(tuple(d.items()) for d in documents)]
|
||||
@ -165,32 +233,52 @@ class Brain(BaseModel):
|
||||
|
||||
def delete_file_from_brain(self, file_name: str):
|
||||
# First, get the vector_ids associated with the file_name
|
||||
vector_response = self.commons["supabase"].table("vectors").select("id").filter("metadata->>file_name", "eq", file_name).execute()
|
||||
vector_response = (
|
||||
self.commons["supabase"]
|
||||
.table("vectors")
|
||||
.select("id")
|
||||
.filter("metadata->>file_name", "eq", file_name)
|
||||
.execute()
|
||||
)
|
||||
vector_ids = [item["id"] for item in vector_response.data]
|
||||
|
||||
# For each vector_id, delete the corresponding entry from the 'brains_vectors' table
|
||||
for vector_id in vector_ids:
|
||||
self.commons["supabase"].table("brains_vectors").delete().filter("vector_id", "eq", vector_id).filter("brain_id", "eq", self.id).execute()
|
||||
self.commons["supabase"].table("brains_vectors").delete().filter(
|
||||
"vector_id", "eq", vector_id
|
||||
).filter("brain_id", "eq", self.id).execute()
|
||||
|
||||
# Check if the vector is still associated with any other brains
|
||||
associated_brains_response = self.commons["supabase"].table("brains_vectors").select("brain_id").filter("vector_id", "eq", vector_id).execute()
|
||||
associated_brains = [item["brain_id"] for item in associated_brains_response.data]
|
||||
associated_brains_response = (
|
||||
self.commons["supabase"]
|
||||
.table("brains_vectors")
|
||||
.select("brain_id")
|
||||
.filter("vector_id", "eq", vector_id)
|
||||
.execute()
|
||||
)
|
||||
associated_brains = [
|
||||
item["brain_id"] for item in associated_brains_response.data
|
||||
]
|
||||
|
||||
# If the vector is not associated with any other brains, delete it from 'vectors' table
|
||||
if not associated_brains:
|
||||
self.commons["supabase"].table("vectors").delete().filter("id", "eq", vector_id).execute()
|
||||
self.commons["supabase"].table("vectors").delete().filter(
|
||||
"id", "eq", vector_id
|
||||
).execute()
|
||||
|
||||
return {"message": f"File {file_name} in brain {self.id} has been deleted."}
|
||||
|
||||
|
||||
|
||||
def get_default_user_brain(user: User):
|
||||
commons = common_dependencies()
|
||||
response = (
|
||||
commons["supabase"]
|
||||
.from_("brains_users") # I'm assuming this is the correct table
|
||||
.select("brain_id")
|
||||
.from_("brains_users") # I'm assuming this is the correct table
|
||||
.select("brain_id")
|
||||
.filter("user_id", "eq", user.id)
|
||||
.filter("default_brain", "eq", True) # Assuming 'default' is the correct column name
|
||||
.filter(
|
||||
"default_brain", "eq", True
|
||||
) # Assuming 'default' is the correct column name
|
||||
.execute()
|
||||
)
|
||||
|
||||
@ -207,8 +295,7 @@ def get_default_user_brain(user: User):
|
||||
.filter("brain_id", "eq", default_brain_id)
|
||||
.execute()
|
||||
)
|
||||
|
||||
|
||||
return brain_response.data[0] if brain_response.data else None
|
||||
|
||||
return None
|
||||
|
||||
|
0
backend/parsers/__init__.py
Normal file
0
backend/parsers/__init__.py
Normal file
0
backend/repository/__init__.py
Normal file
0
backend/repository/__init__.py
Normal file
0
backend/repository/chat/__init__.py
Normal file
0
backend/repository/chat/__init__.py
Normal file
@ -21,6 +21,7 @@ class ApiKeyInfo(BaseModel):
|
||||
|
||||
class ApiKey(BaseModel):
|
||||
api_key: str
|
||||
key_id: str
|
||||
|
||||
|
||||
api_key_router = APIRouter()
|
||||
@ -76,7 +77,7 @@ async def create_api_key(
|
||||
return {"api_key": "Error creating new API key."}
|
||||
logger.info(f"Created new API key for user {current_user.email}.")
|
||||
|
||||
return {"api_key": new_api_key}
|
||||
return {"api_key": new_api_key, "key_id": str(new_key_id)}
|
||||
|
||||
|
||||
@api_key_router.delete(
|
||||
|
@ -41,7 +41,9 @@ async def brain_endpoint(current_user: User = Depends(get_current_user)):
|
||||
return {"brains": brains}
|
||||
|
||||
|
||||
@brain_router.get("/brains/default/", dependencies=[Depends(AuthBearer())], tags=["Brain"])
|
||||
@brain_router.get(
|
||||
"/brains/default/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
|
||||
)
|
||||
async def get_default_brain_endpoint(current_user: User = Depends(get_current_user)):
|
||||
"""
|
||||
Retrieve the default brain for the current user.
|
||||
@ -49,7 +51,7 @@ async def get_default_brain_endpoint(current_user: User = Depends(get_current_us
|
||||
- `current_user`: The current authenticated user.
|
||||
- Returns the default brain for the user.
|
||||
|
||||
This endpoint retrieves the default brain associated with the current authenticated user.
|
||||
This endpoint retrieves the default brain associated with the current authenticated user.
|
||||
The default brain is defined as the brain marked as default in the brains_users table.
|
||||
"""
|
||||
|
||||
@ -87,15 +89,17 @@ async def get_brain_endpoint(brain_id: UUID):
|
||||
@brain_router.delete(
|
||||
"/brains/{brain_id}/", dependencies=[Depends(AuthBearer())], tags=["Brain"]
|
||||
)
|
||||
async def delete_brain_endpoint(brain_id: UUID, current_user: User = Depends(get_current_user),):
|
||||
async def delete_brain_endpoint(
|
||||
brain_id: UUID,
|
||||
current_user: User = Depends(get_current_user),
|
||||
):
|
||||
"""
|
||||
Delete a specific brain by brain ID.
|
||||
"""
|
||||
# [TODO] check if the user is the owner of the brain
|
||||
|
||||
current_user.id,
|
||||
|
||||
brain = Brain(id=brain_id)
|
||||
brain.delete_brain()
|
||||
brain.delete_brain(current_user.id)
|
||||
|
||||
return {"message": f"{brain_id} has been deleted."}
|
||||
|
||||
@ -127,15 +131,21 @@ async def create_brain_endpoint(
|
||||
"""
|
||||
|
||||
brain = Brain(name=brain.name)
|
||||
|
||||
|
||||
brain.create_brain()
|
||||
default_brain = get_default_user_brain(current_user)
|
||||
if default_brain:
|
||||
logger.info(f"Default brain already exists for user {current_user.id}")
|
||||
brain.create_brain_user(user_id=current_user.id, rights="Owner", default_brain=False)
|
||||
brain.create_brain_user(
|
||||
user_id=current_user.id, rights="Owner", default_brain=False
|
||||
)
|
||||
else:
|
||||
logger.info(f"Default brain does not exist for user {current_user.id}. It will be created.")
|
||||
brain.create_brain_user(user_id=current_user.id, rights="Owner", default_brain=True)
|
||||
logger.info(
|
||||
f"Default brain does not exist for user {current_user.id}. It will be created."
|
||||
)
|
||||
brain.create_brain_user(
|
||||
user_id=current_user.id, rights="Owner", default_brain=True
|
||||
)
|
||||
|
||||
return {"id": brain.id, "name": brain.name}
|
||||
|
||||
|
@ -76,10 +76,10 @@ def check_user_limit(
|
||||
):
|
||||
if user.user_openai_api_key is None:
|
||||
date = time.strftime("%Y%m%d")
|
||||
max_requests_number = os.getenv("MAX_REQUESTS_NUMBER")
|
||||
max_requests_number = os.getenv("MAX_REQUESTS_NUMBER", 1000)
|
||||
|
||||
user.increment_user_request_count(date)
|
||||
if user.requests_count >= float(max_requests_number):
|
||||
if user.requests_count >= max_requests_number:
|
||||
raise HTTPException(
|
||||
status_code=429,
|
||||
detail="You have reached the maximum number of requests for today.",
|
||||
|
216
backend/test_main.py
Normal file
216
backend/test_main.py
Normal file
@ -0,0 +1,216 @@
|
||||
import os
|
||||
import random
|
||||
import string
|
||||
import uuid
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
from main import app
|
||||
|
||||
client = TestClient(app)
|
||||
|
||||
API_KEY = os.getenv("CI_TEST_API_KEY")
|
||||
|
||||
|
||||
def test_read_main():
|
||||
response = client.get("/")
|
||||
assert response.status_code == 200
|
||||
assert response.json() == {"status": "OK"}
|
||||
|
||||
|
||||
def test_create_and_delete_api_key():
|
||||
# First, let's create an API key
|
||||
response = client.post(
|
||||
"/api-key",
|
||||
headers={
|
||||
"Authorization": "Bearer " + API_KEY,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
api_key_info = response.json()
|
||||
assert "api_key" in api_key_info
|
||||
|
||||
# Extract the created api_key from the response
|
||||
api_key = api_key_info["api_key"]
|
||||
|
||||
# Now, let's verify the API key
|
||||
verify_response = client.get(
|
||||
"/user",
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
},
|
||||
)
|
||||
assert verify_response.status_code == 200
|
||||
|
||||
# Now, let's delete the API key
|
||||
# Assuming the key_id is part of the api_key_info response. If not, adjust this.
|
||||
key_id = api_key_info["key_id"]
|
||||
delete_response = client.delete(
|
||||
f"/api-key/{key_id}", headers={"Authorization": f"Bearer {API_KEY}"}
|
||||
)
|
||||
assert delete_response.status_code == 200
|
||||
assert delete_response.json() == {"message": "API key deleted."}
|
||||
|
||||
|
||||
def test_retrieve_default_brain():
|
||||
# Making a GET request to the /brains/default/ endpoint
|
||||
response = client.get(
|
||||
"/brains/default/", headers={"Authorization": "Bearer " + API_KEY}
|
||||
)
|
||||
|
||||
# Assert that the response status code is 200 (HTTP OK)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Optionally, assert on specific fields in the response
|
||||
response_data = response.json()
|
||||
# e.g., assert that the response contains a 'brain_id' field
|
||||
assert "brain_id" in response_data
|
||||
|
||||
|
||||
def test_create_brain():
|
||||
# Generate a random UUID for brain_id
|
||||
random_brain_id = str(uuid.uuid4())
|
||||
random_brain_id = str(uuid.uuid4())
|
||||
|
||||
# Generate a random name for the brain
|
||||
random_brain_name = "".join(
|
||||
random.choices(string.ascii_letters + string.digits, k=10)
|
||||
)
|
||||
|
||||
# Set up the request payload
|
||||
payload = {
|
||||
"brain_id": random_brain_id,
|
||||
"name": random_brain_name,
|
||||
"status": "public",
|
||||
"model": "gpt-3.5-turbo-0613",
|
||||
"temperature": 0,
|
||||
"max_tokens": 256,
|
||||
"file_sha1": "",
|
||||
}
|
||||
|
||||
# Making a POST request to the /brains/ endpoint
|
||||
response = client.post(
|
||||
"/brains/", json=payload, headers={"Authorization": "Bearer " + API_KEY}
|
||||
)
|
||||
|
||||
# Assert that the response status code is 200 (HTTP OK)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Optionally, assert on specific fields in the response
|
||||
response_data = response.json()
|
||||
# e.g., assert that the response contains a 'brain_id' field
|
||||
assert "id" in response_data
|
||||
assert "name" in response_data
|
||||
|
||||
# Optionally, assert that the returned 'name' matches the one sent in the request
|
||||
assert response_data["name"] == payload["name"]
|
||||
|
||||
|
||||
def test_retrieve_all_brains():
|
||||
# Making a GET request to the /brains/ endpoint to retrieve all brains for the current user
|
||||
response = client.get("/brains/", headers={"Authorization": "Bearer " + API_KEY})
|
||||
|
||||
# Assert that the response status code is 200 (HTTP OK)
|
||||
assert response.status_code == 200
|
||||
|
||||
response_data = response.json()
|
||||
# Optionally, you can loop through the brains and assert on specific fields in each brain
|
||||
for brain in response_data["brains"]:
|
||||
assert "id" in brain
|
||||
assert "name" in brain
|
||||
|
||||
|
||||
def test_delete_all_brains():
|
||||
# First, retrieve all brains for the current user
|
||||
response = client.get("/brains/", headers={"Authorization": "Bearer " + API_KEY})
|
||||
|
||||
# Assert that the response status code is 200 (HTTP OK)
|
||||
assert response.status_code == 200
|
||||
|
||||
response_data = response.json()
|
||||
|
||||
# Loop through each brain and send a DELETE request
|
||||
for brain in response_data["brains"]:
|
||||
brain_id = brain["id"]
|
||||
|
||||
# Send a DELETE request to delete the specific brain
|
||||
delete_response = client.delete(
|
||||
f"/brains/{brain_id}/", headers={"Authorization": "Bearer " + API_KEY}
|
||||
)
|
||||
|
||||
# Assert that the DELETE response status code is 200 (HTTP OK)
|
||||
assert delete_response.status_code == 200
|
||||
|
||||
|
||||
def test_get_all_chats():
|
||||
# Making a GET request to the /chat endpoint to retrieve all chats
|
||||
response = client.get("/chat", headers={"Authorization": "Bearer " + API_KEY})
|
||||
|
||||
# Assert that the response status code is 200 (HTTP OK)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Assert that the response data is a list
|
||||
response_data = response.json()
|
||||
|
||||
# Optionally, you can loop through the chats and assert on specific fields
|
||||
for chat in response_data["chats"]:
|
||||
# e.g., assert that each chat object contains 'chat_id' and 'chat_name'
|
||||
assert "chat_id" in chat
|
||||
assert "chat_name" in chat
|
||||
|
||||
|
||||
def test_create_chat_and_talk():
|
||||
# Make a POST request to chat with the default brain and a random chat name
|
||||
random_chat_name = "".join(
|
||||
random.choices(string.ascii_letters + string.digits, k=10)
|
||||
)
|
||||
|
||||
brain_response = client.get(
|
||||
"/brains/default", headers={"Authorization": "Bearer " + API_KEY}
|
||||
)
|
||||
assert brain_response.status_code == 200
|
||||
default_brain_id = brain_response.json()["brain_id"]
|
||||
print("Default brain id: " + default_brain_id)
|
||||
|
||||
# Create a chat
|
||||
response = client.post(
|
||||
"/chat",
|
||||
json={"name": random_chat_name},
|
||||
headers={"Authorization": "Bearer " + API_KEY},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
# now talk to the chat with a question
|
||||
response_data = response.json()
|
||||
print(response_data)
|
||||
chat_id = response_data["chat_id"]
|
||||
response = client.post(
|
||||
f"/chat/{chat_id}/question?brain_id={default_brain_id}",
|
||||
json={
|
||||
"model": "gpt-3.5-turbo-0613",
|
||||
"question": "Hello, how are you?",
|
||||
"temperature": "0",
|
||||
"max_tokens": "256",
|
||||
},
|
||||
headers={"Authorization": "Bearer " + API_KEY},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
|
||||
response = client.post(
|
||||
f"/chat/{chat_id}/question?brain_id={default_brain_id}",
|
||||
json={
|
||||
"model": "gpt-4",
|
||||
"question": "Hello, how are you?",
|
||||
"temperature": "0",
|
||||
"max_tokens": "256",
|
||||
},
|
||||
headers={"Authorization": "Bearer " + API_KEY},
|
||||
)
|
||||
print(response)
|
||||
assert response.status_code == 200
|
||||
|
||||
# Now, let's delete the chat
|
||||
# Assuming the chat_id is part of the chat_info response. If not, adjust this.
|
||||
delete_response = client.delete(
|
||||
"/chat/" + chat_id, headers={"Authorization": "Bearer " + API_KEY}
|
||||
)
|
||||
assert delete_response.status_code == 200
|
0
backend/utils/__init__.py
Normal file
0
backend/utils/__init__.py
Normal file
0
backend/vectorstore/__init__.py
Normal file
0
backend/vectorstore/__init__.py
Normal file
Loading…
Reference in New Issue
Block a user