fix: Correctly handle error tokens in chat and deepnews streams

Author: wizrds

asknews_sdk/api/chat.py

@@ -9,6 +9,7 @@
     CreateChatCompletionRequest,
     CreateChatCompletionResponse,
     CreateChatCompletionResponseStream,
+    CreateChatCompletionResponseStreamError,
     FilterParamsResponse,
     ForecastResponse,
     HeadlineQuestionsResponse,
@@ -21,8 +22,10 @@
     CreateDeepNewsResponse,
     CreateDeepNewsResponseStream,
     CreateDeepNewsResponseStreamChunk,
+    CreateDeepNewsResponseStreamError,
     CreateDeepNewsResponseStreamSource,
 )
+from asknews_sdk.errors import APIError
 from asknews_sdk.response import EventSource
 
 
@@ -120,7 +123,23 @@ def _stream():
                 for event in EventSource.from_api_response(response):
                     if event.content == "[DONE]":
                         break
-                    yield CreateChatCompletionResponseStream.model_validate_json(event.content)
+
+                    token = (
+                        TypeAdapter(Union[
+                            CreateChatCompletionResponseStreamError,
+                            CreateChatCompletionResponseStream
+                        ])
+                        .validate_json(event.content)
+                    )
+
+                    if isinstance(token, CreateChatCompletionResponseStreamError):
+                        raise APIError(
+                            response=response,
+                            detail=token.error.message,
+                            code=token.error.code,
+                        )
+
+                    yield token
 
             return _stream()
         else:
@@ -524,7 +543,22 @@ def _stream():
                     if event.content == "[DONE]":
                         break
 
-                    yield TypeAdapter(CreateDeepNewsResponseStream).validate_json(event.content)
+                    token = (
+                        TypeAdapter(Union[
+                            CreateDeepNewsResponseStreamError,
+                            CreateDeepNewsResponseStream
+                        ])
+                        .validate_json(event.content)
+                    )
+
+                    if isinstance(token, CreateDeepNewsResponseStreamError):
+                        raise APIError(
+                            response=response,
+                            detail=token.error.message,
+                            code=token.error.code,
+                        )
+
+                    yield token
 
             return _stream()
         else:
@@ -625,7 +659,23 @@ async def _stream():
                 async for event in EventSource.from_api_response(response):
                     if event.content == "[DONE]":
                         break
-                    yield CreateChatCompletionResponseStream.model_validate_json(event.content)
+
+                    token = (
+                        TypeAdapter(Union[
+                            CreateChatCompletionResponseStreamError,
+                            CreateChatCompletionResponseStream
+                        ])
+                        .validate_json(event.content)
+                    )
+
+                    if isinstance(token, CreateChatCompletionResponseStreamError):
+                        raise APIError(
+                            response=response,
+                            detail=token.error.message,
+                            code=token.error.code,
+                        )
+
+                    yield token
 
             return _stream()
         else:
@@ -1022,13 +1072,27 @@ async def get_deep_news(
         )
 
         if stream:
-
             async def _stream():
                 async for event in EventSource.from_api_response(response):
                     if event.content == "[DONE]":
                         break
 
-                    yield TypeAdapter(CreateDeepNewsResponseStream).validate_json(event.content)
+                    token = (
+                        TypeAdapter(Union[
+                            CreateDeepNewsResponseStreamError,
+                            CreateDeepNewsResponseStream
+                        ])
+                        .validate_json(event.content)
+                    )
+
+                    if isinstance(token, CreateDeepNewsResponseStreamError):
+                        raise APIError(
+                            response=response,
+                            detail=token.error.message,
+                            code=token.error.code,
+                        )
+
+                    yield token
 
             return _stream()
         else:

asknews_sdk/dto/chat.py

@@ -97,6 +97,17 @@ class CreateChatCompletionResponseStream(BaseSchema):
     choices: Annotated[List[CreateChatCompletionResponseStreamChoice], Field(title="Choices")]
 
 
+class CreateChatCompletionResponseStreamErrorDetails(BaseModel):
+    message: Annotated[str, Field(title="Message")]
+    code: Annotated[int, Field(title="Code")]
+
+
+class CreateChatCompletionResponseStreamError(BaseSchema):
+    __content_type__ = "text/event-stream"
+
+    error: Annotated[CreateChatCompletionResponseStreamErrorDetails, Field(title="Error")]
+
+
 class ModelItem(BaseModel):
     id: Annotated[str, Field(title="Id")]
     object: Annotated[Optional[str], Field("model", title="Object")]

asknews_sdk/dto/deepnews.py

@@ -194,6 +194,17 @@ class CreateDeepNewsResponseStreamSource(BaseSchema):
     ]
 
 
+class CreateDeepNewsResponseStreamErrorDetails(BaseModel):
+    code: Annotated[int, Field(title="Code")]
+    message: Annotated[str, Field(title="Message")]
+
+
+class CreateDeepNewsResponseStreamError(BaseSchema):
+    __content_type__ = "text/event-stream"
+
+    error: Annotated[CreateDeepNewsResponseStreamErrorDetails, Field(title="Error")]
+
+
 CreateDeepNewsResponseStream: TypeAlias = Annotated[
     Union[
         Annotated[CreateDeepNewsResponseStreamChunk, Tag("chat.completion.chunk")],

asknews_sdk/errors.py

@@ -49,6 +49,11 @@ class MethodNotAllowed(APIError):
     detail = "Method Not Allowed"
 
 
+class RequestTimeoutError(APIError):
+    code = 408000
+    detail = "Request Timeout"
+
+
 class RateLimitExceededError(APIError):
     code = 429000
     detail = "Rate Limit Exceeded"
@@ -87,6 +92,7 @@ class ServiceUnavailableError(APIError):
     403011: ForbiddenError,
     404000: ResourceNotFoundError,
     405000: MethodNotAllowed,
+    408000: RequestTimeoutError,
     422000: ValidationError,
     429000: RateLimitExceededError,
     429001: ConcurrencyLimitExceededError,