fix(streaming): wrap midstream transport errors in APITimeoutError / APIConnectionError (#3811)

src/openai/_streaming.py

@@ -11,7 +11,16 @@
 import httpx2
 
 from ._utils import is_mapping, extract_type_var_from_base
-from ._exceptions import APIError
+from ._httpx2 import timeout_exceptions, _loaded_legacy_httpx
+from ._exceptions import APIError, APITimeoutError, APIConnectionError
+
+
+def _transport_exceptions() -> tuple[type[BaseException], ...]:
+    module = _loaded_legacy_httpx()
+    if module is None:
+        return (httpx2.TransportError,)
+    return (httpx2.TransportError, module.TransportError)
+
 
 if TYPE_CHECKING:
     from ._client import OpenAI, AsyncOpenAI
@@ -59,6 +68,7 @@ def __stream__(self) -> Iterator[_T]:
         process_data = self._client._process_response_data
         iterator = self._iter_events()
 
+        request = response.request
         try:
             for sse in iterator:
                 if sse.data.startswith("[DONE]"):
@@ -106,6 +116,10 @@ def __stream__(self) -> Iterator[_T]:
                         cast_to=cast_to,
                         response=response,
                     )
+        except timeout_exceptions() as err:
+            raise APITimeoutError(request=request) from err
+        except _transport_exceptions() as err:
+            raise APIConnectionError(request=request) from err
         finally:
             # Ensure the response is closed even if the consumer doesn't read all data
             response.close()
@@ -169,6 +183,7 @@ async def __stream__(self) -> AsyncIterator[_T]:
         process_data = self._client._process_response_data
         iterator = self._iter_events()
 
+        request = response.request
         try:
             async for sse in iterator:
                 if sse.data.startswith("[DONE]"):
@@ -216,6 +231,10 @@ async def __stream__(self) -> AsyncIterator[_T]:
                         cast_to=cast_to,
                         response=response,
                     )
+        except timeout_exceptions() as err:
+            raise APITimeoutError(request=request) from err
+        except _transport_exceptions() as err:
+            raise APIConnectionError(request=request) from err
         finally:
             # Ensure the response is closed even if the consumer doesn't read all data
             await response.aclose()

src/openai/lib/streaming/_assistants.py

@@ -11,6 +11,7 @@
 from ..._models import construct_type
 from ..._streaming import Stream, AsyncStream
 from ...types.beta import AssistantStreamEvent
