mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-15 04:19:08 +00:00
fix(streaming): add CustomStreamWrapper.close() and release the connection when a sync stream ends (#46354)
* fix(streaming): close sync streams so the connection goes back to the pool CustomStreamWrapper gets a sync close() (plus context manager support) that closes the wrapped provider stream, and both the sync and async wrappers close the provider stream when the stream reaches its end. The Anthropic stream iterator keeps the httpx response it reads so close()/aclose() can release the socket after message_stop instead of leaving it checked out of the pool Co-authored-by: Blair Jordan <blair.j.jordan@gmail.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(streaming): keep the basedpyright ceiling flat in CustomStreamWrapper.close() Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(streaming): close the response on the shared sync stream path too BaseModelResponseIterator gets a sync close() next to its aclose(), and the sync branch of BaseLLMHTTPHandler hands it the httpx response the way the async branch already does, so providers built on it release the socket on close() instead of waiting for garbage collection Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(streaming): unit tests for CustomStreamWrapper.close() and the shared iterator close() Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(streaming): close the provider stream before re-raising a mid-stream error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: Blair Jordan <blair.j.jordan@gmail.com>
This commit is contained in:
Kerry Lu
GitHub
parent
6e0f0f7adb
commit
6bd5580425
8 changed files with 347 additions and 2 deletions
|
|
@ -1,6 +1,7 @@
|
|||
import asyncio
|
||||
import collections.abc
|
||||
import datetime
|
||||
import inspect
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
|
|
@ -234,6 +235,7 @@ class CustomStreamWrapper:
|
|||
self.completion_stream = completion_stream
|
||||
self.sent_first_chunk = False
|
||||
self.sent_last_chunk = False
|
||||
self._closed = False
|
||||
self._stream_created_time: float = time.time()
|
||||
|
||||
litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate(
|
||||
|
|
@ -328,6 +330,18 @@ class CustomStreamWrapper:
|
|||
def __aiter__(self) -> AsyncIterator["ModelResponseStream"]:
|
||||
return self
|
||||
|
||||
def __enter__(self) -> "CustomStreamWrapper":
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc_info: object) -> None:
|
||||
self.close()
|
||||
|
||||
async def __aenter__(self) -> "CustomStreamWrapper":
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc_info: object) -> None:
|
||||
await self.aclose()
|
||||
|
||||
def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None:
|
||||
"""Restore trace_id/session_id in the *consuming* thread/task/context.
|
||||
|
||||
|
|
@ -392,11 +406,26 @@ class CustomStreamWrapper:
|
|||
"""
|
||||
self._restore_consumer_correlation_context(guarded=True)
|
||||
|
||||
def close(self) -> None:
|
||||
stream_to_close: Final = cast(object, self.completion_stream) # cast-ok: completion_stream is untyped
|
||||
self.completion_stream = None
|
||||
self._closed = True
|
||||
close_stream: Final = getattr(stream_to_close, "close", None)
|
||||
if callable(close_stream):
|
||||
try:
|
||||
result: Final = close_stream()
|
||||
if inspect.iscoroutine(result):
|
||||
result.close()
|
||||
except Exception as e: # noqa: BLE001 # best-effort cleanup; must not raise into the caller
|
||||
verbose_logger.debug("CustomStreamWrapper.close: error closing completion_stream: %s", e)
|
||||
self._restore_consumer_correlation_context()
|
||||
|
||||
async def aclose(self):
|
||||
# Restore the consumer's outer context only after the underlying
|
||||
# provider stream's own close (and its diagnostic logging below, if
|
||||
# closing fails) completes - not before - so those log lines still
|
||||
# carry this closing stream's own trace_id/session_id.
|
||||
self._closed = True
|
||||
if self.completion_stream is not None:
|
||||
stream_to_close: Final = self.completion_stream
|
||||
self.completion_stream = None
|
||||
|
|
@ -1780,6 +1809,8 @@ class CustomStreamWrapper:
|
|||
response.hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost
|
||||
|
||||
def __next__(self) -> "ModelResponseStream":
|
||||
if self._closed:
|
||||
raise StopIteration
|
||||
cache_hit = False
|
||||
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
|
||||
cache_hit = True
|
||||
|
|
@ -1950,7 +1981,7 @@ class CustomStreamWrapper:
|
|||
if self.sent_stream_usage is False and self.send_stream_usage is True:
|
||||
self.sent_stream_usage = True
|
||||
return response
|
||||
self._restore_consumer_correlation_context()
|
||||
self.close()
|
||||
raise # Re-raise StopIteration
|
||||
else:
|
||||
self.sent_last_chunk = True
|
||||
|
|
@ -1986,6 +2017,7 @@ class CustomStreamWrapper:
|
|||
traceback_exception: Final = traceback.format_exc()
|
||||
# LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated
|
||||
threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start()
|
||||
self.close()
|
||||
self._handle_stream_fallback_error(e)
|
||||
|
||||
def fetch_sync_stream(self):
|
||||
|
|
@ -2005,6 +2037,8 @@ class CustomStreamWrapper:
|
|||
return self.completion_stream
|
||||
|
||||
async def __anext__(self) -> "ModelResponseStream":
|
||||
if self._closed:
|
||||
raise StopAsyncIteration
|
||||
cache_hit = False
|
||||
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
|
||||
cache_hit = True
|
||||
|
|
@ -2125,12 +2159,15 @@ class CustomStreamWrapper:
|
|||
asyncio.create_task(
|
||||
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
|
||||
)
|
||||
await self.aclose()
|
||||
self._handle_stream_fallback_error(e)
|
||||
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
|
||||
if self.received_finish_reason is None:
|
||||
await self.aclose()
|
||||
self._log_stream_failure_and_raise(e)
|
||||
return await self._finalize_completed_stream(cache_hit=cache_hit)
|
||||
except Exception as e:
|
||||
await self.aclose()
|
||||
self._log_stream_failure_and_raise(e)
|
||||
|
||||
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
|
||||
|
|
@ -2224,7 +2261,7 @@ class CustomStreamWrapper:
|
|||
)
|
||||
)
|
||||
|
||||
self._restore_consumer_correlation_context()
|
||||
await self.aclose()
|
||||
raise StopAsyncIteration # Re-raise StopIteration
|
||||
else:
|
||||
self.sent_last_chunk = True
|
||||
|
|
|
|||
|
|
@ -232,6 +232,7 @@ async def make_call(
|
|||
speed=speed,
|
||||
tool_name_reverse_map=tool_name_reverse_map,
|
||||
response_format_schema=response_format_schema,
|
||||
response=response,
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
|
|
@ -302,6 +303,7 @@ def make_sync_call(
|
|||
speed=speed,
|
||||
tool_name_reverse_map=tool_name_reverse_map,
|
||||
response_format_schema=response_format_schema,
|
||||
response=response,
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
|
|
@ -675,9 +677,12 @@ class ModelResponseIterator:
|
|||
speed: str | None = None,
|
||||
tool_name_reverse_map: dict[str, str] | None = None,
|
||||
response_format_schema: ResponseFormatSchema | None = None,
|
||||
response: httpx.Response | None = None,
|
||||
):
|
||||
self.streaming_response = streaming_response
|
||||
self.response_iterator = self.streaming_response
|
||||
self.sync_stream = sync_stream
|
||||
self._response = response
|
||||
self.content_blocks: list[ContentBlockDelta] = []
|
||||
self.tool_index = -1
|
||||
self.json_mode = json_mode
|
||||
|
|
@ -1375,6 +1380,22 @@ class ModelResponseIterator:
|
|||
return self._handle_accumulated_json_chunk(data_str)
|
||||
|
||||
# Sync iterator
|
||||
def close(self) -> None:
|
||||
response: Final = self._response
|
||||
self._response = None
|
||||
if response is not None:
|
||||
response.close()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
response: Final = self._response
|
||||
self._response = None
|
||||
if response is None:
|
||||
return
|
||||
if self.sync_stream:
|
||||
response.close()
|
||||
return
|
||||
await response.aclose()
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
|
|
|
|||
|
|
@ -81,6 +81,10 @@ class BaseModelResponseIterator:
|
|||
if self.http_response is not None:
|
||||
await self.http_response.aclose()
|
||||
|
||||
def close(self) -> None:
|
||||
if self.http_response is not None:
|
||||
self.http_response.close()
|
||||
|
||||
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream:
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
|
|
|
|||
|
|
@ -1008,6 +1008,8 @@ class BaseLLMHTTPHandler:
|
|||
json_mode=json_mode,
|
||||
optional_params=_OPTIONAL_PARAMS.validate_python(optional_params),
|
||||
)
|
||||
if isinstance(completion_stream, BaseModelResponseIterator):
|
||||
completion_stream.http_response = response
|
||||
|
||||
# LOGGING
|
||||
logging_obj.post_call(
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ class Wire:
|
|||
received: SimpleQueue[Request]
|
||||
disconnected: SimpleQueue[str]
|
||||
connected: SimpleQueue[str]
|
||||
closed: SimpleQueue[str]
|
||||
|
||||
def drain(self) -> tuple[Request, ...]:
|
||||
return tuple(self.received.get_nowait() for _ in range(self.received.qsize()))
|
||||
|
|
@ -77,6 +78,9 @@ class Wire:
|
|||
def connections(self) -> int:
|
||||
return self.connected.qsize()
|
||||
|
||||
def open_connections(self) -> int:
|
||||
return self.connected.qsize() - self.closed.qsize()
|
||||
|
||||
|
||||
def _await_release(gate: threading.Event, closing: threading.Event) -> bool:
|
||||
while not closing.is_set():
|
||||
|
|
@ -98,6 +102,7 @@ def wire_server(
|
|||
errors: Final[SimpleQueue[Exception]] = SimpleQueue()
|
||||
disconnected: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
connected: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
closed: Final[SimpleQueue[str]] = SimpleQueue()
|
||||
closing: Final = threading.Event()
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
|
|
@ -108,6 +113,10 @@ def wire_server(
|
|||
super().setup()
|
||||
connected.put(f"{self.client_address[0]}:{self.client_address[1]}")
|
||||
|
||||
def finish(self) -> None:
|
||||
super().finish()
|
||||
closed.put(f"{self.client_address[0]}:{self.client_address[1]}")
|
||||
|
||||
def respond(self) -> None:
|
||||
headers: Final = {name.lower(): value for name, value in self.headers.items()}
|
||||
try:
|
||||
|
|
@ -194,6 +203,7 @@ def wire_server(
|
|||
received,
|
||||
disconnected,
|
||||
connected,
|
||||
closed,
|
||||
)
|
||||
finally:
|
||||
closing.set()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,214 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.client import eventually
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
import litellm
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
_API_KEY: Final = "synthetic-stream-release-key"
|
||||
_PROMPT: Final = "does the socket go back to the pool once the stream is done with it?"
|
||||
_FIRST_WORDS: Final = "hello "
|
||||
_LAST_WORD: Final = "world"
|
||||
_RELEASE_SECONDS: Final = 2
|
||||
|
||||
|
||||
def _anthropic_event(event: str, payload: Mapping[str, object]) -> bytes:
|
||||
return f"event: {event}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
def _openai_delta(text: str, finish_reason: str | None) -> bytes:
|
||||
chunk: Final = {
|
||||
"id": "chatcmpl-stream-release-wire",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": "gpt-5.6",
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": finish_reason}],
|
||||
}
|
||||
return b"data: " + json.dumps(chunk).encode() + b"\n\n"
|
||||
|
||||
|
||||
_ANTHROPIC_OPENING: Final = b"".join(
|
||||
(
|
||||
_anthropic_event(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_stream_release_wire",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-5-5",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 12, "output_tokens": 1},
|
||||
},
|
||||
},
|
||||
),
|
||||
_anthropic_event(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
_anthropic_event(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _FIRST_WORDS}},
|
||||
),
|
||||
)
|
||||
)
|
||||
_ANTHROPIC_CLOSING: Final = (
|
||||
_anthropic_event(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": _LAST_WORD}},
|
||||
),
|
||||
_anthropic_event("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
_anthropic_event(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 3},
|
||||
},
|
||||
),
|
||||
_anthropic_event("message_stop", {"type": "message_stop"}),
|
||||
)
|
||||
_OPENAI_OPENING: Final = _openai_delta(_FIRST_WORDS, None)
|
||||
_OPENAI_CLOSING: Final = (_openai_delta(_LAST_WORD, None), _openai_delta("", "stop"), b"data: [DONE]\n\n")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Shape:
|
||||
name: str
|
||||
model: str
|
||||
opening: bytes
|
||||
closing: tuple[bytes, ...]
|
||||
|
||||
|
||||
_SHAPES: Final = (
|
||||
_Shape("anthropic", "anthropic/claude-sonnet-5-5", _ANTHROPIC_OPENING, _ANTHROPIC_CLOSING),
|
||||
_Shape("openai", "openai/gpt-5.6", _OPENAI_OPENING, _OPENAI_CLOSING),
|
||||
_Shape("databricks", "databricks/databricks-llama", _OPENAI_OPENING, _OPENAI_CLOSING),
|
||||
)
|
||||
|
||||
|
||||
def _peer(shape: _Shape, gate_after_first: threading.Event | None = None):
|
||||
def respond(request: Request) -> Reply:
|
||||
return Reply(
|
||||
content_type="text/event-stream", chunks=(shape.opening, *shape.closing), gate_after_first=gate_after_first
|
||||
)
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _request_kwargs(shape: _Shape, wire: Wire) -> Mapping[str, object]:
|
||||
client: Final = openai.OpenAI(api_key=_API_KEY, base_url=wire.url) if shape.name == "openai" else None
|
||||
return {
|
||||
"model": shape.model,
|
||||
"api_base": wire.url,
|
||||
"api_key": _API_KEY,
|
||||
"messages": ({"role": "user", "content": _PROMPT},),
|
||||
"stream": True,
|
||||
**({"client": client} if client is not None else {}),
|
||||
}
|
||||
|
||||
|
||||
def _text_of(chunk: ModelResponseStream) -> str:
|
||||
return chunk.choices[0].delta.content or ""
|
||||
|
||||
|
||||
def _open_sessions_settling_to(wire: Wire, at_most: int) -> int:
|
||||
return eventually(wire.open_connections, lambda open_sessions: open_sessions <= at_most, seconds=_RELEASE_SECONDS)
|
||||
|
||||
|
||||
async def _open_sessions_settling_to_while_the_loop_runs(wire: Wire, at_most: int) -> int:
|
||||
deadline: Final = time.monotonic() + _RELEASE_SECONDS
|
||||
while wire.open_connections() > at_most and time.monotonic() < deadline:
|
||||
await asyncio.sleep(0.05)
|
||||
return wire.open_connections()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", _SHAPES, ids=lambda shape: shape.name)
|
||||
def test_sync_stream_read_to_its_last_chunk_releases_the_connection(shape: _Shape) -> None:
|
||||
with wire_server(_peer(shape), keep_alive=True) as wire:
|
||||
request: Final = _request_kwargs(shape, wire)
|
||||
first_text: Final = "".join(_text_of(chunk) for chunk in litellm.completion(**request))
|
||||
second_text: Final = "".join(_text_of(chunk) for chunk in litellm.completion(**request))
|
||||
assert (first_text, second_text) == (_FIRST_WORDS + _LAST_WORD, _FIRST_WORDS + _LAST_WORD)
|
||||
assert _open_sessions_settling_to(wire, at_most=1) <= 1, "a finished stream left its TCP session open"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", _SHAPES, ids=lambda shape: shape.name)
|
||||
def test_sync_stream_closed_after_its_first_chunk_releases_the_connection(shape: _Shape) -> None:
|
||||
gate: Final = threading.Event()
|
||||
with wire_server(_peer(shape, gate_after_first=gate), keep_alive=True) as wire:
|
||||
request: Final = _request_kwargs(shape, wire)
|
||||
abandoned: Final = litellm.completion(**request)
|
||||
first_chunk: Final = next(iter(abandoned))
|
||||
assert _text_of(first_chunk) == _FIRST_WORDS
|
||||
abandoned.close()
|
||||
gate.set()
|
||||
assert _open_sessions_settling_to(wire, at_most=0) == 0, "close() left the TCP session open"
|
||||
gate.clear()
|
||||
follow_up: Final = litellm.completion(**request)
|
||||
assert _text_of(next(iter(follow_up))) == _FIRST_WORDS
|
||||
gate.set()
|
||||
assert "".join(_text_of(chunk) for chunk in follow_up) == _LAST_WORD
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_read_to_its_last_chunk_releases_the_connection() -> None:
|
||||
shape: Final = _SHAPES[0]
|
||||
with wire_server(_peer(shape), keep_alive=True) as wire:
|
||||
request: Final = _request_kwargs(shape, wire)
|
||||
first_text: Final = "".join([_text_of(chunk) async for chunk in await litellm.acompletion(**request)])
|
||||
second_text: Final = "".join([_text_of(chunk) async for chunk in await litellm.acompletion(**request)])
|
||||
assert (first_text, second_text) == (_FIRST_WORDS + _LAST_WORD, _FIRST_WORDS + _LAST_WORD)
|
||||
assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=1) <= 1, (
|
||||
"a finished stream left its TCP session open"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_stream_aclosed_after_its_first_chunk_releases_the_connection() -> None:
|
||||
shape: Final = _SHAPES[0]
|
||||
gate: Final = threading.Event()
|
||||
with wire_server(_peer(shape, gate_after_first=gate), keep_alive=True) as wire:
|
||||
request: Final = _request_kwargs(shape, wire)
|
||||
abandoned: Final = await litellm.acompletion(**request)
|
||||
assert _text_of(await abandoned.__anext__()) == _FIRST_WORDS
|
||||
await abandoned.aclose()
|
||||
gate.set()
|
||||
assert await _open_sessions_settling_to_while_the_loop_runs(wire, at_most=0) == 0, (
|
||||
"aclose() left the TCP session open"
|
||||
)
|
||||
gate.clear()
|
||||
follow_up: Final = await litellm.acompletion(**request)
|
||||
assert _text_of(await follow_up.__anext__()) == _FIRST_WORDS
|
||||
gate.set()
|
||||
assert "".join([_text_of(chunk) async for chunk in follow_up]) == _LAST_WORD
|
||||
|
||||
|
||||
_ANTHROPIC_ERROR_MIDDLE: Final = _anthropic_event(
|
||||
"error", {"type": "error", "error": {"type": "overloaded_error", "message": "peer overloaded mid stream"}}
|
||||
)
|
||||
|
||||
|
||||
def test_sync_stream_that_raises_mid_stream_releases_the_connection() -> None:
|
||||
shape: Final = _SHAPES[0]
|
||||
broken: Final = _Shape(shape.name, shape.model, shape.opening, (_ANTHROPIC_ERROR_MIDDLE, *shape.closing))
|
||||
with wire_server(_peer(broken), keep_alive=True) as wire:
|
||||
request: Final = _request_kwargs(broken, wire)
|
||||
with pytest.raises(Exception, match="overloaded"):
|
||||
"".join(_text_of(chunk) for chunk in litellm.completion(**request))
|
||||
assert _open_sessions_settling_to(wire, at_most=0) == 0, (
|
||||
"a stream that raised mid-way left its TCP session open"
|
||||
)
|
||||
|
|
@ -6415,3 +6415,51 @@ def validate_last_format(chunk):
|
|||
assert isinstance(
|
||||
choice["finish_reason"], str
|
||||
), "'finish_reason' should be a string."
|
||||
|
||||
|
||||
class _RecordingProviderStream:
|
||||
def __init__(self) -> None:
|
||||
self.close_calls = 0
|
||||
|
||||
def __iter__(self) -> "_RecordingProviderStream":
|
||||
return self
|
||||
|
||||
def __next__(self) -> object:
|
||||
raise StopIteration
|
||||
|
||||
def close(self) -> None:
|
||||
self.close_calls += 1
|
||||
|
||||
|
||||
def _wrapper_over(stream: _RecordingProviderStream) -> CustomStreamWrapper:
|
||||
return CustomStreamWrapper(
|
||||
completion_stream=stream,
|
||||
model="gpt-5.6",
|
||||
logging_obj=MagicMock(),
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
|
||||
def test_close_releases_the_provider_stream_once_and_ends_iteration():
|
||||
stream: Final = _RecordingProviderStream()
|
||||
wrapper: Final = _wrapper_over(stream)
|
||||
wrapper.close()
|
||||
wrapper.close()
|
||||
assert stream.close_calls == 1
|
||||
with pytest.raises(StopIteration):
|
||||
next(wrapper)
|
||||
|
||||
|
||||
def test_leaving_a_with_block_closes_the_provider_stream():
|
||||
stream: Final = _RecordingProviderStream()
|
||||
with _wrapper_over(stream):
|
||||
pass
|
||||
assert stream.close_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_leaving_an_async_with_block_closes_the_provider_stream():
|
||||
stream: Final = _RecordingProviderStream()
|
||||
async with _wrapper_over(stream):
|
||||
pass
|
||||
assert stream.close_calls == 1
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import json
|
|||
from collections.abc import AsyncIterator
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -305,3 +306,11 @@ async def test_bytes_chunk_with_text_before_data_prefix_is_parsed(sync_mode: boo
|
|||
assert chunks[0].id == "chatcmpl-bytes"
|
||||
assert chunks[0].choices[0].delta.content == "Hello"
|
||||
assert chunks[1]["is_finished"] is True
|
||||
|
||||
|
||||
def test_close_releases_the_attached_http_response():
|
||||
response: Final = httpx.Response(200, stream=httpx.ByteStream(b"data: [DONE]\n\n"))
|
||||
iterator: Final = BaseModelResponseIterator(streaming_response=iter(()), sync_stream=True)
|
||||
iterator.http_response = response
|
||||
iterator.close()
|
||||
assert response.is_closed
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue