feat: migrated to DIAL SDK exceptions with headers

epam · Nov 26, 2024 · d5f6b59 · d5f6b59
1 parent a78d201
commit d5f6b59
Show file tree

Hide file tree

Showing 4 changed files with 62 additions and 86 deletions.
diff --git a/aidial_interceptors_sdk/utils/_exceptions.py b/aidial_interceptors_sdk/utils/_exceptions.py
@@ -3,68 +3,39 @@
 from typing import Dict
 
 from aidial_sdk.exceptions import HTTPException as DialException
-from fastapi import HTTPException as FastAPIException
-from fastapi.responses import JSONResponse as FastAPIResponse
 from openai import APIConnectionError, APIError, APIStatusError, APITimeoutError
-from typing_extensions import override
 
 _log = logging.getLogger(__name__)
 
 
-# TODO: support headers in DIAL SDK exception
-class DialExceptionWithHeaders(DialException):
-    headers: Dict[str, str] | None = None
-
-    def __init__(self, *, headers: Dict[str, str] | None = None, **kwargs):
-        super().__init__(**kwargs)
-        self.headers = headers
-
-    @classmethod
-    def create(
-        cls,
-        status_code: int,
-        content: dict | str,
-        headers: Dict[str, str] | None = None,
+def _parse_dial_exception(
+    status_code: int,
+    content: dict | str,
+    headers: Dict[str, str] | None = None,
+):
+    if (
+        isinstance(content, dict)
+        and (error := content.get("error"))
+        and isinstance(error, dict)
     ):
-        if (
-            isinstance(content, dict)
-            and (error := content.get("error"))
-            and isinstance(error, dict)
-        ):
-            message = error.get("message") or "Unknown error"
-            code = error.get("code")
-            type = error.get("type")
-            param = error.get("param")
-            display_message = error.get("display_message")
-        else:
-            message = content
-            code = type = param = display_message = None
-
-        return cls(
-            status_code=status_code,
-            message=message,
-            type=type,
-            param=param,
-            code=code,
-            display_message=display_message,
-            headers=headers,
-        )
+        message = error.get("message") or "Unknown error"
+        code = error.get("code")
+        type = error.get("type")
+        param = error.get("param")
+        display_message = error.get("display_message")
+    else:
+        message = str(content)
+        code = type = param = display_message = None
 
-    @override
-    def to_fastapi_response(self) -> FastAPIResponse:
-        return FastAPIResponse(
-            status_code=self.status_code,
-            content=self.json_error(),
-            headers=self.headers,
-        )
-
-    @override
-    def to_fastapi_exception(self) -> FastAPIException:
-        return FastAPIException(
-            status_code=self.status_code,
-            detail=self.json_error(),
-            headers=self.headers,
-        )
+    return DialException(
+        status_code=status_code,
+        message=message,
+        type=type,
+        param=param,
+        code=code,
+        display_message=display_message,
+        headers=headers,
+    )
 
 
 def to_dial_exception(exc: Exception) -> DialException:
@@ -95,7 +66,7 @@ def to_dial_exception(exc: Exception) -> DialException:
         except Exception:
             content = r.text
 
-        return DialExceptionWithHeaders.create(
+        return _parse_dial_exception(
             status_code=r.status_code,
             headers=plain_headers,
             content=content,
@@ -110,7 +81,7 @@ def to_dial_exception(exc: Exception) -> DialException:
             except Exception:
                 pass
 
-        return DialExceptionWithHeaders.create(
+        return _parse_dial_exception(
             status_code=status_code,
             headers={},
             content={"error": exc.body or {}},

diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -25,7 +25,9 @@ python = ">=3.11,<4.0"
 fastapi = ">=0.51,<1.0"
 httpx = ">=0.25.0,<1.0"
 openai = ">=1.32.0,<2.0"
-aidial-sdk = { version = "^0.15.0", extras = ["telemetry"] }
+# FIXME: revert to a release version
+# aidial-sdk = { version = "^0.15.0", extras = ["telemetry"] }
+aidial-sdk = { git = "https://github.com/epam/ai-dial-sdk.git", branch = "feat/support-headers-in-dial-exception", extras = ["telemetry"] }
 
 # Extras for examples
 aiostream = { version = "^0.6.2", optional = true }

diff --git a/tests/test_exceptions.py b/tests/test_exceptions.py
@@ -6,13 +6,13 @@
 from aidial_interceptors_sdk.chat_completion.base import (
     ChatCompletionNoOpInterceptor,
 )
-from aidial_interceptors_sdk.utils._exceptions import DialExceptionWithHeaders
+from aidial_interceptors_sdk.utils._exceptions import _parse_dial_exception
 from tests.utils.applications import create_broken_application
 from tests.utils.chunks import create_chunk_checker, create_sse_stream_checker
 from tests.utils.dial_app import create_httpx_client
 from tests.utils.json import has_type, match_objects, memorize
 
-to_many_requests_error = DialExceptionWithHeaders.create(
+to_many_requests_error = _parse_dial_exception(
     status_code=http.HTTPStatus.TOO_MANY_REQUESTS,
     content={"error": {"message": "Too many requests"}},
     headers={"retry-after": "42"},
@@ -52,11 +52,10 @@ async def test_interceptor_errors(stream: bool, repeats: int):
         actual_headers = {
             k.decode(): v.decode() for k, v in response.headers.raw
         }
-        # FIXME: Retry-After should actually be propagated
-        # See https://github.com/epam/ai-dial-sdk/blob/45681f3763679e115d95bc5ce32cf382e0083420/aidial_sdk/_errors.py#L21C1-L26C6
         assert match_objects(
             actual_headers,
             {
+                "retry-after": "42",
                 "content-length": has_type(str),
                 "content-type": "application/json",
             },