mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
321 lines
10 KiB
Python
321 lines
10 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_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()
|