Update app.py

app.py

@@ -13,7 +13,8 @@ from pymongo import MongoClient
 import gridfs
 from bson.objectid import ObjectId
 from PIL import Image
-
+from fastapi.concurrency import run_in_threadpool
+import shutil
 import firebase_admin
 from firebase_admin import credentials, auth
 
@@ -93,12 +94,77 @@ class HealthResponse(BaseModel):
 # ---------------------------------------------------------------------
 # Endpoints
 # ---------------------------------------------------------------------
-@app.get("/health", response_model=HealthResponse)
-def health():
-    """Public health check"""
-    mongo.admin.command("ping")
-    return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
+# @app.get("/health", response_model=HealthResponse)
+# def health():
+#     """Public health check"""
+#     mongo.admin.command("ping")
+#     return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
+@app.get("/health")
+async def health():
+    checks = {
+        "api": "ok",
+        "mongo": "unknown",
+        "gridfs": "unknown",
+        "firebase": "unknown",
+        "huggingface": "unknown",
+        "timestamp": datetime.utcnow().isoformat()
+    }
+
+    errors = {}
+
+    # ------------------ MONGO CHECK ------------------
+    try:
+        def mongo_ping():
+            mongo.admin.command("ping")
+
+        await run_in_threadpool(mongo_ping)
+        checks["mongo"] = "connected"
+    except Exception as e:
+        checks["mongo"] = "failed"
+        errors["mongo"] = str(e)
+
+    # ------------------ GRIDFS CHECK ------------------
+    try:
+        def gridfs_test():
+            test_id = fs.put(b"healthcheck", filename="health.txt")
+            fs.delete(test_id)
+
+        await run_in_threadpool(gridfs_test)
+        checks["gridfs"] = "ok"
+    except Exception as e:
+        checks["gridfs"] = "failed"
+        errors["gridfs"] = str(e)
 
+    # ------------------ FIREBASE CHECK ------------------
+    try:
+        _ = firebase_admin.get_app()
+        checks["firebase"] = "initialized"
+    except Exception as e:
+        checks["firebase"] = "failed"
+        errors["firebase"] = str(e)
+
+    # ------------------ HUGGINGFACE CHECK ------------------
+    try:
+        _ = hf_client.model_info("Qwen/Qwen-Image-Edit")
+        checks["huggingface"] = "reachable"
+    except Exception as e:
+        checks["huggingface"] = "failed"
+        errors["huggingface"] = str(e)
+    # ------------------ FINAL STATUS ------------------
+    is_healthy = len(errors) == 0
+
+    if is_healthy:
+        checks["status"] = "healthy"
+        return checks
+
+    raise HTTPException(
+        status_code=503,
+        detail={
+            "status": "unhealthy",
+            "checks": checks,
+            "errors": errors
+        }
+    )
 
 @app.post("/generate")
 async def generate(