|
|
|
import os, random, logging, pickle, shutil |
|
from dotenv import load_dotenv, find_dotenv |
|
from typing import Optional |
|
from pydantic import BaseModel, Field |
|
|
|
from fastapi import FastAPI, HTTPException, File, UploadFile, status |
|
from fastapi.responses import HTMLResponse |
|
from fastapi.middleware.cors import CORSMiddleware |
|
|
|
from engine.processing import process_pdf, index_data, empty_collection, vector_search |
|
from rag.rag import rag_it |
|
|
|
from engine.logger import logger |
|
|
|
from settings import datadir |
|
|
|
os.makedirs(datadir, exist_ok=True) |
|
|
|
app = FastAPI() |
|
|
|
environment = os.getenv("ENVIRONMENT", "dev") |
|
|
|
if environment == "dev": |
|
logger.warning("Running in development mode - allowing CORS for all origins") |
|
app.add_middleware( |
|
CORSMiddleware, |
|
allow_origins=["*"], |
|
allow_credentials=True, |
|
allow_methods=["*"], |
|
allow_headers=["*"], |
|
) |
|
|
|
try: |
|
load_dotenv(find_dotenv('env')) |
|
|
|
except Exception as e: |
|
pass |
|
|
|
|
|
@app.get("/", response_class=HTMLResponse) |
|
def read_root(): |
|
logger.info("Title displayed on home page") |
|
return """ |
|
<html> |
|
<body> |
|
<h1>Welcome to FinExpert, a RAG system designed by JP Bianchi!</h1> |
|
</body> |
|
</html> |
|
""" |
|
|
|
|
|
@app.get("/ping/") |
|
def ping(): |
|
""" Testing """ |
|
logger.info("Someone is pinging the server") |
|
return {"answer": str(int(random.random() * 100))} |
|
|
|
|
|
@app.delete("/erase_data/") |
|
def erase_data(): |
|
""" Erase all files in the data directory, but not the vector store """ |
|
if len(os.listdir(datadir)) == 0: |
|
logger.info("No data to erase") |
|
return {"message": "No data to erase"} |
|
|
|
shutil.rmtree(datadir, ignore_errors=True) |
|
os.mkdir(datadir) |
|
logger.warning("All data has been erased") |
|
return {"message": "All data has been erased"} |
|
|
|
|
|
@app.delete("/empty_collection/") |
|
def delete_vectors(): |
|
""" Empty the collection in the vector store """ |
|
try: |
|
status = empty_collection() |
|
return {f"""message": "Collection{'' if status else ' NOT'} erased!"""} |
|
except Exception as e: |
|
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)) |
|
|
|
@app.get("/list_files/") |
|
def list_files(): |
|
""" List all files in the data directory """ |
|
files = os.listdir(datadir) |
|
logger.info(f"Files in data directory: {files}") |
|
return {"files": files} |
|
|
|
|
|
@app.post("/upload/") |
|
|
|
async def upload_file(file: UploadFile = File(...)): |
|
""" Uploads a file in data directory, for later indexing """ |
|
try: |
|
filepath = os.path.join(datadir, file.filename) |
|
logger.info(f"Fiename detected: {file.filename}") |
|
if os.path.exists(filepath): |
|
logger.warning(f"File {file.filename} already exists: no processing done") |
|
return {"message": f"File {file.filename} already exists: no processing done"} |
|
|
|
else: |
|
logger.info(f"Receiving file: {file.filename}") |
|
contents = await file.read() |
|
logger.info(f"File reception complete!") |
|
|
|
except Exception as e: |
|
logger.error(f"Error during file upload: {str(e)}") |
|
return {"message": f"Error during file upload: {str(e)}"} |
|
|
|
if file.filename.endswith('.pdf'): |
|
|
|
|
|
with open(filepath, 'wb') as f: |
|
f.write(contents) |
|
|
|
try: |
|
logger.info(f"Starting to process {file.filename}") |
|
new_content = process_pdf(filepath) |
|
success = {"message": f"Successfully uploaded {file.filename}"} |
|
success.update(new_content) |
|
return success |
|
|
|
except Exception as e: |
|
return {"message": f"Failed to extract text from PDF: {str(e)}"} |
|
else: |
|
return {"message": "Only PDF files are accepted"} |
|
|
|
|
|
@app.post("/create_index/") |
|
async def create_index(): |
|
""" Create an index for the uploaded files """ |
|
|
|
logger.info("Creating index for uploaded files") |
|
try: |
|
msg = index_data() |
|
return {"message": msg} |
|
except Exception as e: |
|
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)) |
|
|
|
|
|
class Question(BaseModel): |
|
question: str |
|
|
|
@app.post("/ask/") |
|
async def hybrid_search(question: Question): |
|
logger.info(f"Processing question: {question.question}") |
|
try: |
|
search_results = vector_search(question.question) |
|
logger.info(f"Answer: {search_results}") |
|
return {"answer": search_results} |
|
except Exception as e: |
|
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)) |
|
|
|
|
|
@app.post("/ragit/") |
|
async def ragit(question: Question): |
|
logger.info(f"Processing question: {question.question}") |
|
try: |
|
search_results = vector_search(question.question) |
|
logger.info(f"Search results generated: {search_results}") |
|
|
|
answer = rag_it(question.question, search_results) |
|
|
|
logger.info(f"Answer: {answer}") |
|
return {"answer": answer} |
|
except Exception as e: |
|
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=str(e)) |
|
|
|
if __name__ == '__main__': |
|
import uvicorn |
|
from os import getenv |
|
port = int(getenv("PORT", 80)) |
|
print(f"Starting server on port {port}") |
|
reload = True if environment == "dev" else False |
|
uvicorn.run("main:app", host="0.0.0.0", port=port, reload=reload) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
