fix lint issues

epam · Feb 12, 2024 · 1ce19f0 · 1ce19f0
1 parent d67ac48
commit 1ce19f0
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/aidial_sdk/application.py b/aidial_sdk/application.py
@@ -192,8 +192,11 @@ async def _healthcheck() -> JSONResponse:
         return JSONResponse(content={"status": "ok"})
 
     @staticmethod
-    def _exception_handler(request: Request, exc: HTTPException):
-        return JSONResponse(
-            status_code=exc.status_code,
-            content=exc.detail,
-        )
+    def _exception_handler(request: Request, exc: Exception):
+        if isinstance(exc, HTTPException):
+            return JSONResponse(
+                status_code=exc.status_code,
+                content=exc.detail,
+            )
+        else:
+            raise exc
diff --git a/aidial_sdk/header_propagator.py b/aidial_sdk/header_propagator.py
@@ -8,10 +8,11 @@
 from fastapi import FastAPI
 from requests import PreparedRequest
 from requests.sessions import Session
+from starlette.middleware import _MiddlewareClass
 from starlette.middleware.exceptions import ExceptionMiddleware
 
 
-class FastAPIMiddleware:
+class FastAPIMiddleware(_MiddlewareClass):
     def __init__(
         self,
         app: ExceptionMiddleware,