diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 78e2d20aa7..cc499c318a 100644 --- a/src/openai/_streaming.py +++ b/src/openai/_streaming.py @@ -11,7 +11,8 @@ import httpx2 from ._utils import is_mapping, extract_type_var_from_base -from ._exceptions import APIError +from ._httpx2 import timeout_exceptions +from ._exceptions import APIError, OpenAIError, APITimeoutError, APIConnectionError if TYPE_CHECKING: from ._client import OpenAI, AsyncOpenAI @@ -106,6 +107,12 @@ def __stream__(self) -> Iterator[_T]: cast_to=cast_to, response=response, ) + except timeout_exceptions() as err: + raise APITimeoutError(request=response.request) from err + except OpenAIError: + raise + except Exception as err: + raise APIConnectionError(request=response.request) from err finally: # Ensure the response is closed even if the consumer doesn't read all data response.close() @@ -216,6 +223,12 @@ async def __stream__(self) -> AsyncIterator[_T]: cast_to=cast_to, response=response, ) + except timeout_exceptions() as err: + raise APITimeoutError(request=response.request) from err + except OpenAIError: + raise + except Exception as err: + raise APIConnectionError(request=response.request) from err finally: # Ensure the response is closed even if the consumer doesn't read all data await response.aclose() diff --git a/tests/test_streaming.py b/tests/test_streaming.py index ae6c0590f7..df409de617 100644 --- a/tests/test_streaming.py +++ b/tests/test_streaming.py @@ -5,9 +5,60 @@ import httpx2 import pytest -from openai import OpenAI, AsyncOpenAI +from openai import OpenAI, AsyncOpenAI, APITimeoutError from openai._streaming import Stream, AsyncStream, ServerSentEvent +FIRST_CHUNK = ( + b'data: {"id":"c1","object":"chat.completion.chunk","created":0,"model":"gpt-5.2",' + b'"choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},"finish_reason":null}]}\n\n' +) + + +class _DiesMidStream(httpx2.SyncByteStream): + def __iter__(self) -> Iterator[bytes]: + yield FIRST_CHUNK + raise httpx2.ReadTimeout("timed out while reading the stream") + + +class _AsyncDiesMidStream(httpx2.AsyncByteStream): + async def __aiter__(self) -> AsyncIterator[bytes]: + yield FIRST_CHUNK + raise httpx2.ReadTimeout("timed out while reading the stream") + + +def test_sync_stream_wraps_mid_stream_transport_error() -> None: + def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, headers={"content-type": "text/event-stream"}, stream=_DiesMidStream()) + + client = OpenAI( + api_key="My API Key", + http_client=httpx2.Client(transport=httpx2.MockTransport(handler)), + max_retries=0, + ) + + with pytest.raises(APITimeoutError): + for _ in client.chat.completions.create( + model="gpt-5.2", messages=[{"role": "user", "content": "hi"}], stream=True + ): + pass + + +async def test_async_stream_wraps_mid_stream_transport_error() -> None: + async def handler(request: httpx2.Request) -> httpx2.Response: + return httpx2.Response(200, headers={"content-type": "text/event-stream"}, stream=_AsyncDiesMidStream()) + + client = AsyncOpenAI( + api_key="My API Key", + http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler)), + max_retries=0, + ) + + with pytest.raises(APITimeoutError): + async for _ in await client.chat.completions.create( + model="gpt-5.2", messages=[{"role": "user", "content": "hi"}], stream=True + ): + pass + @pytest.mark.asyncio @pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])