From 5a9e98daddf1b9f2c7757c62fd3727ef9b2b821f Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Thu, 5 Oct 2023 17:54:11 -0700
Subject: [PATCH] handle WyvernError in except handler

---
 wyvern/web_frameworks/fastapi.py | 36 +++++++++++++-------------------
 1 file changed, 15 insertions(+), 21 deletions(-)

diff --git a/wyvern/web_frameworks/fastapi.py b/wyvern/web_frameworks/fastapi.py
index 5f16a14..baba3d3 100644
--- a/wyvern/web_frameworks/fastapi.py
+++ b/wyvern/web_frameworks/fastapi.py
@@ -72,27 +72,6 @@ def __init__(self, host: str = "127.0.0.1", port: int = 8000) -> None:
         self.host = host
         self.port = port
 
-        @self.app.exception_handler(WyvernError)
-        async def wyvern_exception_handler(request: Request, exc: WyvernError):
-            json = await request.json()
-            logger.warning(f"wyvern_error={exc} request_payload={json}")
-            return JSONResponse(
-                status_code=400,
-                content={
-                    "detail": str(exc),
-                    "error_code": exc.error_code,
-                },
-            )
-
-        @self.app.exception_handler(ValidationError)
-        async def pydantic_validation_error_handler(
-            request: Request,
-            exc: ValidationError,
-        ):
-            json = await request.json()
-            logger.exception(f"Validation error error={exc} request_payload={json}")
-            raise HTTPException(status_code=422, detail=exc.errors())
-
         @self.app.get("/healthcheck")
         async def healthcheck() -> Dict[str, str]:
             return {"status": "OK"}
@@ -189,6 +168,21 @@ async def post(
 
                 # profiler.stop()
                 # profiler.print(show_all=True)
+            except ValidationError as e:
+                logger.exception(f"Validation error error={e} request_payload={json}")
+                raise HTTPException(status_code=422, detail=e.errors())
+            except WyvernError as e:
+                logger.warning(f"Wyvern error error={e} request_payload={json}")
+                return JSONResponse(
+                    status_code=400,
+                    content={
+                        "detail": str(e),
+                        "error_code": e.error_code,
+                    },
+                )
+            except Exception as e:
+                logger.exception(f"Unexpected error error={e} request_payload={json}")
+                raise HTTPException(status_code=500, detail=str(e))
             finally:
                 request_context.reset()
             if not output: