fix: wrap transport failures while consuming a stream

src/openai/_httpx2.py

@@ -17,6 +17,7 @@ class _LegacyHttpxModule(Protocol):
     Timeout: type[httpx2.Timeout]
     Limits: type[httpx2.Limits]
     TimeoutException: type[httpx2.TimeoutException]
+    TransportError: type[httpx2.TransportError]
     HTTPStatusError: type[httpx2.HTTPStatusError]
     StreamConsumed: type[httpx2.StreamConsumed]
     RequestNotRead: type[httpx2.RequestNotRead]
@@ -99,6 +100,11 @@ def timeout_exceptions() -> tuple[type[httpx2.TimeoutException], ...]:
     return (httpx2.TimeoutException,) if module is None else (httpx2.TimeoutException, module.TimeoutException)
 
 
+def transport_exceptions() -> tuple[type[httpx2.TransportError], ...]:
+    module = _loaded_legacy_httpx()
+    return (httpx2.TransportError,) if module is None else (httpx2.TransportError, module.TransportError)
+
+
 def status_exceptions() -> tuple[type[httpx2.HTTPStatusError], ...]:
     module = _loaded_legacy_httpx()
     return (httpx2.HTTPStatusError,) if module is None else (httpx2.HTTPStatusError, module.HTTPStatusError)

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, transport_exceptions
+from ._exceptions import APIError, APITimeoutError, APIConnectionError
 
 if TYPE_CHECKING:
     from ._client import OpenAI, AsyncOpenAI
@@ -106,6 +107,10 @@ def __stream__(self) -> Iterator[_T]:
                         cast_to=cast_to,
                         response=response,
                     )
+        except timeout_exceptions() as err:
+            raise APITimeoutError(request=response.request) from err
+        except transport_exceptions() 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 +221,10 @@ 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 transport_exceptions() 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()

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
 from ...types.beta.threads import (
     Run,
     Text,
@@ -25,7 +26,7 @@
 
 
 def _timeout_exceptions() -> tuple[type[Exception], ...]:
-    return (*timeout_exceptions(), asyncio.TimeoutError)
+    return (*timeout_exceptions(), asyncio.TimeoutError, APITimeoutError)
 
 
 class AssistantEventHandler:

tests/test_httpx2.py

@@ -638,11 +638,13 @@ async def async_response(request: httpx2.Request) -> httpx2.Response:
         with sync_client.beta.threads.runs.stream(  # pyright: ignore[reportDeprecated]
             assistant_id="asst_test", thread_id="thread_test", event_handler=sync_handler
         ) as stream:
-            with pytest.raises(httpx2.ReadTimeout, match="assistant stream timeout"):
+            with pytest.raises(APITimeoutError) as sync_exc_info:
                 stream.until_done()
 
+    assert isinstance(sync_exc_info.value.__cause__, httpx2.ReadTimeout)
     assert sync_handler.timed_out
-    assert isinstance(sync_handler.exception, httpx2.ReadTimeout)
+    assert isinstance(sync_handler.exception, APITimeoutError)
+    assert isinstance(sync_handler.exception.__cause__, httpx2.ReadTimeout)
 
     async_handler = AsyncHandler()
     async with AsyncOpenAI(
@@ -654,11 +656,13 @@ async def async_response(request: httpx2.Request) -> httpx2.Response:
         async with async_client.beta.threads.runs.stream(  # pyright: ignore[reportDeprecated]
             assistant_id="asst_test", thread_id="thread_test", event_handler=async_handler
         ) as async_stream:
-            with pytest.raises(httpx2.ReadTimeout, match="assistant stream timeout"):
+            with pytest.raises(APITimeoutError) as async_exc_info:
                 await async_stream.until_done()
 
+    assert isinstance(async_exc_info.value.__cause__, httpx2.ReadTimeout)
     assert async_handler.timed_out
-    assert isinstance(async_handler.exception, httpx2.ReadTimeout)
+    assert isinstance(async_handler.exception, APITimeoutError)
+    assert isinstance(async_handler.exception.__cause__, httpx2.ReadTimeout)
 
 
 async def test_sigv4_provider_preserves_httpx2_family_and_rejects_one_shot_bodies() -> None:

tests/test_streaming.py

@@ -5,7 +5,7 @@
 import httpx2
 import pytest
 
-from openai import OpenAI, AsyncOpenAI
+from openai import OpenAI, APIError, AsyncOpenAI, APITimeoutError, APIConnectionError
 from openai._streaming import Stream, AsyncStream, ServerSentEvent
 
 
@@ -216,6 +216,35 @@ def body() -> Iterator[bytes]:
     assert sse.json() == {"content": "известни"}
 
 
+@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
+@pytest.mark.parametrize(
+    "failure,expected",
+    [
+        (httpx2.ReadTimeout("timed out while reading the stream"), APITimeoutError),
+        (httpx2.RemoteProtocolError("peer closed connection"), APIConnectionError),
+    ],
+    ids=["timeout", "connection"],
+)
+async def test_transport_error_mid_stream(
+    sync: bool,
+    failure: httpx2.TransportError,
+    expected: type[APIError],
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> None:
+    def body() -> Iterator[bytes]:
+        yield b'data: {"foo":true}\n'
+        yield b"\n"
+        raise failure
+
+    stream = make_stream(content=body(), sync=sync, client=client, async_client=async_client)
+
+    with pytest.raises(expected) as exc_info:
+        await consume_stream(stream)
+
+    assert exc_info.value.__cause__ is failure
+
+
 async def to_aiter(iter: Iterator[bytes]) -> AsyncIterator[bytes]:
     for chunk in iter:
         yield chunk
@@ -246,3 +275,30 @@ def make_event_iterator(
     return AsyncStream(
         cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(content))
     )._iter_events()
+
+
+def make_stream(
+    content: Iterator[bytes],
+    *,
+    sync: bool,
+    client: OpenAI,
+    async_client: AsyncOpenAI,
+) -> Stream[object] | AsyncStream[object]:
+    request = httpx2.Request("POST", "https://example.test/v1/chat/completions")
+
+    if sync:
+        return Stream(cast_to=object, client=client, response=httpx2.Response(200, content=content, request=request))
+
+    return AsyncStream(
+        cast_to=object, client=async_client, response=httpx2.Response(200, content=to_aiter(content), request=request)
+    )
+
+
+async def consume_stream(stream: Stream[object] | AsyncStream[object]) -> None:
+    if isinstance(stream, AsyncStream):
+        async for _ in stream:
+            pass
+        return
+
+    for _ in stream:
+        pass