diff --git a/src/openai/_streaming.py b/src/openai/_streaming.py index 78e2d20aa7..e8a119c66c 100644 --- a/src/openai/_streaming.py +++ b/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() diff --git a/src/openai/lib/streaming/_assistants.py b/src/openai/lib/streaming/_assistants.py index 314961230d..9ec40f7c21 100644 --- a/src/openai/lib/streaming/_assistants.py +++ b/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) diff --git a/tests/test_httpx2.py b/tests/test_httpx2.py index 764fc00f0b..e4a245f604 100644 --- a/tests/test_httpx2.py +++ b/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)