app/main.py

import os
import shutil
import time
from typing import List
from fastapi import FastAPI, UploadFile
from fastapi.middleware import Middleware
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from .rag import ChatPDF

middleware = [
    Middleware(
        CORSMiddleware,
        allow_origins=["*"],
        allow_methods=['*'],
        allow_headers=['*']
    )
]

app = FastAPI(middleware=middleware)

files_dir = os.path.expanduser("~/wtp_be_files/")
session_assistant = ChatPDF()


def astreamer(generator):
    t0 = time.time()
    for i in generator:
        print(f"Chunk being yielded (time {int((time.time()-t0)*1000)}ms)", flush=True)
        yield i
    print(f"Over (time {int((time.time()-t0)*1000)}ms)", flush=True)


@app.get("/query")
async def process_input(text: str):
    if text and len(text.strip()) > 0:
        text = text.strip()
        streaming_response = session_assistant.ask(text)
        return StreamingResponse(astreamer(streaming_response.response_gen), media_type='text/event-stream')


@app.post("/upload")
def upload(files: list[UploadFile]):  
    try:
        os.makedirs(files_dir)
        for file in files:
            try:
                path = f"{files_dir}/{file.filename}"
                file.file.seek(0)
                with open(path, 'wb') as destination:
                    shutil.copyfileobj(file.file, destination)
            finally:
                file.file.close()
    finally:
        session_assistant.ingest(files_dir)
        shutil.rmtree(files_dir)

    return "Files inserted!"


@app.get("/clear")
def ping():
    session_assistant.clear()
    return "All files have been cleared."


@app.get("/")
def ping():
    return "Pong!"