feat: improve error handling

Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>

src/genai/exceptions/genai_exception.py

@@ -10,8 +10,11 @@
 
 
 class GenAiException(Exception):
-    def __init__(self, error: Union[Exception, Response, str]) -> None:
-        if isinstance(error, Response):
+    def __init__(self, error: Union[Exception, Response, str, ErrorResponse]) -> None:
+        if isinstance(error, ErrorResponse):
+            self.error = error
+            self.error_message = self.error.message
+        elif isinstance(error, Response):
             try:
                 self.error = ErrorResponse(**error.json())
                 self.error_message = self.error.message

src/genai/schemas/responses.py

@@ -3,7 +3,7 @@
 from enum import Enum
 from typing import Any, List, Optional, Type, Union
 
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic import AliasChoices, BaseModel, ConfigDict, Field, model_validator
 
 from genai.schemas.generate_params import GenerateParams
 
@@ -179,12 +179,13 @@ class ErrorExtensionState(GenAiResponseModel):
 
 class ErrorExtensions(GenAiResponseModel):
     code: str
-    state: list[ErrorExtensionState]
+    state: Optional[list[ErrorExtensionState]] = None
+    reason: Optional[str] = None
 
 
 class ErrorResponse(GenAiResponseModel):
-    status_code: int
-    error: str
+    status_code: int = Field(validation_alias=AliasChoices("status_code", "statusCode"))
+    error: str = ""
     message: str
     extensions: Optional[ErrorExtensions] = None
 

src/genai/services/request_handler.py

@@ -1,3 +1,4 @@
+import json
 import logging
 from typing import Optional
 
@@ -7,6 +8,7 @@
 from genai._version import version
 from genai.exceptions import GenAiException
 from genai.options import Options
+from genai.schemas.responses import ErrorResponse
 from genai.services.connection_manager import ConnectionManager
 from genai.utils.http_provider import HttpProvider
 
@@ -297,6 +299,14 @@ def post_stream(endpoint, headers, json_data, files):
             ) as event_source:
                 try:
                     for sse in event_source.iter_sse():
+                        if sse.event == "error":
+                            if sse.data.startswith("{") and sse.data.endswith("}"):
+                                raise GenAiException(ErrorResponse(**json.loads(sse.data)))
+
+                            raise GenAiException(
+                                f"Invalid server response during streaming!\nRetrieved data: {sse.data}"
+                            )
+
                         yield sse.data
                 except SSEError as e:
                     response: Response = event_source.response