app.py

import io
from functions import *
from PyPDF2 import PdfReader
import pandas as pd
from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel
from fastapi.middleware.cors import CORSMiddleware
from langchain_community.document_loaders import UnstructuredURLLoader
from src.api.speech_api import speech_translator_router
from functions import client as supabase
from fastapi.exception_handlers import HTTPException

app = FastAPI(title="ConversAI", root_path="/api/v1")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(speech_translator_router, prefix="/speech")


@app.post("/signup")
async def sign_up(email, username, password):
    try:
        res, _ = supabase.auth.sign_up(
            {"email": email, "password": password, "role": "user"}
        )
        createUser(username=username)
        response = {
            "status": "success",
            "code": 200,
            "message": "Please check you email address for email verification",
        }

        return response
    except Exception as e:
        raise HTTPException(status_code=400, detail="Sign-up failed")


@app.post("/session-check")
async def check_session():
    res = supabase.auth.get_session()

    return res


@app.post("/login")
async def sign_in(email, password):
    try:
        res = supabase.auth.sign_in_with_password(
            {"email": email, "password": password}
        )
        user_id = res.user.id
        access_token = res.session.access_token
        refresh_token = res.session.refresh_token
        createUser(username=user_id)

        store_session_check = supabase.table("Stores").select("*").filter("StoreID", "eq", user_id).execute()
        try:
            store_id = store_session_check[1][0]["StoreID"]
        except:
            store_id = None

        if not store_id:

            response, _ = (
                supabase.table("Stores").insert(
                    {"AccessToken": access_token, "StoreID": user_id, "RefreshToken": refresh_token, }).execute())

            message = {
                "message": "success",
                "code": 200,
                "store_id": user_id,
                "access_token": access_token,
                "refresh_token": refresh_token}

            return message


        elif store_id == user_id:
            message = {
                "message": "You are already signed in please sign out first to sign in again",
                "code": 400,

            }
            return message

        else:
            message = {
                "message": "failed",
                "code": 400,
                "error": "Failed to sign in please check your credentials",
            }

        return message

    except:
        message = {
            "message": "You are already signed in please sign out first to sign in again",
            "code": 400,
            "error": "Failed to sign in please check your credentials",
        }

        return message


@app.post("/set-session-data")
async def set_session_data(access_token, refresh_token):
    res = supabase.auth.set_session(access_token, refresh_token)

    return res


@app.post("/logout")
async def sign_out(store_id):
    try:
        supabase.table("Stores").delete().eq(
            "StoreID", store_id
        ).execute()
        res = supabase.auth.sign_out()
        response = {"message": "success"}

        return response
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))


@app.post("/oauth")
async def oauth(provider):
    res = supabase.auth.sign_in_with_oauth(provider)

    return res


@app.post("/newChatbot")
async def newChatbot(chatbotName: str, username: str):
    currentBotCount = len(listTables(username=username)["output"])
    limit = client.table("ConversAI_UserConfig").select("chatbotLimit").eq("user_id", username).execute().data[0][
        "chatbotLimit"]
    if currentBotCount >= int(limit):
        return {
            "output": "CHATBOT LIMIT EXCEEDED"
        }
    client.table("ConversAI_ChatbotInfo").insert({"user_id": username, "chatbotname": chatbotName}).execute()
    chatbotName = f"convai-{username}-{chatbotName}"
    return createTable(tablename=chatbotName)


@app.post("/addPDF")
async def addPDFData(vectorstore: str, pdf: UploadFile = File(...)):
    pdf = await pdf.read()
    reader = PdfReader(io.BytesIO(pdf))
    text = ""
    for page in reader.pages:
        text += page.extract_text()
    username, chatbotname = vectorstore.split("-")[1], vectorstore.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    currentCount = df[(df["user_id"] == username) & (df["chatbotname"] == chatbotname)]["charactercount"].iloc[0]
    limit = client.table("ConversAI_UserConfig").select("tokenLimit").eq("user_id", username).execute().data[0][
        "tokenLimit"]
    newCount = currentCount + len(text)
    if newCount < int(limit):
        client.table("ConversAI_ChatbotInfo").update({"charactercount": str(newCount)}).eq("user_id", username).eq(
            "chatbotname", chatbotname).execute()
        return addDocuments(text=text, vectorstore=vectorstore)
    else:
        return {
            "output": "DOCUMENT EXCEEDING LIMITS, PLEASE TRY WITH A SMALLER DOCUMENT."
        }


@app.post("/scanAndReturnText")
async def returnText(pdf: UploadFile = File(...)):
    pdf = await pdf.read()
    text = getTextFromImagePDF(pdfBytes=pdf)
    return text


@app.post("/addText")
async def addText(vectorstore: str, text: str):
    username, chatbotname = vectorstore.split("-")[1], vectorstore.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    currentCount = df[(df["user_id"] == username) & (df["chatbotname"] == chatbotname)]["charactercount"].iloc[0]
    newCount = currentCount + len(text)
    limit = client.table("ConversAI_UserConfig").select("tokenLimit").eq("user_id", username).execute().data[0][
        "tokenLimit"]
    if newCount < int(limit):
        client.table("ConversAI_ChatbotInfo").update({"charactercount": str(newCount)}).eq("user_id", username).eq(
            "chatbotname", chatbotname).execute()
        return addDocuments(text=text, vectorstore=vectorstore)
    else:
        return {
            "output": "WEBSITE EXCEEDING LIMITS, PLEASE TRY WITH A SMALLER DOCUMENT."
        }


class AddQAPair(BaseModel):
    vectorstore: str
    question: str
    answer: str


@app.post("/addQAPair")
async def addText(addQaPair: AddQAPair):
    username, chatbotname = addQaPair.vectorstore.split("-")[1], addQaPair.vectorstore.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    currentCount = df[(df["user_id"] == username) & (df["chatbotname"] == chatbotname)]["charactercount"].iloc[0]
    qa = f"QUESTION: {addQaPair.question}\tANSWER: {addQaPair.answer}"
    newCount = currentCount + len(qa)
    limit = client.table("ConversAI_UserConfig").select("tokenLimit").eq("user_id", username).execute().data[0][
        "tokenLimit"]
    if newCount < int(limit):
        client.table("ConversAI_ChatbotInfo").update({"charactercount": str(newCount)}).eq("user_id", username).eq(
            "chatbotname", chatbotname).execute()
        return addDocuments(text=qa, vectorstore=addQaPair.vectorstore)
    else:
        return {
            "output": "WEBSITE EXCEEDING LIMITS, PLEASE TRY WITH A SMALLER DOCUMENT."
        }


@app.post("/addWebsite")
async def addWebsite(vectorstore: str, websiteUrls: list[str]):
    urls = websiteUrls
    loader = UnstructuredURLLoader(urls=urls)
    docs = loader.load()
    text = "\n\n".join(
        [f"Metadata:\n{docs[doc].metadata} \nPage Content:\n {docs[doc].page_content}" for doc in range(len(docs))])
    username, chatbotname = vectorstore.split("-")[1], vectorstore.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    currentCount = df[(df["user_id"] == username) & (df["chatbotname"] == chatbotname)]["charactercount"].iloc[0]
    newCount = currentCount + len(text)
    limit = client.table("ConversAI_UserConfig").select("tokenLimit").eq("user_id", username).execute().data[0][
        "tokenLimit"]
    if newCount < int(limit):
        client.table("ConversAI_ChatbotInfo").update({"charactercount": str(newCount)}).eq("user_id", username).eq(
            "chatbotname", chatbotname).execute()
        return addDocuments(text=text, vectorstore=vectorstore)
    else:
        return {
            "output": "WEBSITE EXCEEDING LIMITS, PLEASE TRY WITH A SMALLER DOCUMENT."
        }


@app.post("/answerQuery")
async def answerQuestion(query: str, vectorstore: str, llmModel: str = "llama3-70b-8192"):
    return answerQuery(query=query, vectorstore=vectorstore, llmModel=llmModel)


@app.post("/deleteChatbot")
async def delete(chatbotName: str):
    username, chatbotName = chatbotName.split("-")[1], chatbotName.split("-")[2]
    client.table('ConversAI_ChatbotInfo').delete().eq('user_id', username).eq('chatbotname', chatbotName).execute()
    return deleteTable(tableName=chatbotName)


@app.post("/listChatbots")
async def delete(username: str):
    return listTables(username=username)


@app.post("/getLinks")
async def crawlUrl(baseUrl: str):
    return {
        "urls": getLinks(url=baseUrl, timeout=30)
    }


@app.post("/getCurrentCount")
async def getCount(vectorstore: str):
    username, chatbotName = vectorstore.split("-")[1], vectorstore.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    return {
        "currentCount": df[(df['user_id'] == username) & (df['chatbotname'] == chatbotName)]['charactercount'].iloc[0]
    }


@app.post("/getYoutubeTranscript")
async def getYTTranscript(urls: str):
    return {
        "transcript": getTranscript(urls=urls)
    }


@app.post("/analyzeData")
async def analyzeAndAnswer(query: str, file: UploadFile = File(...)):
    extension = file.filename.split(".")[-1]
    try:
        if extension in ["xls", "xlsx", "xlsm", "xlsb"]:
            df = pd.read_excel(io.BytesIO(await file.read()))
            response = analyzeData(query=query, dataframe=df)
        elif extension == "csv":
            df = pd.read_csv(io.BytesIO(await file.read()))
            response = analyzeData(query=query, dataframe=df)
        else:
            response = "INVALID FILE TYPE"
        return {
            "output": response
        }
    except:
        return {
            "output": "UNABLE TO ANSWER QUERY"
        }