mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
fix(anthropic_messages): bound streaming relay queue and cap detached drains
The relay queue was unbounded, so a client reading a long stream more slowly than Bedrock produced it let the pump accumulate every pending SSE chunk in memory, and detached post-disconnect drains had no concurrency bound, so an authenticated client could open many large streams and read slowly to pin unbounded worker state. Bound the relay queue and make the pump apply backpressure while the client is connected (it blocks on a full queue, racing the disconnect signal), so a slow reader throttles the upstream read exactly as the old direct yield did. Cap how many detached drains run at once; over the cap a disconnected pump bills what it collected instead of draining further. Detached-drain lifetime is otherwise bounded by the upstream stream/read timeout. Both limits are tunable via env.
This commit is contained in:
nuernber
parent
ce25486702
commit
739447fa4b
3 changed files with 309 additions and 61 deletions
|
|
@ -445,6 +445,18 @@ FIREWORKS_AI_80_B: Final = int(os.getenv("FIREWORKS_AI_80_B", 80))
|
|||
#### Logging callback constants ####
|
||||
REDACTED_BY_LITELM_STRING: Final = "REDACTED_BY_LITELM"
|
||||
MAX_LANGFUSE_INITIALIZED_CLIENTS: Final = int(os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50))
|
||||
# Backpressure + lifetime bounds for the /v1/messages streaming relay (see
|
||||
# BaseAnthropicMessagesStreamingIterator.async_sse_wrapper). The relay queue is
|
||||
# bounded so a slow client throttles the upstream pump instead of letting it
|
||||
# buffer the whole response in memory; the detached-drain cap bounds how many
|
||||
# post-disconnect drains may run concurrently so client behavior can't create
|
||||
# unbounded worker state.
|
||||
ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE: Final = int(
|
||||
os.getenv("ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", "1024")
|
||||
)
|
||||
ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS: Final = int(
|
||||
os.getenv("ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", "100")
|
||||
)
|
||||
LOGGING_WORKER_CONCURRENCY: Final = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0
|
||||
LOGGING_WORKER_MAX_QUEUE_SIZE: Final = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000))
|
||||
LOGGING_WORKER_MAX_TIME_PER_COROUTINE: Final = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0))
|
||||
|
|
|
|||
|
|
@ -8,6 +8,10 @@ import httpx
|
|||
from pydantic import TypeAdapter
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from litellm.constants import (
|
||||
ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS,
|
||||
ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE,
|
||||
)
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
|
|
@ -25,6 +29,13 @@ GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging()
|
|||
# guidance and drop it again from the done callback.
|
||||
_UPSTREAM_PUMP_TASKS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: stdlib strong-ref set for pump tasks
|
||||
|
||||
# Rooted set of pumps still draining upstream AFTER their client disconnected.
|
||||
# Bounds how many detached drains run at once so a burst of slow/abandoned
|
||||
# streams can't pin unbounded worker memory; a pump over the cap bills what it
|
||||
# already collected instead of continuing to drain. Only ever touched from the
|
||||
# event loop, so a plain set + len() check needs no lock.
|
||||
_DETACHED_STREAM_DRAINS: Final[set[asyncio.Task[None]]] = set() # mutable-ok: bounded strong-ref set, detached drains
|
||||
|
||||
INCOMPLETE_STREAM_ERROR_MESSAGE: Final = (
|
||||
"Provider stream ended before emitting a message_stop event; "
|
||||
"the response is incomplete and any partial content (e.g. tool_use input JSON) may be truncated."
|
||||
|
|
@ -51,6 +62,23 @@ def _is_terminal_stream_chunk(chunk: object) -> bool:
|
|||
return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk)
|
||||
|
||||
|
||||
def _try_claim_detached_drain_slot() -> bool:
|
||||
"""Claim a detached-drain slot for the current task, bounding concurrency.
|
||||
|
||||
Returns True if a slot was claimed (the caller may keep draining upstream
|
||||
for billing) or False if the cap is already reached (the caller should stop
|
||||
and bill what it has). Only touched from the event loop, so the check +
|
||||
insert need no lock.
|
||||
"""
|
||||
if len(_DETACHED_STREAM_DRAINS) >= ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS:
|
||||
return False
|
||||
current_task: Final = asyncio.current_task()
|
||||
if current_task is not None:
|
||||
_DETACHED_STREAM_DRAINS.add(current_task)
|
||||
current_task.add_done_callback(_DETACHED_STREAM_DRAINS.discard)
|
||||
return True
|
||||
|
||||
|
||||
def _incomplete_stream_error_sse_event() -> bytes:
|
||||
payload: Final = json.dumps(
|
||||
{
|
||||
|
|
@ -205,9 +233,18 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
generating and billing the full response regardless of the client, so
|
||||
draining it to completion is what lets spend tracking see the real
|
||||
terminal ``message_delta`` / ``message_stop`` usage instead of a
|
||||
truncated placeholder count. Chunks reach the client through a queue;
|
||||
once the client goes away the pump only buffers for billing so the queue
|
||||
can't grow unbounded.
|
||||
truncated placeholder count.
|
||||
|
||||
Chunks reach the client through a bounded queue. While the client is
|
||||
connected the pump blocks on a full queue (racing the disconnect
|
||||
signal), so a slow reader throttles the upstream read exactly as the old
|
||||
direct ``yield`` did instead of letting the whole response buffer in
|
||||
memory. Once the client goes away the pump stops enqueueing and only
|
||||
keeps a single ``collected_chunks`` copy for billing, and the number of
|
||||
such post-disconnect drains running at once is capped so client behavior
|
||||
can't create unbounded worker state; over the cap the pump bills what it
|
||||
has rather than draining further. Detached-drain lifetime is otherwise
|
||||
bounded by the upstream stream/read timeout.
|
||||
|
||||
An upstream failure (Bedrock read / decode / chunk-conversion error)
|
||||
that happens while the client is still connected is forwarded through
|
||||
|
|
@ -217,63 +254,12 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
|
||||
This method provides the common logic for both Anthropic and Bedrock implementations.
|
||||
"""
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue()
|
||||
queue: Final[asyncio.Queue[bytes | None | BaseException]] = asyncio.Queue(
|
||||
maxsize=ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE
|
||||
)
|
||||
client_detached: Final = asyncio.Event()
|
||||
|
||||
async def _pump_upstream() -> None:
|
||||
collected_chunks: Final[list[bytes]] = []
|
||||
saw_terminal_event = False # rebind-ok: accumulates across the upstream loop
|
||||
|
||||
async def _bill() -> None:
|
||||
try:
|
||||
await self._handle_streaming_logging(collected_chunks)
|
||||
except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper billing failed after %d chunks: %s(%s)",
|
||||
len(collected_chunks),
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
|
||||
try:
|
||||
async for chunk in completion_stream:
|
||||
if self.completion_start_time is None:
|
||||
self.completion_start_time = datetime.now()
|
||||
saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk)
|
||||
encoded_chunk = self._convert_chunk_to_sse_format(chunk)
|
||||
collected_chunks.append(encoded_chunk)
|
||||
if not client_detached.is_set():
|
||||
queue.put_nowait(encoded_chunk)
|
||||
except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend
|
||||
if not client_detached.is_set():
|
||||
# Preserve the provider-specific failure: hand the original
|
||||
# exception to the client-facing generator so it re-raises
|
||||
# and the proxy's failure handling (status code,
|
||||
# post_call_failure_hook) runs. The failure path owns
|
||||
# logging here, so don't also success-bill.
|
||||
queue.put_nowait(exc)
|
||||
return
|
||||
# Client already disconnected: nothing to propagate to and no
|
||||
# failure hook will run, so salvage the partial spend instead
|
||||
# of dropping the request entirely.
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)",
|
||||
len(collected_chunks),
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
await _bill()
|
||||
return
|
||||
|
||||
if not client_detached.is_set():
|
||||
if not saw_terminal_event:
|
||||
queue.put_nowait(_incomplete_stream_error_sse_event())
|
||||
queue.put_nowait(None)
|
||||
await _bill()
|
||||
|
||||
pump_task: Final = asyncio.create_task(_pump_upstream())
|
||||
pump_task: Final = asyncio.create_task(self._pump_upstream_to_queue(completion_stream, queue, client_detached))
|
||||
_UPSTREAM_PUMP_TASKS.add(pump_task)
|
||||
pump_task.add_done_callback(_UPSTREAM_PUMP_TASKS.discard)
|
||||
|
||||
|
|
@ -288,7 +274,130 @@ class BaseAnthropicMessagesStreamingIterator:
|
|||
finally:
|
||||
# Client-facing generator is being torn down (normal end, a
|
||||
# re-raised upstream error, or a disconnect GeneratorExit). Signal
|
||||
# the pump to stop enqueueing so the queue can't grow unbounded, but
|
||||
# let it keep draining upstream to its terminal usage event for
|
||||
# accurate billing.
|
||||
# the pump to stop enqueueing and unblock any backpressure-blocked
|
||||
# put; the pump then either finishes billing or drains detached
|
||||
# (subject to the cap) for accurate usage.
|
||||
client_detached.set()
|
||||
|
||||
async def _bill_collected_chunks(
|
||||
self,
|
||||
collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _handle_streaming_logging
|
||||
) -> None:
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
try:
|
||||
await self._handle_streaming_logging(collected_chunks)
|
||||
except Exception as exc: # noqa: BLE001 # billing is best-effort; never crash the pump
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper billing failed after %d chunks: %s(%s)",
|
||||
len(collected_chunks),
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _enqueue_for_client(
|
||||
queue: "asyncio.Queue[bytes | None | BaseException]",
|
||||
client_detached: "asyncio.Event",
|
||||
item: bytes | None | BaseException,
|
||||
) -> bool:
|
||||
"""Deliver one item to the client, applying backpressure.
|
||||
|
||||
Returns True if the item was queued, False if the client disconnected
|
||||
before there was room (the item is then dropped, since a gone client
|
||||
can't receive it). Never blocks once the client has detached.
|
||||
"""
|
||||
if client_detached.is_set():
|
||||
return False
|
||||
try:
|
||||
queue.put_nowait(item)
|
||||
return True
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
put_task: Final = asyncio.ensure_future(queue.put(item))
|
||||
detached_task: Final = asyncio.ensure_future(client_detached.wait())
|
||||
try:
|
||||
await asyncio.wait(frozenset((put_task, detached_task)), return_when=asyncio.FIRST_COMPLETED)
|
||||
finally:
|
||||
if not detached_task.done():
|
||||
detached_task.cancel()
|
||||
if put_task.done() and not put_task.cancelled():
|
||||
return True
|
||||
put_task.cancel()
|
||||
return False
|
||||
|
||||
async def _pump_upstream_to_queue(
|
||||
self,
|
||||
completion_stream: AsyncIterator[bytes | GenericStreamingChunk | ModelResponseStream | dict],
|
||||
queue: "asyncio.Queue[bytes | None | BaseException]",
|
||||
client_detached: "asyncio.Event",
|
||||
) -> None:
|
||||
"""Drain the whole upstream into ``queue`` (backpressured) and bill once.
|
||||
|
||||
Runs detached so a client disconnect can't interrupt the upstream read;
|
||||
see ``async_sse_wrapper`` for the full rationale. Returns after billing.
|
||||
"""
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
collected_chunks: Final[list[bytes]] = [] # mutable-ok: SSE billing buffer appended to across the drain
|
||||
saw_terminal_event = False # rebind-ok: accumulates across the upstream loop
|
||||
draining_detached = False # rebind-ok: set once this pump claims a detached-drain slot
|
||||
try:
|
||||
async for chunk in completion_stream:
|
||||
if self.completion_start_time is None:
|
||||
self.completion_start_time = datetime.now()
|
||||
saw_terminal_event = saw_terminal_event or _is_terminal_stream_chunk(chunk)
|
||||
encoded_chunk = self._convert_chunk_to_sse_format(chunk)
|
||||
collected_chunks.append(encoded_chunk)
|
||||
if not client_detached.is_set():
|
||||
await self._enqueue_for_client(queue, client_detached, encoded_chunk)
|
||||
continue
|
||||
# Client has gone: keep draining only to reach the terminal usage
|
||||
# event for billing, but claim a detached-drain slot first; over
|
||||
# the cap, bill what we have rather than pinning more memory.
|
||||
if not draining_detached:
|
||||
if not _try_claim_detached_drain_slot():
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper: detached-drain cap (%d) reached; billing %d partial "
|
||||
"chunks without draining the rest of the upstream stream",
|
||||
ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS,
|
||||
len(collected_chunks),
|
||||
)
|
||||
await self._bill_collected_chunks(collected_chunks)
|
||||
return
|
||||
draining_detached = True
|
||||
except Exception as exc: # noqa: BLE001 # forward the provider error to a live client, else salvage spend
|
||||
await self._handle_pump_upstream_error(queue, client_detached, collected_chunks, exc)
|
||||
return
|
||||
|
||||
if not client_detached.is_set():
|
||||
if not saw_terminal_event:
|
||||
await self._enqueue_for_client(queue, client_detached, _incomplete_stream_error_sse_event())
|
||||
await self._enqueue_for_client(queue, client_detached, None)
|
||||
await self._bill_collected_chunks(collected_chunks)
|
||||
|
||||
async def _handle_pump_upstream_error(
|
||||
self,
|
||||
queue: "asyncio.Queue[bytes | None | BaseException]",
|
||||
client_detached: "asyncio.Event",
|
||||
collected_chunks: list[bytes], # mutable-ok: SSE buffer forwarded to list-typed _bill_collected_chunks
|
||||
exc: BaseException,
|
||||
) -> None:
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
||||
if not client_detached.is_set() and await self._enqueue_for_client(queue, client_detached, exc):
|
||||
# Preserve the provider-specific failure: the client-facing
|
||||
# generator re-raises it and the proxy's failure handling (status
|
||||
# code, post_call_failure_hook) runs. That path owns logging, so
|
||||
# don't also success-bill.
|
||||
return
|
||||
# Client already gone (or disconnected before the error reached it): no
|
||||
# failure hook will run, so salvage the partial spend instead of
|
||||
# dropping the request entirely.
|
||||
verbose_proxy_logger.warning(
|
||||
"async_sse_wrapper upstream pump failed after client disconnect (%d chunks): %s(%s)",
|
||||
len(collected_chunks),
|
||||
type(exc).__name__,
|
||||
exc,
|
||||
)
|
||||
await self._bill_collected_chunks(collected_chunks)
|
||||
|
|
|
|||
|
|
@ -9,6 +9,7 @@ import pytest
|
|||
sys.path.insert(0, os.path.abspath("../../../../.."))
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages import streaming_iterator as streaming_iterator_module
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
INCOMPLETE_STREAM_ERROR_MESSAGE,
|
||||
BaseAnthropicMessagesStreamingIterator,
|
||||
|
|
@ -440,3 +441,129 @@ async def test_async_sse_wrapper_salvages_partial_spend_on_upstream_error_after_
|
|||
assert len(received) == 2
|
||||
# Partial spend was still recorded rather than the whole request being dropped.
|
||||
assert iterator.logged_chunks == received
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_applies_backpressure_to_slow_client(monkeypatch):
|
||||
"""
|
||||
Regression: the relay queue is bounded, so a slow client throttles the
|
||||
upstream read instead of letting the pump buffer the whole response in
|
||||
memory. With a tiny queue and a client that reads a single chunk, the pump
|
||||
must stall after producing only a queue's worth of chunks ahead, not race
|
||||
to the end of a large stream.
|
||||
"""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 2)
|
||||
|
||||
total = 200
|
||||
produced = 0
|
||||
|
||||
async def _fast_stream():
|
||||
nonlocal produced
|
||||
for i in range(total):
|
||||
produced += 1
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"t{i}"}}
|
||||
|
||||
iterator = _make_iterator("test_backpressure_slow_client")
|
||||
gen = iterator.async_sse_wrapper(_fast_stream())
|
||||
try:
|
||||
await gen.__anext__() # read exactly one chunk, then stall
|
||||
# Let the pump run as far as the bounded queue permits.
|
||||
for _ in range(500):
|
||||
await asyncio.sleep(0)
|
||||
# Bounded by queue maxsize + the one in-flight put + the one delivered
|
||||
# chunk; nowhere near the full 200-chunk stream.
|
||||
assert produced <= 2 + 3, f"pump ran ahead unthrottled: produced {produced} of {total}"
|
||||
assert produced < total
|
||||
finally:
|
||||
await gen.aclose()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_bills_partial_when_detached_drain_cap_reached(monkeypatch):
|
||||
"""
|
||||
Regression: when the concurrent detached-drain cap is already reached, a
|
||||
pump whose client has disconnected must bill what it collected instead of
|
||||
continuing to drain (and accumulating) the rest of a large upstream stream,
|
||||
so slow/abandoned clients can't pin unbounded worker state.
|
||||
|
||||
The cap slot set is pre-occupied so the single slot is unavailable when this
|
||||
pump reaches its first post-disconnect chunk; that isolates the cap decision
|
||||
from multi-pump scheduling races.
|
||||
"""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1)
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4)
|
||||
|
||||
# Occupy the only detached-drain slot with a placeholder task.
|
||||
async def _hold_slot():
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
holder = asyncio.ensure_future(_hold_slot())
|
||||
streaming_iterator_module._DETACHED_STREAM_DRAINS.add(holder)
|
||||
tail_reached = False
|
||||
|
||||
async def _long_stream():
|
||||
nonlocal tail_reached
|
||||
yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}}
|
||||
# These arrive only after the client has disconnected.
|
||||
for i in range(100):
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"more{i}"}}
|
||||
tail_reached = True
|
||||
yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}}
|
||||
yield {"type": "message_stop"}
|
||||
|
||||
iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drain_cap_full"), request_body={})
|
||||
try:
|
||||
gen = iterator.async_sse_wrapper(_long_stream())
|
||||
await gen.__anext__() # message_start
|
||||
await gen.__anext__() # first delta
|
||||
await gen.aclose() # client disconnects; 100+ chunks remain upstream
|
||||
|
||||
for _ in range(200):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0)
|
||||
|
||||
# Billed a small bounded partial (the prefix plus at most a queue's
|
||||
# worth the pump ran ahead before disconnect) without draining the
|
||||
# 100-chunk tail. The exact count depends on how far the bounded queue
|
||||
# let the pump run ahead, so assert the bound, not an exact number.
|
||||
assert iterator.logged_chunks, "capped pump never billed"
|
||||
assert len(iterator.logged_chunks) <= 2 + streaming_iterator_module.ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE
|
||||
assert len(iterator.logged_chunks) < 100
|
||||
assert not any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks)
|
||||
assert tail_reached is False, "pump kept draining past the cap instead of stopping"
|
||||
finally:
|
||||
holder.cancel()
|
||||
streaming_iterator_module._DETACHED_STREAM_DRAINS.discard(holder)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_sse_wrapper_drains_detached_when_cap_available(monkeypatch):
|
||||
"""Complement to the cap test: with a slot free, a disconnected pump drains
|
||||
the full upstream and bills the terminal usage, and releases its slot after."""
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 1)
|
||||
monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4)
|
||||
|
||||
async def _stream():
|
||||
yield {"type": "message_start", "message": {"id": "m", "usage": {"input_tokens": 5, "output_tokens": 1}}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "x"}}
|
||||
for i in range(20):
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": f"m{i}"}}
|
||||
yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 42}}
|
||||
yield {"type": "message_stop"}
|
||||
|
||||
iterator = _RecordingLoggingIterator(litellm_logging_obj=_make_logging_obj("drain_cap_free"), request_body={})
|
||||
gen = iterator.async_sse_wrapper(_stream())
|
||||
await gen.__anext__()
|
||||
await gen.__anext__()
|
||||
await gen.aclose()
|
||||
|
||||
for _ in range(300):
|
||||
if iterator.logged_chunks:
|
||||
break
|
||||
await asyncio.sleep(0.01)
|
||||
|
||||
assert any(c.startswith(b"event: message_stop\n") for c in iterator.logged_chunks)
|
||||
# Slot released once the drain finished.
|
||||
assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue