changing to fast api

app.py

@@ -1,76 +1,61 @@
-from flask import Flask, Response,request, jsonify
-from werkzeug.utils import secure_filename
-import os 
+from fastapi import FastAPI, UploadFile, File, HTTPException, Request 
+from fastapi.middleware.cors import CORSMiddleware
+import os
 from rag_engine import RagEngine
-from flask_cors import CORS, cross_origin
+from starlette.responses import JSONResponse
+from starlette.status import HTTP_400_BAD_REQUEST
+from fastapi.responses import StreamingResponse
+import asyncio
 
-app = Flask(__name__)
-cors = CORS(app)
-app.config["CORS_HEADERS"]= 'Content-Type'
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # or specify your allowed origins
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],)
+
+UPLOAD_FOLDER = "uploads"
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 
-app.config["UPLOAD_FOLDER"]= "uploads"
-os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)
 
 rag = RagEngine()
 
-@app.route("/upload", methods=["POST"])
-@cross_origin()
-def upload_pdf():
-    file = request.files.get('file')
+@app.post("/upload")
+async def upload_pdf(file: UploadFile = File(...)):
     if not file.filename.endswith(".pdf"):
-        return jsonify({"error":"Only pdf files are support"}),400
-    filename = secure_filename(file.filename)
-    filepath = os.path.join(app.config["UPLOAD_FOLDER"], filename)
-    file.save(filepath)
-    
-    try :
+        raise HTTPException(status_code=HTTP_400_BAD_REQUEST, detail="Only pdf files are supported")
+
+    filename = os.path.basename(file.filename)  # simple sanitization
+    filepath = os.path.join(UPLOAD_FOLDER, filename)
+
+    # Save uploaded file to disk
+    with open(filepath, "wb") as buffer:
+        content = await file.read()
+        buffer.write(content)
+
+    try:
         rag.index_pdf(filepath)
     except ValueError as ve:
-        return jsonify({"error": str(ve)}), 400
-    
-    return jsonify({"message":f"file {filename} uploaded and indexed successfully"})
+        raise HTTPException(status_code=HTTP_400_BAD_REQUEST, detail=str(ve))
 
+    return JSONResponse(content={"message": f"file {filename} uploaded and indexed successfully"})
 
-@app.route ("/stream", methods=["POST"])
-@cross_origin()
-def stream_answer():
-    
-    question = request.json.get("question", "")
-    print (question)
-    if not question.strip():
-        return jsonify({"error": "Empty question"}), 400
 
-    def generate():
-        for token in rag.ask_question_stream(question):
-            yield token
 
-    return Response(generate(), mimetype='text/plain')
-    
-    
-@app.route("/ask", methods=["POST"])
-@cross_origin()
-def ask():
-    question = request.json.get("question", "")
+@app.post("/stream")
+async def stream_answer(request:Request ):
+    data = await request.json()
+    question = data.get("question", "")
+    print(question)
     if not question.strip():
-        return jsonify({"error": "Empty question"}), 400
-    try :
-        answer = rag.ask_question(question)
-    except Exception as e:
-        return jsonify({"error": str(e)}),500
-    return jsonify({"message": answer})
+        raise HTTPException(status_code=400, detail="Empty question")
 
-@app.route("/stream_answer",methods=["POST"])
-@cross_origin()
-def stream_question():
-    data = request.get_json()
-    question = data.get("question","")
-    if not question:
-        return jsonify({"error": "No question provided"}),400
-    def event_stream():
-        for token in rag.stream_answer(question=question):
+    async def generate():
+        # Assuming rag.ask_question_stream is a generator
+        for token in rag.stream_answer(question):
             yield token
-    return Response(event_stream(), content_type ="text/event-stream")
-
+            await asyncio.sleep(0)  # yield control to event loop    
 
-if __name__ == "__main__":
-      app.run(host="0.0.0.0", port=7860)
+    return StreamingResponse(generate(), media_type="text/plain")