Skip to content

Commit b8f52dc

Browse files
Merge b8826d8 into c5452af
2 parents c5452af + b8826d8 commit b8f52dc

3 files changed

Lines changed: 203 additions & 1 deletion

File tree

‎src/openai/_streaming.py‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
from ._utils import is_mapping, extract_type_var_from_base
1414
from ._httpx2 import request_exceptions, timeout_exceptions
1515
from ._exceptions import APIError, APITimeoutError, APIConnectionError
16+
from ._streaming_cleanup import drain_response_async
1617

1718
if TYPE_CHECKING:
1819
from ._client import OpenAI, AsyncOpenAI
@@ -183,6 +184,11 @@ async def __stream__(self) -> AsyncIterator[_T]:
183184
try:
184185
async for sse in iterator:
185186
if sse.data.startswith("[DONE]"):
187+
# HTTP/1.1 needs the remaining body consumed for connection reuse.
188+
# HTTP/2 can close this stream and keep the connection, so draining
189+
# it would add unnecessary latency after [DONE].
190+
if response.http_version == "HTTP/1.1":
191+
await drain_response_async(response)
186192
break
187193

188194
# we have to special case the Assistants `thread.` events since we won't have an "event" key in the data

‎src/openai/_streaming_cleanup.py‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
from __future__ import annotations
2+
3+
from typing import AsyncIterator, cast
4+
5+
import anyio
6+
import httpx2
7+
from anyio.lowlevel import checkpoint
8+
9+
from ._httpx2 import request_exceptions, stream_consumed_exceptions
10+
11+
_DRAIN_TIMEOUT = 0.1
12+
13+
14+
async def drain_response_async(response: httpx2.Response) -> None:
15+
if response.is_closed:
16+
return
17+
try:
18+
with anyio.move_on_after(_DRAIN_TIMEOUT):
19+
# Discard raw transport bytes without decoding trailing SSE or content.
20+
async for _ in cast(AsyncIterator[bytes], response.stream):
21+
# Also enforce the deadline on transports yielding buffered data.
22+
await checkpoint()
23+
except request_exceptions() + stream_consumed_exceptions():
24+
pass # Cleanup must not turn completed output into a request error.
25+
finally:
26+
# Close outside our timeout scope and shield cleanup from caller
27+
# cancellation. Cancellation still propagates after the response closes.
28+
with anyio.CancelScope(shield=True):
29+
await response.aclose()

‎tests/test_streaming.py‎

Lines changed: 168 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,21 @@
22

33
import os
44
import ssl
5+
import time
6+
import socket
7+
import asyncio
58
import importlib
9+
import threading
610
from typing import Any, Iterator, AsyncIterator
711
from contextlib import aclosing, nullcontext
12+
from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
13+
from typing_extensions import override
814

915
import anyio
1016
import httpx2
1117
import pytest
1218

13-
from openai import OpenAI, AsyncOpenAI, APITimeoutError, APIConnectionError
19+
from openai import OpenAI, AsyncOpenAI, APITimeoutError, APIConnectionError, DefaultAioHttpClient
1420
from openai._streaming import Stream, AsyncStream, ServerSentEvent
1521

1622

@@ -29,6 +35,167 @@ def http_module(request: pytest.FixtureRequest) -> Any:
2935
return importlib.import_module(request.param)
3036

3137