+from ..._exceptions import APITimeoutError, APIConnectionError
 from ...types.beta.threads import (
     Run,
     Text,
@@ -410,6 +411,15 @@ def __stream__(self) -> Iterator[AssistantStreamEvent]:
                 self._emit_sse_event(event)
 
                 yield event
+        except APITimeoutError as exc:
+            cause = exc.__cause__ if isinstance(exc.__cause__, _timeout_exceptions()) else exc
+            self.on_timeout()
+            self.on_exception(cause)
+            raise cause from None
+        except APIConnectionError as exc:
+            cause = exc.__cause__ if isinstance(exc.__cause__, Exception) else exc
+            self.on_exception(cause)
+            raise cause from None
         except _timeout_exceptions() as exc:
             self.on_timeout()
             self.on_exception(exc)
@@ -842,6 +852,15 @@ async def __stream__(self) -> AsyncIterator[AssistantStreamEvent]:
                 await self._emit_sse_event(event)
 
                 yield event
+        except APITimeoutError as exc:
+            cause = exc.__cause__ if isinstance(exc.__cause__, _timeout_exceptions()) else exc
+            await self.on_timeout()
+            await self.on_exception(cause)
+            raise cause from None
+        except APIConnectionError as exc:
+            cause = exc.__cause__ if isinstance(exc.__cause__, Exception) else exc
+            await self.on_exception(cause)
+            raise cause from None
         except _timeout_exceptions() as exc:
             await self.on_timeout()
             await self.on_exception(exc)

tests/test_httpx2.py

@@ -710,3 +710,105 @@ async def body():
     assert "Credential=fixture-access-key/" in async_requests[0].headers["authorization"]
     assert sync_requests[0].headers["x-amz-security-token"] == "fixture-session-token"
     assert async_requests[0].headers["x-amz-security-token"] == "fixture-session-token"
+
+
+class _DiesMidStreamSync(httpx2.SyncByteStream):
+    @override
+    def __iter__(self):  # type: ignore[no-untyped-def]
+        yield b'data: {"id":"c1","object":"chat.completion.chunk","created":0,"model":"x","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},"finish_reason":null}]}\n\n'
+        raise httpx2.ReadTimeout("timed out while reading the stream")
+
+
+class _DiesMidStreamAsync(httpx2.AsyncByteStream):
+    @override
+    async def __aiter__(self):  # type: ignore[no-untyped-def]
+        yield b'data: {"id":"c1","object":"chat.completion.chunk","created":0,"model":"x","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},"finish_reason":null}]}\n\n'
+        raise httpx2.RemoteProtocolError("peer closed connection without sending complete message body")
+
+
+def test_chat_stream_midstream_timeout_wrapped() -> None:
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(
+            200, headers={"content-type": "text/event-stream"}, stream=_DiesMidStreamSync(), request=request
+        )
+
+    with OpenAI(
+        api_key="test",
+        base_url="https://example.test/v1",
+        http_client=httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False),
+        max_retries=0,
+    ) as client:
+        with pytest.raises(APITimeoutError) as exc_info:
+            for _ in client.chat.completions.create(
+                model="x", messages=[{"role": "user", "content": "hi"}], stream=True
+            ):
+                pass
+        assert isinstance(exc_info.value.__cause__, httpx2.ReadTimeout)
+
+
+async def test_chat_stream_midstream_connection_error_wrapped() -> None:
+    async def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(
+            200, headers={"content-type": "text/event-stream"}, stream=_DiesMidStreamAsync(), request=request
+        )
+
+    async with AsyncOpenAI(
+        api_key="test",
+        base_url="https://example.test/v1",
+        http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False),
+        max_retries=0,
+    ) as client:
+        with pytest.raises(APIConnectionError) as exc_info:
+            async for _ in await client.chat.completions.create(
+                model="x", messages=[{"role": "user", "content": "hi"}], stream=True
+            ):
+                pass
+        assert isinstance(exc_info.value.__cause__, httpx2.RemoteProtocolError)
+
+
+def test_chat_stream_midstream_connection_error_wrapped_sync() -> None:
+    class _DiesSync(httpx2.SyncByteStream):
+        @override
+        def __iter__(self):  # type: ignore[no-untyped-def]
+            yield b'data: {"id":"c1","object":"chat.completion.chunk","created":0,"model":"x","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},"finish_reason":null}]}\n\n'
+            raise httpx2.RemoteProtocolError("peer closed connection without sending complete message body")
+
+    def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(200, headers={"content-type": "text/event-stream"}, stream=_DiesSync(), request=request)
+
+    with OpenAI(
+        api_key="test",
+        base_url="https://example.test/v1",
+        http_client=httpx2.Client(transport=httpx2.MockTransport(handler), trust_env=False),
+        max_retries=0,
+    ) as client:
+        with pytest.raises(APIConnectionError) as exc_info:
+            for _ in client.chat.completions.create(
+                model="x", messages=[{"role": "user", "content": "hi"}], stream=True
+            ):
+                pass
+        assert isinstance(exc_info.value.__cause__, httpx2.RemoteProtocolError)
+
+
+async def test_chat_stream_midstream_timeout_wrapped_async() -> None:
+    class _DiesAsync(httpx2.AsyncByteStream):
+        @override
+        async def __aiter__(self):  # type: ignore[no-untyped-def]
+            yield b'data: {"id":"c1","object":"chat.completion.chunk","created":0,"model":"x","choices":[{"index":0,"delta":{"role":"assistant","content":"hi"},"finish_reason":null}]}\n\n'
+            raise httpx2.ReadTimeout("timed out while reading the stream")
+
+    async def handler(request: httpx2.Request) -> httpx2.Response:
+        return httpx2.Response(200, headers={"content-type": "text/event-stream"}, stream=_DiesAsync(), request=request)
+
+    async with AsyncOpenAI(
+        api_key="test",
+        base_url="https://example.test/v1",
+        http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(handler), trust_env=False),
+        max_retries=0,
+    ) as client:
+        with pytest.raises(APITimeoutError) as exc_info:
+            async for _ in await client.chat.completions.create(
+                model="x", messages=[{"role": "user", "content": "hi"}], stream=True
+            ):
+                pass
+        assert isinstance(exc_info.value.__cause__, httpx2.ReadTimeout)