Update app/main.py

app/main.py

@@ -1,9 +1,14 @@
 from contextlib import asynccontextmanager
+import asyncio
+from typing import Any, Dict
+
+from fastapi.responses import JSONResponse
 
 from fastapi import FastAPI
 
 from app.api.routes import router as api_router
-from app.db.mongo import close_mongo_connection, connect_to_mongo
+from app.core.openai_client import openai_client
+from app.db.mongo import close_mongo_connection, connect_to_mongo, get_client
 
 
 @asynccontextmanager
@@ -25,5 +30,70 @@ app.include_router(api_router, prefix="/api/v1")
 
 
 @app.get("/health", tags=["Health"])
-async def health_check() -> dict[str, str]:
-    return {"status": "ok"}
+async def health_check() -> JSONResponse | Dict[str, Any]:
+    checks: Dict[str, bool] = {"mongo": False, "openai": False}
+    errors: Dict[str, str] = {}
+
+    # Run checks with conservative timeouts to avoid hanging the health endpoint.
+    mongo_ok, mongo_err = await _run_check(_check_mongo)
+    checks["mongo"] = mongo_ok
+    if mongo_err:
+        errors["mongo"] = mongo_err
+
+    openai_ok, openai_err = await _run_check(_check_openai)
+    checks["openai"] = openai_ok
+    if openai_err:
+        errors["openai"] = openai_err
+
+    if errors:
+        return JSONResponse(status_code=503, content={"status": "unhealthy", "checks": checks, "errors": errors})
+
+    return {"status": "healthy", "checks": checks}
+
+
+async def _run_check(fn, timeout: float = 5.0) -> tuple[bool, str | None]:
+    try:
+        await asyncio.wait_for(fn(), timeout=timeout)
+        return True, None
+    except Exception as exc:  # pragma: no cover - defensive
+        return False, str(exc)
+
+
+async def _check_mongo() -> None:
+    await connect_to_mongo()
+    client = get_client()
+    await client.admin.command("ping")
+
+
+async def _check_openai() -> None:
+    await openai_client.models.list()
+
+# from contextlib import asynccontextmanager
+
+# from fastapi import FastAPI
+
+# from app.api.routes import router as api_router
+# from app.db.mongo import close_mongo_connection, connect_to_mongo
+
+
+# @asynccontextmanager
+# async def lifespan(_: FastAPI):
+#     await connect_to_mongo()
+#     try:
+#         yield
+#     finally:
+#         await close_mongo_connection()
+
+
+# app = FastAPI(
+#     title="Auto Categorizer API",
+#     version="1.0.0",
+#     lifespan=lifespan,
+# )
+
+# app.include_router(api_router, prefix="/api/v1")
+
+
+# @app.get("/health", tags=["Health"])
+# async def health_check() -> dict[str, str]:
+#     return {"status": "ok"}