Update app.py

app.py

@@ -1,45 +1,36 @@
-from fastapi import FastAPI, Request
-from fastapi.responses import HTMLResponse, StreamingResponse
-from fastapi.templating import Jinja2Templates
+from flask import Flask, request, Response, render_template
 from pmbl import PMBL
+import os
+import json
 
-
-app = FastAPI()
+app = Flask(__name__, template_folder='.')
 pmbl = PMBL("./PMB-7b.Q6_K.gguf")  # Replace with the path to your model
 
-templates = Jinja2Templates(directory=".")
-
-@app.post("/chat")
-async def chat(request: Request):
+@app.route("/chat", methods=['POST'])
+def chat():
     try:
-        data = await request.json()
+        data = request.get_json()
         user_input = data["user_input"]
         mode = data["mode"]
         history = pmbl.get_chat_history(mode, user_input)
         response_generator = pmbl.generate_response(user_input, history, mode)
-        return StreamingResponse(response_generator, media_type="text/plain")
+        return Response(response_generator, mimetype='text/plain')
     except Exception as e:
         print(f"[SYSTEM] Error: {str(e)}")
-        return {"error": str(e)}
+        return json.dumps({"error": str(e)})
 
-@app.get("/", response_class=HTMLResponse)
-async def root(request: Request):
-    return templates.TemplateResponse("index.html", {"request": request})
+@app.route("/", methods=['GET'])
+def root():
+    return render_template("index.html")
 
-@app.post("/sleep")
-async def sleep():
+@app.route("/sleep", methods=['POST'])
+def sleep():
     try:
         pmbl.sleep_mode()
-        return {"message": "Sleep mode completed successfully"}
+        return json.dumps({"message": "Sleep mode completed successfully"})
     except Exception as e:
         print(f"[SYSTEM] Error: {str(e)}")
-        return {"error": str(e)}
-
-# Remove the main function as Hugging Face Spaces will automatically run the app. Using here to test
+        return json.dumps({"error": str(e)})
 
 if __name__ == "__main__":
-    import uvicorn
-    import asyncio
-
-    loop = asyncio.get_event_loop()
-    loop.run_until_complete(uvicorn.run(app, host="0.0.0.0", port=1771))
+    app.run(host="0.0.0.0", port=int(os.environ.get("PORT", 8000)))