fix(streaming): handle top-level SSE error messages

Author: raashish1601

src/openai/_streaming.py

@@ -20,6 +20,35 @@
 _T = TypeVar("_T")
 
 
+def _build_streaming_api_error(
+    *,
+    data: object,
+    request: httpx.Request,
+    is_error_event: bool,
+) -> APIError | None:
+    if not is_mapping(data):
+        return None
+
+    error = data.get("error")
+    if is_error_event:
+        body = error if error is not None else data
+    elif error is not None:
+        body = error
+    else:
+        return None
+
+    message = data.get("message")
+    if not isinstance(message, str) and is_mapping(error):
+        nested_message = error.get("message")
+        if isinstance(nested_message, str):
+            message = nested_message
+
+    if not isinstance(message, str) or not message:
+        message = "An error occurred during streaming"
+
+    return APIError(message=message, request=request, body=body)
+
+
 class Stream(Generic[_T]):
     """Provides the core interface to iterate over a synchronous stream response."""
 
@@ -63,41 +92,19 @@ def __stream__(self) -> Iterator[_T]:
                 if sse.data.startswith("[DONE]"):
                     break
 
+                data = sse.json()
+                api_error = _build_streaming_api_error(
+                    data=data,
+                    request=self.response.request,
+                    is_error_event=sse.event == "error",
+                )
+                if api_error is not None:
+                    raise api_error
+
                 # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
                 if sse.event and sse.event.startswith("thread."):
-                    data = sse.json()
-
-                    if sse.event == "error" and is_mapping(data) and data.get("error"):
-                        message = None
-                        error = data.get("error")
-                        if is_mapping(error):
-                            message = error.get("message")
-                        if not message or not isinstance(message, str):
-                            message = "An error occurred during streaming"
-
-                        raise APIError(
-                            message=message,
-                            request=self.response.request,
-                            body=data["error"],
-                        )
-
                     yield process_data(data={"data": data, "event": sse.event}, cast_to=cast_to, response=response)
                 else:
-                    data = sse.json()
-                    if is_mapping(data) and data.get("error"):
-                        message = None
-                        error = data.get("error")
-                        if is_mapping(error):
-                            message = error.get("message")
-                        if not message or not isinstance(message, str):
-                            message = "An error occurred during streaming"
-
-                        raise APIError(
-                            message=message,
-                            request=self.response.request,
-                            body=data["error"],
-                        )
-
                     yield process_data(
                         data={"data": data, "event": sse.event}
                         if self._options is not None and self._options.synthesize_event_and_data
@@ -173,41 +180,19 @@ async def __stream__(self) -> AsyncIterator[_T]:
                 if sse.data.startswith("[DONE]"):
                     break
 
+                data = sse.json()
+                api_error = _build_streaming_api_error(
+                    data=data,
+                    request=self.response.request,
+                    is_error_event=sse.event == "error",
+                )
+                if api_error is not None:
+                    raise api_error
+
                 # we have to special case the Assistants `thread.` events since we won't have an "event" key in the data
                 if sse.event and sse.event.startswith("thread."):
-                    data = sse.json()
-
-                    if sse.event == "error" and is_mapping(data) and data.get("error"):
-                        message = None
-                        error = data.get("error")
-                        if is_mapping(error):
-                            message = error.get("message")
-                        if not message or not isinstance(message, str):
-                            message = "An error occurred during streaming"
-
-                        raise APIError(
-                            message=message,
-                            request=self.response.request,
-                            body=data["error"],
-                        )
-
                     yield process_data(data={"data": data, "event": sse.event}, cast_to=cast_to, response=response)
                 else:
-                    data = sse.json()
-                    if is_mapping(data) and data.get("error"):
-                        message = None
-                        error = data.get("error")
-                        if is_mapping(error):
-                            message = error.get("message")
-                        if not message or not isinstance(message, str):
-                            message = "An error occurred during streaming"
-
-                        raise APIError(
-                            message=message,
-                            request=self.response.request,
-                            body=data["error"],
-                        )
-
                     yield process_data(
                         data={"data": data, "event": sse.event}
                         if self._options is not None and self._options.synthesize_event_and_data

tests/test_streaming.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
-from typing import Iterator, AsyncIterator
+from typing import TypeVar, Iterator, AsyncIterator
 
 import httpx
 import pytest
 
-from openai import OpenAI, AsyncOpenAI
+from openai import OpenAI, APIError, AsyncOpenAI
 from openai._streaming import Stream, AsyncStream, ServerSentEvent
 
+_ItemT = TypeVar("_ItemT")
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
@@ -216,23 +218,78 @@ def body() -> Iterator[bytes]:
     assert sse.json() == {"content": "известни"}
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_error_event_uses_top_level_message(sync: bool, client: OpenAI, async_client: AsyncOpenAI) -> None:
+    def body() -> Iterator[bytes]:
+        yield b"event: error\n"
+        yield b'data: {"type":"error","code":"server_error","message":"Something went wrong"}\n'
+        yield b"\n"
+
+    stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client)
+
+    with pytest.raises(APIError, match="Something went wrong") as exc_info:
+        await iter_next(stream)
+
+    assert exc_info.value.body == {"type": "error", "code": "server_error", "message": "Something went wrong"}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+async def test_error_event_keeps_nested_error_message_fallback(
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> None:
+    def body() -> Iterator[bytes]:
+        yield b"event: error\n"
+        yield b'data: {"error":{"type":"error","code":"nested_error","message":"Nested failure"}}\n'
+        yield b"\n"
+
+    stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client)
+
+    with pytest.raises(APIError, match="Nested failure") as exc_info:
+        await iter_next(stream)
+
+    assert exc_info.value.body == {"type": "error", "code": "nested_error", "message": "Nested failure"}
+
+
 async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]:
     for chunk in iter:
         yield chunk
 
 
-async def iter_next(iter: Iterator[ServerSentEvent] | AsyncIterator[ServerSentEvent]) -> ServerSentEvent:
+async def iter_next(iter: Iterator[_ItemT] | AsyncIterator[_ItemT]) -> _ItemT:
     if isinstance(iter, AsyncIterator):
         return await iter.__anext__()
 
     return next(iter)
 
 
-async def assert_empty_iter(iter: Iterator[ServerSentEvent] | AsyncIterator[ServerSentEvent]) -> None:
+async def assert_empty_iter(iter: Iterator[object] | AsyncIterator[object]) -> None:
     with pytest.raises((StopAsyncIteration, RuntimeError)):
         await iter_next(iter)
 
 
+def make_stream(
+    content: Iterator[bytes],
+    *,
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> Stream[object] | AsyncStream[object]:
+    request = httpx.Request("GET", "https://example.com/stream")
+
+    if sync:
+        return Stream(cast_to=object, client=client, response=httpx.Response(200, content=content, request=request))
+
+    return AsyncStream(
+        cast_to=object,
+        client=async_client,
+        response=httpx.Response(200, content=to_aiter(content), request=request),
+    )
+
+
 def make_event_iterator(
     content: Iterator[bytes],
     *,