Update app.py

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Depends
 from fastapi.responses import HTMLResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from modules.pmbl import PMBL
@@ -14,7 +14,12 @@ app = FastAPI(docs_url=None, redoc_url=None)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 app.mount("/templates", StaticFiles(directory="templates"), name="templates")
 
-pmbl = PMBL("./PMB-7b.Q6_K.gguf", gpu_layers=50)
+async def get_pmbl():
+    pmbl = PMBL("./PMB-7b.Q6_K.gguf", gpu_layers=50)
+    try:
+        yield pmbl
+    finally:
+        await pmbl.close()
 
 @app.head("/")
 @app.get("/")
@@ -23,12 +28,12 @@ def index() -> HTMLResponse:
         return HTMLResponse(content=f.read())
 
 @app.post("/chat")
-async def chat(request: Request):
+async def chat(request: Request, pmbl: PMBL = Depends(get_pmbl)):
     try:
         data = await request.json()
         user_input = data["user_input"]
         mode = data["mode"]
-        history = pmbl.get_chat_history(mode, user_input)
+        history = await pmbl.get_chat_history(mode, user_input)
         response_generator = pmbl.generate_response(user_input, history, mode)
         return StreamingResponse(response_generator, media_type="text/plain")
     except Exception as e:
@@ -36,9 +41,9 @@ async def chat(request: Request):
         return {"error": str(e)}
 
 @app.post("/sleep")
-async def sleep():
+async def sleep(pmbl: PMBL = Depends(get_pmbl)):
     try:
-        pmbl.sleep_mode()
+        await pmbl.sleep_mode()
         return {"message": "Sleep mode completed successfully"}
     except Exception as e:
         print(f"[SYSTEM] Error: {str(e)}")