fix: improve http error handling (#320)

src/genai/_utils/http_client/httpx_client.py

@@ -48,10 +48,11 @@ def post_stream(
                 response: Response = event_source.response
                 if "application/json" in response.headers["content-type"]:
                     response.read()
-                    raise ApiResponseException(  # noqa: B904
+                    raise ApiResponseException.from_http_response(
                         message="Invalid data chunk retrieved during streaming.", response=response
-                    )
-                raise e
+                    ) from None
+                else:
+                    raise e
 
 
 class AsyncHttpxClient(httpx.AsyncClient):

src/genai/_utils/http_client/retry_transport.py

@@ -96,7 +96,7 @@ def _create_exception(
                 else f"Failed to handle request to {request.url}."
             )
 
-            return ApiResponseException(
+            return ApiResponseException.from_http_response(
                 message=message,
                 response=exception.response,
             )

src/genai/_utils/responses.py

@@ -1,7 +1,6 @@
 from typing import Any, Optional
 
 import httpx
-from httpx import Response
 from typing_extensions import TypeGuard
 
 from genai.schema import (
@@ -13,7 +12,12 @@
     UnauthorizedResponse,
 )
 
-__all__ = ["is_api_error_response", "get_api_error_class_by_status_code", "to_api_error", "BaseErrorResponse"]
+__all__ = [
+    "is_api_error_response",
+    "get_api_error_class_by_status_code",
+    "to_api_error",
+    "BaseErrorResponse",
+]
 
 
 def is_api_error_response(input: Any) -> TypeGuard[BaseErrorResponse]:
@@ -31,10 +35,7 @@ def get_api_error_class_by_status_code(code: int) -> Optional[type[BaseErrorResp
     return response_class_mapping.get(code)
 
 
-def to_api_error(response: Response) -> BaseErrorResponse:
-    if response.is_success:
-        raise ValueError("Cannot convert succeed HTTP response to error response.")
-
-    cls: type[BaseErrorResponse] = get_api_error_class_by_status_code(response.status_code) or BaseErrorResponse
-    model = cls.model_validate(response.json())
+def to_api_error(body: dict) -> BaseErrorResponse:
+    cls: type[BaseErrorResponse] = get_api_error_class_by_status_code(body["status_code"]) or BaseErrorResponse
+    model = cls.model_validate(body)
     return model

src/genai/exceptions.py

@@ -1,6 +1,5 @@
 import logging
-import typing
-from typing import Union
+from typing import Optional, Union
 
 from httpx import HTTPError, Response
 from pydantic import ValidationError as _ValidationError
@@ -34,19 +33,30 @@ class ApiResponseException(BaseApiException):
 
     def __init__(
         self,
-        response: Union[Response, BaseErrorResponse],
-        message: str = "Server Error",
+        response: Union[BaseErrorResponse, dict],
+        message: Optional[str] = None,
         *args,
     ) -> None:
         if is_api_error_response(response):
             self.response = response
-        elif isinstance(response, Response):
+        elif isinstance(response, dict):
             self.response = to_api_error(response)
         else:
             raise TypeError(f"Expected either Response or Api Error Response, but {type(response)} received.")
 
-        message = f"{message}\n{self.response.model_dump_json(indent=2)}"
-        super().__init__(message, *args)
+        self.message = f"{message or 'Server Error'}\n{self.response.model_dump_json(indent=2)}"
+        super().__init__(self.message, *args)
+
+    @classmethod
+    def from_http_response(cls, response: Response, message: Optional[str] = None):
+        if response.is_success:
+            raise ValueError("Cannot convert succeed HTTP response to error response.")
+
+        response_body = to_api_error(response.json())
+        return cls(message=message, response=response_body)
+
+    def __reduce__(self):
+        return self.__class__, (self.response.model_dump(), self.message)
 
 
 class ApiNetworkException(BaseApiException):
@@ -57,8 +67,8 @@ class ApiNetworkException(BaseApiException):
         message (str): Explanation of the error.
     """
 
-    __cause__: typing.Optional[HTTPError] = None
+    __cause__: Optional[HTTPError] = None
 
-    def __init__(self, message: typing.Optional[str] = None, *args) -> None:
+    def __init__(self, message: Optional[str] = None, *args) -> None:
         self.message = message or "Network Exception has occurred. Try again later."
         super().__init__(self.message, *args)

src/genai/text/embedding/embedding_service.py

@@ -125,8 +125,8 @@ def create(
         async def handler(input: str, http_client: AsyncClient, limiter: BaseLimiter) -> TextEmbeddingCreateResponse:
             self._log_method_execution("Embedding Create - processing input", input=input)
 
-            async def handle_retry(ex: HTTPStatusError):
-                if ex.response.status_code == httpx.codes.TOO_MANY_REQUESTS:
+            async def handle_retry(ex: Exception):
+                if isinstance(ex, HTTPStatusError) and ex.response.status_code == httpx.codes.TOO_MANY_REQUESTS:
                     await limiter.report_error()
 
             async def handle_success(*args):

src/genai/text/generation/generation_service.py

@@ -172,8 +172,8 @@ def create(
         async def handler(input: str, http_client: AsyncClient, limiter: BaseLimiter) -> TextGenerationCreateResponse:
             self._log_method_execution("Generate Create - processing input", input=input)
 
-            async def handle_retry(ex: HTTPStatusError):
-                if ex.response.status_code == httpx.codes.TOO_MANY_REQUESTS:
+            async def handle_retry(ex: Exception):
+                if isinstance(ex, HTTPStatusError) and ex.response.status_code == httpx.codes.TOO_MANY_REQUESTS:
                     await limiter.report_error()
 
             async def handle_success(*args):