mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
ef1cb19351
* test(integration): Converse stream post_call guardrail relays frames as they arrive (LIT-9515) * fix(bedrock): stream Converse passthrough frames through post_call guardrails instead of buffering (LIT-9515) A post_call guardrail that scans through apply_guardrail now runs on /bedrock/.../converse-stream the way it runs on /v1/messages: frames reach the caller as the upstream sends them, every fifth frame and the end of the stream are scanned, and a block after the first frame ends the stream with an AWS exception frame. streaming_buffer_until_moderated still holds the stream and applies masks per content block. Guardrails that need the whole response (own async_post_call_success_hook, mask_response_content, native hooks, policy pipelines) keep the buffered path. * fix(bedrock): apply held Converse stream masks when the upstream ends with an exception frame (LIT-9515) * fix(proxy): log a relayed passthrough stream's spend before its mid-stream failure hook (LIT-9515) * fix(bedrock): close and report every way a relayed Converse stream ends (LIT-9515) - The framer ends the stream with an empty item once the upstream closes, so held masks are applied whether the upstream ends with messageStop, an exception or error frame, or a dropped connection - A guardrail HTTPException after the first frame reaches the relay, which logs the stream's spend, runs the failure hook and ends the stream with the exception frame, instead of ending inside the guardrail with no failure row - A failure before the first frame waits for the stream's spend to be logged, as a mid-stream failure does - Closing the response before its body starts closes the guardrail layers and the upstream - Each stacked guardrail layer gets its own copy of every frame, so an earlier layer's end-of-stream rewrite never reaches frames a later layer still holds - Deltas whose text ends in half a surrogate pair are decoded with the standard JSON parser * fix(bedrock): relay a Converse frame whose payload is not UTF-8 instead of ending the stream (LIT-9515) * test(integration): cover Converse stream relay failure, end-of-stream-only, stacked, disconnect and burst cells (LIT-9515) * fix(bedrock): estimate prompt tokens for a passthrough stream cut before Bedrock reports usage (LIT-9515) Bedrock sends usage only in the last frame of a stream, so a stream ended by a guardrail block or a client disconnect logged 0 prompt tokens while AWS still charges for the prompt. The stream collector now estimates them from the request the way chat completions streams do * fix(bedrock): read tool result blocks without recursion in the passthrough prompt estimate (LIT-9515) * fix(passthrough): keep the spend flush running when the wait for it is cancelled (LIT-9515) * fix(bedrock): keep a converse-stream buffered when its guardrail releases a held stream on scan (LIT-9515) A held post_call guardrail with streaming_buffer_release_on_scan releases each window as received once its scan passes, so the relay delivered the original text instead of the guardrail's masks. Such a stream keeps the buffered path, which applies masks per content block as on main. * fix(bedrock): keep a converse-stream buffered when a post_call guardrail runs in parallel (LIT-9515) The relay ran every guardrail in config order, so a run_in_parallel guardrail listed before a masking guardrail scanned the unmasked text. Main's buffered path runs parallel guardrails after the sequential ones and discards their rewrites, which the relay cannot reproduce, so these streams stay buffered * fix(bedrock): mark Noma v2 relay scans like chat and stop reading Bedrock when a held stream's client leaves - Noma v2 scans on the converse-stream relay carry stream, and the last one carries the finished Converse response with its stopReason, as chat sends them, so Noma keeps one record per reply. Other guardrails get the request unchanged. When the last sampled scan already covered the whole answer, the end-of-stream scan still runs for Noma so it gets the finished answer - A held relay stream whose client leaves before the first frame stops reading Bedrock and returns 499, like chat. When Bedrock sent nothing yet, the budget reservation and the parallel-request slot are released - A cancelled wait for the first streamed chunk now stops the upstream read instead of leaving it running - A Noma v2 scan cancelled before Noma answers records not_run instead of success with an empty response * fix(bedrock): send Noma the finished Converse reply's tool calls and reasoning as Converse blocks * fix(bedrock): give Noma the ended scan of a Converse reply that only calls a tool without input * test(bedrock): expect the ValidationException main now returns for a held block before the first frame Main's AWS-shaped passthrough errors (#46522) answer a guardrail block before the first frame with x-amzn-errortype ValidationException, so boto3 reads the code and the guardrail message instead of a bare 400. * fix(bedrock): keep held masks when an earlier sampled guardrail ends the Converse stream with its message A sampled guardrail that blocks with disable_exception_on_block ends the stream with its message as raw event-stream bytes. A held guardrail after it on the same key only read frame objects, so it never saw the messageStop, skipped its write-back and released the held text without its masks. The Converse handler now reads those bytes back into frames, so the held guardrail sees the end of the message and masks the text * fix(proxy): free the parallel slot when a caller leaves a Converse relay holding part of the first frame The relay skipped releasing the slot and budget reservation whenever a spend flush was scheduled, but a flush of a partial first frame builds no response and logs nothing, so the slot stayed held until it expired. The flush now says whether it logged the stream. --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com>
434 lines
14 KiB
Python
434 lines
14 KiB
Python
"""Regression tests for LIT-2642 — interrupted streams must still flush usage."""
|
|
|
|
import asyncio
|
|
from typing import List
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
|
|
def _make_streaming_response(chunks: List[bytes]):
|
|
mock = MagicMock(spec=httpx.Response)
|
|
mock.status_code = 200
|
|
mock.headers = httpx.Headers({"content-type": "application/vnd.amazon.eventstream"})
|
|
mock.raise_for_status = MagicMock(return_value=None)
|
|
|
|
async def _aiter_bytes():
|
|
for chunk in chunks:
|
|
yield chunk
|
|
|
|
mock.aiter_bytes = _aiter_bytes
|
|
mock.aclose = AsyncMock()
|
|
return mock
|
|
|
|
|
|
def _make_logging_obj():
|
|
mock = MagicMock()
|
|
mock.async_flush_passthrough_collected_chunks = AsyncMock()
|
|
return mock
|
|
|
|
|
|
class _ImmediateExecutor:
|
|
def submit(self, fn, *args, **kwargs):
|
|
fn(*args, **kwargs)
|
|
|
|
|
|
class _RecordingCollector:
|
|
def __init__(self) -> None:
|
|
self.chunks: List[bytes] = []
|
|
|
|
def add(self, chunk: bytes) -> None:
|
|
self.chunks.append(chunk)
|
|
|
|
def build_logged_response(self, litellm_logging_obj: MagicMock) -> bytes:
|
|
return b"".join(self.chunks)
|
|
|
|
|
|
class _FailingCollector(_RecordingCollector):
|
|
def add(self, chunk: bytes) -> None:
|
|
raise ValueError("bad frame")
|
|
|
|
|
|
def _provider_config(collector: _RecordingCollector) -> MagicMock:
|
|
provider_config = MagicMock()
|
|
provider_config.create_stream_collector.return_value = collector
|
|
return provider_config
|
|
|
|
|
|
def _spend_payload(flush_mock: MagicMock) -> bytes:
|
|
flush_mock.assert_called_once()
|
|
collector = flush_mock.call_args.kwargs["collector"]
|
|
return collector.build_logged_response(litellm_logging_obj=MagicMock())
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_flushes_on_normal_completion():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
|
|
mock_response = _make_streaming_response(chunks)
|
|
mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})
|
|
|
|
async def response_coro():
|
|
return mock_response
|
|
|
|
mock_logging_obj = _make_logging_obj()
|
|
|
|
received = []
|
|
received_response = AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
)
|
|
|
|
async for chunk in received_response:
|
|
received.append(chunk)
|
|
|
|
assert received == chunks
|
|
|
|
assert received_response.headers["content-type"] == "application/octet-stream"
|
|
assert received_response.headers["x-request-id"] == "req-123"
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == b"".join(chunks)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_flushes_on_client_disconnect():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
chunks = [
|
|
b'{"chunk": 1, "outputTokens": 10}',
|
|
b'{"chunk": 2, "outputTokens": 12}',
|
|
b'{"chunk": 3, "outputTokens": 8}',
|
|
]
|
|
mock_response = _make_streaming_response(chunks)
|
|
mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})
|
|
|
|
async def response_coro():
|
|
return mock_response
|
|
|
|
mock_logging_obj = _make_logging_obj()
|
|
|
|
gen = AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
)
|
|
|
|
received = [await gen.__anext__()]
|
|
await gen.aclose()
|
|
|
|
assert received == [chunks[0]]
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == chunks[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_wait_for_spend_flush_returns_once_the_flush_has_finished():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
mock_response = _make_streaming_response([b"chunk-1", b"chunk-2"])
|
|
|
|
async def response_coro():
|
|
return mock_response
|
|
|
|
flushed: List[bytes] = []
|
|
|
|
async def flush(collector: _RecordingCollector) -> bool:
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
flushed.append(collector.build_logged_response(litellm_logging_obj=MagicMock()))
|
|
return True
|
|
|
|
mock_logging_obj = MagicMock()
|
|
mock_logging_obj.async_flush_passthrough_collected_chunks = flush
|
|
gen = AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
)
|
|
|
|
await gen.__anext__()
|
|
await gen.aclose()
|
|
|
|
assert await gen.wait_for_spend_flush() is True
|
|
assert flushed == [b"chunk-1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_says_it_logged_no_spend_when_the_flush_had_nothing_to_log():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
async def response_coro():
|
|
return _make_streaming_response([b"part-of-a-frame", b"rest-of-the-frame"])
|
|
|
|
mock_logging_obj = MagicMock()
|
|
mock_logging_obj.async_flush_passthrough_collected_chunks = AsyncMock(return_value=False)
|
|
gen = AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
)
|
|
|
|
await gen.__anext__()
|
|
await gen.aclose()
|
|
|
|
assert await gen.wait_for_spend_flush() is False
|
|
mock_logging_obj.async_flush_passthrough_collected_chunks.assert_awaited_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_closed_before_any_chunk_says_it_logged_no_spend():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
async def response_coro():
|
|
return _make_streaming_response([b"chunk-1"])
|
|
|
|
mock_logging_obj = _make_logging_obj()
|
|
gen = AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
)
|
|
|
|
await gen
|
|
await gen.aclose()
|
|
|
|
assert await gen.wait_for_spend_flush() is False
|
|
mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_spend_flush_survives_a_cancelled_wait():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
mock_response = _make_streaming_response([b"chunk-1", b"chunk-2"])
|
|
|
|
async def response_coro():
|
|
return mock_response
|
|
|
|
release = asyncio.Event()
|
|
flushed: List[bytes] = []
|
|
|
|
async def flush(collector: _RecordingCollector) -> None:
|
|
await release.wait()
|
|
flushed.append(collector.build_logged_response(litellm_logging_obj=MagicMock()))
|
|
|
|
mock_logging_obj = MagicMock()
|
|
mock_logging_obj.async_flush_passthrough_collected_chunks = flush
|
|
gen = AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
)
|
|
|
|
await gen.__anext__()
|
|
await gen.aclose()
|
|
waiter = asyncio.create_task(gen.wait_for_spend_flush())
|
|
await asyncio.sleep(0)
|
|
await asyncio.sleep(0)
|
|
waiter.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await waiter
|
|
release.set()
|
|
await gen.wait_for_spend_flush()
|
|
|
|
assert flushed == [b"chunk-1"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_does_not_flush_on_4xx():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
err_response = MagicMock(spec=httpx.Response)
|
|
err_response.status_code = 429
|
|
err_response.headers = httpx.Headers({"content-type": "application/octet-stream"})
|
|
|
|
def _raise():
|
|
raise httpx.HTTPStatusError(
|
|
"429",
|
|
request=httpx.Request("POST", "https://example.com"),
|
|
response=httpx.Response(429, request=httpx.Request("POST", "https://example.com")),
|
|
)
|
|
|
|
err_response.raise_for_status = _raise
|
|
err_response.aclose = AsyncMock()
|
|
|
|
async def response_coro():
|
|
return err_response
|
|
|
|
mock_logging_obj = _make_logging_obj()
|
|
|
|
with pytest.raises(httpx.HTTPStatusError):
|
|
async for _ in AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=MagicMock(),
|
|
):
|
|
pass
|
|
|
|
mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_flushes_on_upstream_exception_with_partial_data():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
partial_chunks = [b"partial-chunk-1", b"partial-chunk-2"]
|
|
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.status_code = 200
|
|
mock_response.raise_for_status = MagicMock(return_value=None)
|
|
mock_response.aclose = AsyncMock()
|
|
mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})
|
|
|
|
async def _aiter_bytes_then_raise():
|
|
for c in partial_chunks:
|
|
yield c
|
|
raise httpx.ReadError("upstream disconnected")
|
|
|
|
mock_response.aiter_bytes = _aiter_bytes_then_raise
|
|
|
|
async def response_coro():
|
|
return mock_response
|
|
|
|
mock_logging_obj = _make_logging_obj()
|
|
|
|
received = []
|
|
|
|
async def _drain():
|
|
async for chunk in AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
):
|
|
received.append(chunk)
|
|
|
|
with pytest.raises(httpx.ReadError):
|
|
await _drain()
|
|
|
|
assert received == partial_chunks
|
|
|
|
await asyncio.sleep(0)
|
|
|
|
assert _spend_payload(mock_logging_obj.async_flush_passthrough_collected_chunks) == b"".join(partial_chunks)
|
|
|
|
|
|
def test_passthroughstreamingresponse_flushes_on_normal_completion():
|
|
from litellm.passthrough.main import PassthroughStreamingResponse
|
|
|
|
chunks = [b"a", b"b", b"c"]
|
|
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.status_code = 200
|
|
mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})
|
|
|
|
def _iter_bytes():
|
|
yield from chunks
|
|
|
|
mock_response.iter_bytes = _iter_bytes
|
|
|
|
mock_logging_obj = MagicMock()
|
|
mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
|
|
|
|
received_responce = PassthroughStreamingResponse(
|
|
response=mock_response,
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
)
|
|
|
|
with patch("litellm.utils.executor", _ImmediateExecutor()):
|
|
received = list(received_responce)
|
|
|
|
assert received == chunks
|
|
|
|
assert received_responce.headers["content-type"] == "application/octet-stream"
|
|
assert received_responce.headers["x-request-id"] == "req-123"
|
|
|
|
assert _spend_payload(mock_logging_obj.flush_passthrough_collected_chunks) == b"".join(chunks)
|
|
|
|
|
|
def test_passthroughstreamingresponse_flushes_on_early_close():
|
|
from litellm.passthrough.main import PassthroughStreamingResponse
|
|
|
|
chunks = [b"first", b"second", b"third"]
|
|
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.status_code = 200
|
|
mock_response.headers = httpx.Headers({"content-type": "application/octet-stream", "x-request-id": "req-123"})
|
|
|
|
def _iter_bytes():
|
|
yield from chunks
|
|
|
|
mock_response.iter_bytes = _iter_bytes
|
|
|
|
mock_logging_obj = MagicMock()
|
|
mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
|
|
|
|
with patch("litellm.utils.executor", _ImmediateExecutor()):
|
|
gen = PassthroughStreamingResponse(
|
|
response=mock_response,
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_RecordingCollector()),
|
|
)
|
|
|
|
first = next(gen)
|
|
gen.close()
|
|
|
|
assert first == chunks[0]
|
|
assert _spend_payload(mock_logging_obj.flush_passthrough_collected_chunks) == chunks[0]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_asyncpassthroughstreamingresponse_relays_the_stream_when_spend_parsing_fails():
|
|
from litellm.passthrough.main import AsyncPassthroughStreamingResponse
|
|
|
|
chunks = [b"chunk-1", b"chunk-2", b"chunk-3"]
|
|
mock_response = _make_streaming_response(chunks)
|
|
|
|
async def response_coro():
|
|
return mock_response
|
|
|
|
mock_logging_obj = _make_logging_obj()
|
|
|
|
received = [
|
|
chunk
|
|
async for chunk in AsyncPassthroughStreamingResponse(
|
|
response=response_coro(),
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_FailingCollector()),
|
|
)
|
|
]
|
|
await asyncio.sleep(0)
|
|
|
|
assert received == chunks
|
|
mock_logging_obj.async_flush_passthrough_collected_chunks.assert_not_called()
|
|
|
|
|
|
def test_passthroughstreamingresponse_relays_the_stream_when_spend_parsing_fails():
|
|
from litellm.passthrough.main import PassthroughStreamingResponse
|
|
|
|
chunks = [b"a", b"b", b"c"]
|
|
mock_response = MagicMock(spec=httpx.Response)
|
|
mock_response.status_code = 200
|
|
mock_response.headers = httpx.Headers({"content-type": "application/octet-stream"})
|
|
mock_response.iter_bytes = lambda: iter(chunks)
|
|
|
|
mock_logging_obj = MagicMock()
|
|
mock_logging_obj.flush_passthrough_collected_chunks = MagicMock()
|
|
|
|
received = list(
|
|
PassthroughStreamingResponse(
|
|
response=mock_response,
|
|
litellm_logging_obj=mock_logging_obj,
|
|
provider_config=_provider_config(_FailingCollector()),
|
|
)
|
|
)
|
|
|
|
assert received == chunks
|
|
mock_logging_obj.flush_passthrough_collected_chunks.assert_not_called()
|