|
import os
|
|
import shutil
|
|
import uvicorn
|
|
|
|
from pydantic import BaseModel
|
|
from fastapi import FastAPI, File, UploadFile, HTTPException
|
|
|
|
from app.app import data_ingestion, handle_query, DATA_DIR
|
|
|
|
class Query(BaseModel):
|
|
question: str
|
|
|
|
app = FastAPI()
|
|
|
|
@app.post("/upload")
|
|
async def upload_file(file: UploadFile = File(...)):
|
|
file_extension = os.path.splitext(file.filename)[1].lower()
|
|
if file_extension not in [".pdf", ".docx", ".txt"]:
|
|
raise HTTPException(status_code=400, detail="Invalid file type. Only PDF, DOCX, and TXT are allowed.")
|
|
|
|
file_path = os.path.join(DATA_DIR, file.filename)
|
|
with open(file_path, "wb") as buffer:
|
|
shutil.copyfileobj(file.file, buffer)
|
|
|
|
data_ingestion()
|
|
return {"message": "File uploaded and processed successfully"}
|
|
|
|
@app.post("/query")
|
|
async def query_document(query: Query):
|
|
if not os.listdir(DATA_DIR):
|
|
raise HTTPException(status_code=400, detail="No document has been uploaded yet.")
|
|
|
|
response = handle_query(query.question)
|
|
return {"response": response}
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="0.0.0.0", port=8000) |
