mirror of
https://github.com/StanGirard/quivr.git
synced 2024-10-26 14:00:37 +03:00
227 lines
10 KiB
Python
227 lines
10 KiB
Python
import os
|
|
import shutil
|
|
import time
|
|
from tempfile import SpooledTemporaryFile
|
|
|
|
import pypandoc
|
|
from auth.auth_bearer import JWTBearer
|
|
from crawl.crawler import CrawlWebsite
|
|
from fastapi import Depends, FastAPI, UploadFile
|
|
from llm.qa import get_qa_llm
|
|
from llm.summarization import llm_evaluate_summaries
|
|
from logger import get_logger
|
|
from middlewares.cors import add_cors_middleware
|
|
from models.chats import ChatMessage
|
|
from models.users import User
|
|
from parsers.github import process_github
|
|
from utils.file import convert_bytes, get_file_size
|
|
from utils.processors import filter_file
|
|
from utils.vectors import (CommonsDep, create_user, similarity_search,
|
|
update_user_request_count)
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
add_cors_middleware(app)
|
|
|
|
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
pypandoc.download_pandoc()
|
|
|
|
|
|
|
|
|
|
@app.post("/upload", dependencies=[Depends(JWTBearer())])
|
|
async def upload_file(commons: CommonsDep, file: UploadFile, enable_summarization: bool = False, credentials: dict = Depends(JWTBearer())):
|
|
max_brain_size = os.getenv("MAX_BRAIN_SIZE")
|
|
|
|
user = User(email=credentials.get('email', 'none'))
|
|
user_vectors_response = commons['supabase'].table("vectors").select(
|
|
"name:metadata->>file_name, size:metadata->>file_size", count="exact") \
|
|
.filter("user_id", "eq", user.email)\
|
|
.execute()
|
|
documents = user_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
|
|
user_unique_vectors = [dict(t) for t in set(tuple(d.items()) for d in documents)]
|
|
|
|
current_brain_size = sum(float(doc['size']) for doc in user_unique_vectors)
|
|
|
|
file_size = get_file_size(file)
|
|
|
|
remaining_free_space = float(max_brain_size) - (current_brain_size)
|
|
|
|
if remaining_free_space - file_size < 0:
|
|
message = {"message": f"❌ User's brain will exceed maximum capacity with this upload. Maximum file allowed is : {convert_bytes(remaining_free_space)}", "type": "error"}
|
|
else:
|
|
message = await filter_file(file, enable_summarization, commons['supabase'], user)
|
|
|
|
return message
|
|
|
|
|
|
@app.post("/chat/", dependencies=[Depends(JWTBearer())])
|
|
async def chat_endpoint(commons: CommonsDep, chat_message: ChatMessage, credentials: dict = Depends(JWTBearer())):
|
|
user = User(email=credentials.get('email', 'none'))
|
|
date = time.strftime("%Y%m%d")
|
|
max_requests_number = os.getenv("MAX_REQUESTS_NUMBER")
|
|
response = commons['supabase'].from_('users').select(
|
|
'*').filter("user_id", "eq", user.email).filter("date", "eq", date).execute()
|
|
|
|
|
|
userItem = next(iter(response.data or []), {"requests_count": 0})
|
|
old_request_count = userItem['requests_count']
|
|
|
|
history = chat_message.history
|
|
history.append(("user", chat_message.question))
|
|
|
|
qa = get_qa_llm(chat_message, user.email)
|
|
|
|
if old_request_count == 0:
|
|
create_user(user_id= user.email, date=date)
|
|
elif old_request_count < float(max_requests_number) :
|
|
update_user_request_count(user_id=user.email, date=date, requests_count= old_request_count+1)
|
|
else:
|
|
history.append(('assistant', "You have reached your requests limit"))
|
|
return {"history": history }
|
|
|
|
|
|
if chat_message.use_summarization:
|
|
# 1. get summaries from the vector store based on question
|
|
summaries = similarity_search(
|
|
chat_message.question, table='match_summaries')
|
|
# 2. evaluate summaries against the question
|
|
evaluations = llm_evaluate_summaries(
|
|
chat_message.question, summaries, chat_message.model)
|
|
# 3. pull in the top documents from summaries
|
|
logger.info('Evaluations: %s', evaluations)
|
|
if evaluations:
|
|
reponse = commons['supabase'].from_('vectors').select(
|
|
'*').in_('id', values=[e['document_id'] for e in evaluations]).execute()
|
|
# 4. use top docs as additional context
|
|
additional_context = '---\nAdditional Context={}'.format(
|
|
'---\n'.join(data['content'] for data in reponse.data)
|
|
) + '\n'
|
|
model_response = qa(
|
|
{"question": additional_context + chat_message.question})
|
|
else:
|
|
model_response = qa({"question": chat_message.question})
|
|
history.append(("assistant", model_response["answer"]))
|
|
|
|
return {"history": history}
|
|
|
|
|
|
@app.post("/crawl/", dependencies=[Depends(JWTBearer())])
|
|
async def crawl_endpoint(commons: CommonsDep, crawl_website: CrawlWebsite, enable_summarization: bool = False, credentials: dict = Depends(JWTBearer())):
|
|
max_brain_size = os.getenv("MAX_BRAIN_SIZE")
|
|
|
|
user = User(email=credentials.get('email', 'none'))
|
|
user_vectors_response = commons['supabase'].table("vectors").select(
|
|
"name:metadata->>file_name, size:metadata->>file_size", count="exact") \
|
|
.filter("user_id", "eq", user.email)\
|
|
.execute()
|
|
documents = user_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
|
|
user_unique_vectors = [dict(t) for t in set(tuple(d.items()) for d in documents)]
|
|
|
|
current_brain_size = sum(float(doc['size']) for doc in user_unique_vectors)
|
|
|
|
file_size = 1000000
|
|
|
|
remaining_free_space = float(max_brain_size) - (current_brain_size)
|
|
|
|
if remaining_free_space - file_size < 0:
|
|
message = {"message": f"❌ User's brain will exceed maximum capacity with this upload. Maximum file allowed is : {convert_bytes(remaining_free_space)}", "type": "error"}
|
|
else:
|
|
user = User(email=credentials.get('email', 'none'))
|
|
if not crawl_website.checkGithub():
|
|
|
|
file_path, file_name = crawl_website.process()
|
|
|
|
# Create a SpooledTemporaryFile from the file_path
|
|
spooled_file = SpooledTemporaryFile()
|
|
with open(file_path, 'rb') as f:
|
|
shutil.copyfileobj(f, spooled_file)
|
|
|
|
# Pass the SpooledTemporaryFile to UploadFile
|
|
file = UploadFile(file=spooled_file, filename=file_name)
|
|
message = await filter_file(file, enable_summarization, commons['supabase'], user=user)
|
|
return message
|
|
else:
|
|
message = await process_github(crawl_website.url, "false", user=user, supabase=commons['supabase'])
|
|
|
|
|
|
@app.get("/explore", dependencies=[Depends(JWTBearer())])
|
|
async def explore_endpoint(commons: CommonsDep,credentials: dict = Depends(JWTBearer()) ):
|
|
user = User(email=credentials.get('email', 'none'))
|
|
response = commons['supabase'].table("vectors").select(
|
|
"name:metadata->>file_name, size:metadata->>file_size", count="exact").filter("user_id", "eq", user.email).execute()
|
|
documents = 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_data = [dict(t) for t in set(tuple(d.items()) for d in documents)]
|
|
# Sort the list of documents by size in decreasing order
|
|
unique_data.sort(key=lambda x: int(x['size']), reverse=True)
|
|
|
|
return {"documents": unique_data}
|
|
|
|
|
|
@app.delete("/explore/{file_name}", dependencies=[Depends(JWTBearer())])
|
|
async def delete_endpoint(commons: CommonsDep, file_name: str, credentials: dict = Depends(JWTBearer())):
|
|
user = User(email=credentials.get('email', 'none'))
|
|
# Cascade delete the summary from the database first, because it has a foreign key constraint
|
|
commons['supabase'].table("summaries").delete().match(
|
|
{"metadata->>file_name": file_name}).execute()
|
|
commons['supabase'].table("vectors").delete().match(
|
|
{"metadata->>file_name": file_name, "user_id": user.email}).execute()
|
|
return {"message": f"{file_name} of user {user.email} has been deleted."}
|
|
|
|
|
|
@app.get("/explore/{file_name}", dependencies=[Depends(JWTBearer())])
|
|
async def download_endpoint(commons: CommonsDep, file_name: str,credentials: dict = Depends(JWTBearer()) ):
|
|
user = User(email=credentials.get('email', 'none'))
|
|
response = commons['supabase'].table("vectors").select(
|
|
"metadata->>file_name, metadata->>file_size, metadata->>file_extension, metadata->>file_url", "content").match({"metadata->>file_name": file_name, "user_id": user.email}).execute()
|
|
documents = response.data
|
|
# Returns all documents with the same file name
|
|
return {"documents": documents}
|
|
|
|
@app.get("/user", dependencies=[Depends(JWTBearer())])
|
|
async def get_user_endpoint(commons: CommonsDep, credentials: dict = Depends(JWTBearer())):
|
|
# Create a function that returns the unique documents out of the vectors
|
|
# Create a function that returns the list of documents that can take in what to put in the select + the filter
|
|
user = User(email=credentials.get('email', 'none'))
|
|
# Cascade delete the summary from the database first, because it has a foreign key constraint
|
|
user_vectors_response = commons['supabase'].table("vectors").select(
|
|
"name:metadata->>file_name, size:metadata->>file_size", count="exact") \
|
|
.filter("user_id", "eq", user.email)\
|
|
.execute()
|
|
documents = user_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
|
|
user_unique_vectors = [dict(t) for t in set(tuple(d.items()) for d in documents)]
|
|
|
|
current_brain_size = sum(float(doc['size']) for doc in user_unique_vectors)
|
|
|
|
max_brain_size = os.getenv("MAX_BRAIN_SIZE")
|
|
|
|
# Create function get user request stats -> nombre de requetes par jour + max number of requests -> svg to display the number of requests ? une fusee ?
|
|
user = User(email=credentials.get('email', 'none'))
|
|
date = time.strftime("%Y%m%d")
|
|
max_requests_number = os.getenv("MAX_REQUESTS_NUMBER")
|
|
requests_stats = commons['supabase'].from_('users').select(
|
|
'*').filter("user_id", "eq", user.email).execute()
|
|
|
|
return {"email":user.email,
|
|
"max_brain_size": max_brain_size,
|
|
"current_brain_size": current_brain_size,
|
|
"max_requests_number": max_requests_number,
|
|
"requests_stats" : requests_stats.data,
|
|
"date": date,
|
|
}
|
|
|
|
|
|
@app.get("/")
|
|
async def root():
|
|
return {"status": "OK"}
|