38+
@pytest.mark.parametrize("mode", ["async", "aiohttp"])
39+
@pytest.mark.parametrize("ending", ["complete", "truncate", "stall", "trickle"])
40+
async def test_done_http11_connection_reuse(mode: str, ending: str, http_module: Any) -> None:
41+
if mode == "aiohttp" and http_module.__name__ == "httpx":
42+
pytest.skip("aiohttp uses HTTPX2")
43+
connections = 0
44+
requests = 0
45+
release = threading.Event()
46+
47+
class Handler(BaseHTTPRequestHandler):
48+
protocol_version = "HTTP/1.1"
49+
50+
@override
51+
def setup(self) -> None:
52+
nonlocal connections
53+
super().setup()
54+
connections += 1
55+
self.connection.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
56+
57+
def do_POST(self) -> None:
58+
nonlocal requests
59+
self.rfile.read(int(self.headers["Content-Length"]))
60+
requests += 1
61+
self.send_response(200)
62+
self.send_header("Content-Type", "text/event-stream")
63+
self.send_header("Transfer-Encoding", "chunked")
64+
self.end_headers()
65+
payload = (
66+
b'data: {"id":"synthetic","object":"chat.completion.chunk","created":0,"model":"synthetic",'
67+
b'"choices":[{"index":0,"delta":{"content":"hello"}}]}\n\ndata: [DONE]\n\n'
68+
)
69+
try:
70+
self.wfile.write(f"{len(payload):x}\r\n".encode() + payload + b"\r\n")
71+
if requests == 1:
72+
if ending == "truncate":
73+
self.close_connection = True
74+
return
75+
if ending == "stall":
76+
release.wait(timeout=0.6)
77+
if ending == "trickle":
78+
# Trailer bytes do not yield body chunks, but must still
79+
# obey the total cleanup deadline, not a per-read timeout.
80+
self.wfile.write(b"0\r\nX-Synthetic: ")
81+
for _ in range(15):
82+
if release.wait(timeout=0.04):
83+
return
84+
self.wfile.write(b"a")
85+
self.wfile.write(b"\r\n\r\n")
86+
return
87+
self.wfile.write(b"0\r\n\r\n")
88+
except ConnectionError:
89+
self.close_connection = True
90+
91+
@override
92+
def log_message(self, format: str, *args: object) -> None: # noqa: A002
93+
pass
94+
95+
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
96+
thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01})
97+
thread.start()
98+
options: dict[str, Any] = {
99+
"api_key": "synthetic",
100+
"base_url": f"http://127.0.0.1:{server.server_port}/v1",
101+
"max_retries": 0,
102+
"timeout": http_module.Timeout(5, read=None),
103+
}
104+
transport_options: dict[str, Any] = {"trust_env": False, "limits": http_module.Limits(max_connections=1)}
105+
try:
106+
http_client = (
107+
DefaultAioHttpClient(**transport_options)
108+
if mode == "aiohttp"
109+
else http_module.AsyncClient(**transport_options)
110+
)
111+
async with AsyncOpenAI(**options, http_client=http_client) as async_client:
112+
for _ in range(2):
113+
async with await async_client.chat.completions.create(
114+
model="synthetic", messages=[], stream=True
115+
) as async_stream:
116+
assert (await async_stream.__anext__()).choices[0].delta.content == "hello"
117+
start = time.monotonic()
118+
assert [chunk async for chunk in async_stream] == []
119+
assert time.monotonic() - start < 0.4
120+
finally:
121+
release.set()
122+
server.shutdown()
123+
thread.join(timeout=5)
124+
server.server_close()
125+
126+
assert not thread.is_alive()
127+
assert requests == 2
128+
# A failed drain must discard the connection and release the only pool slot.
129+
assert connections == (1 if ending == "complete" else 2)
130+
131+
132+
@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
133+
@pytest.mark.parametrize("http_version", [b"HTTP/1.1", b"HTTP/2"])
134+
async def test_done_discards_bytes_only_for_async_http11(
135+
sync: bool, http_version: bytes, http_module: Any, client: OpenAI, async_client: AsyncOpenAI
136+
) -> None:
137+
read_tail = False
138+
139+
def body() -> Iterator[bytes]:
140+
nonlocal read_tail
141+
yield b"data: [DONE]\n\n"
142+
read_tail = True
143+
yield b"\xff\n\n" # Trailing data must not be decoded as SSE.
144+
145+
iterator = body()
146+
async_iterator = to_aiter(iterator)
147+
148+
class Body(http_module.SyncByteStream, http_module.AsyncByteStream):
149+
def __iter__(self) -> Iterator[bytes]:
150+
return iterator
151+
152+
def __aiter__(self) -> AsyncIterator[bytes]:
153+
return async_iterator
154+
155+
response = http_module.Response(200, stream=Body(), extensions={"http_version": http_version})
156+
if sync:
157+
assert list(Stream(cast_to=object, response=response, client=client)) == []
158+
else:
159+
assert [event async for event in AsyncStream(cast_to=object, response=response, client=async_client)] == []
160+
assert response.is_closed
161+
assert read_tail == (not sync and http_version == b"HTTP/1.1")
162+
163+
164+
@pytest.mark.parametrize("cancel", [True, False], ids=["caller-cancel", "deadline"])
165+
async def test_cancellation_during_done_drain_closes_response(
166+
cancel: bool, http_module: Any, async_client: AsyncOpenAI
167+
) -> None:
168+
draining = asyncio.Event()
169+
170+
class Body(http_module.AsyncByteStream):
171+
closed = False
172+
173+
async def __aiter__(self) -> AsyncIterator[bytes]:
174+
yield b"data: [DONE]\n\n"
175+
draining.set()
176+
await asyncio.Event().wait()
177+
178+
async def aclose(self) -> None:
179+
await anyio.sleep(0)
180+
self.closed = True
181+
182+
body = Body()
183+
response = http_module.Response(200, stream=body)
184+
stream = AsyncStream(cast_to=object, response=response, client=async_client)
185+
task = asyncio.create_task(stream.__anext__())
186+
try:
187+
await asyncio.wait_for(draining.wait(), timeout=5)
188+
if cancel:
189+
task.cancel()
190+
with pytest.raises(asyncio.CancelledError if cancel else StopAsyncIteration):
191+
await task
192+
assert response.is_closed
193+
assert body.closed
194+
finally:
195+
task.cancel()
196+
await asyncio.gather(task, return_exceptions=True)
197+
198+
32199
@pytest.mark.parametrize("sync", [True, False], ids=["sync", "async"])
33200
@pytest.mark.parametrize("delivered", [False, True], ids=["before-first-event", "after-first-event"])
34201
@pytest.mark.parametrize(

0 commit comments

Comments
 (0)