quivr/backend/main.py

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, Request, 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)
max_brain_size = os.getenv("MAX_BRAIN_SIZE")
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()




@app.post("/upload", dependencies=[Depends(JWTBearer())])
async def upload_file(request: Request, commons: CommonsDep,  file: UploadFile, enable_summarization: bool = False, credentials: dict = Depends(JWTBearer())):
   
    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 = 0
    if request.headers.get('Openai-Api-Key'):
        remaining_free_space =  float(max_brain_size_with_own_key) - (current_brain_size)
    else:
        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(request: Request, 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")
    user_openai_api_key = request.headers.get('Openai-Api-Key')

    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, user_openai_api_key)

    if user_openai_api_key is None:
        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})
    
    answer = model_response["answer"]
    
    # append sources (file_name) to answer
    if "source_documents" in model_response:
        logger.debug('Source Documents: %s', model_response["source_documents"])
        sources = [
            doc.metadata["file_name"] for doc in model_response["source_documents"]
            if "file_name" in doc.metadata]
        logger.debug('Sources: %s', sources)
        if sources:
            files = dict.fromkeys(sources)
            # # shall provide file links until pages available
            # files = [f"[{f}](/explore/{f})" for f in files]
            answer = answer + "\n\nRef: " + "; ".join(files)
    
    history.append(("assistant", answer))

    return {"history": history}


@app.post("/crawl/", dependencies=[Depends(JWTBearer())])
async def crawl_endpoint(request: Request,commons: CommonsDep, crawl_website: CrawlWebsite, enable_summarization: bool = False, credentials: dict = Depends(JWTBearer())):
    max_brain_size = os.getenv("MAX_BRAIN_SIZE")
    if request.headers.get('Openai-Api-Key'):
        max_brain_size = os.getenv("MAX_BRAIN_SIZE_WITH_KEY",209715200)
    


    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(request: Request, 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")
    if request.headers.get('Openai-Api-Key'):
        max_brain_size = max_brain_size_with_own_key


    # 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")
    if request.headers.get('Openai-Api-Key'):
        max_requests_number = 1000000
    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"}