diff --git a/litellm/__init__.py b/litellm/__init__.py index 5f14c605cd0..aed13486d22 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -249,6 +249,10 @@ anthropic_messages_estimate_initial_input_tokens: bool = False strip_anthropic_total_tokens: bool = False anthropic_sse_ping_interval_seconds: float = 15.0 sse_keepalive_ping_interval_seconds: float | None = None +# Most bytes of one streamed pass-through response kept in memory for end-of-stream logging and cost +# tracking. Past it the middle is dropped and the first and last halves are kept, so the final usage +# event survives. None keeps everything. Set via `litellm_settings.pass_through_max_logged_stream_bytes`. +pass_through_max_logged_stream_bytes: int | None = 16 * 1024 * 1024 route_all_chat_openai_to_responses: bool = ( os.getenv("LITELLM_ROUTE_ALL_CHAT_OPENAI_TO_RESPONSES", "false").lower() == "true" ) # When True, routes all OpenAI /chat/completions requests through the Responses API bridge diff --git a/litellm/proxy/pass_through_endpoints/bounded_chunk_buffer.py b/litellm/proxy/pass_through_endpoints/bounded_chunk_buffer.py new file mode 100644 index 00000000000..2dd980af5ba --- /dev/null +++ b/litellm/proxy/pass_through_endpoints/bounded_chunk_buffer.py @@ -0,0 +1,123 @@ +from collections import deque +from collections.abc import Iterator, Sequence +from typing import Final, overload + +_FRAME_SEPARATORS: Final = (b"\r\n\r\n", b"\n\n") + + +def _end_of_first_frame(data: bytes) -> int: + """Index just past the first complete frame in ``data``; 0 when it holds none.""" + ends: Final = [i + len(sep) for sep in _FRAME_SEPARATORS if (i := data.find(sep)) != -1] + if ends: + return min(ends) + newline: Final = data.find(b"\n") + return 0 if newline == -1 else newline + 1 + + +def _end_of_last_frame(data: bytes) -> int: + """Index just past the last complete frame in ``data``; 0 when it holds none.""" + ends: Final = [i + len(sep) for sep in _FRAME_SEPARATORS if (i := data.rfind(sep)) != -1] + if ends: + return max(ends) + newline: Final = data.rfind(b"\n") + return 0 if newline == -1 else newline + 1 + + +class BoundedChunkBuffer(Sequence[bytes]): + """ + Collects the chunks of a streamed response for end-of-stream logging without letting + an unbounded stream exhaust memory. + + Half of ``max_bytes`` keeps the start of the stream and half keeps the end, so the + first event (input usage, model) and the terminal event (final usage) both survive. + When the stream fits, every chunk is kept untouched. When it does not, the bytes in + between are dropped and both ends are trimmed back to a frame boundary, so the + retained bytes never contain a half-cut event. + + ``max_bytes=None`` keeps everything. + """ + + def __init__(self, max_bytes: int | None = None) -> None: + self._head_limit: Final = None if max_bytes is None else max_bytes // 2 + self._tail_limit: Final = None if max_bytes is None else max_bytes - (max_bytes // 2) + self._head: Final[list[bytes]] = [] # mutable-ok: append-only buffer owned by this class + self._tail: Final[deque[bytes]] = deque() # mutable-ok: append-only buffer owned by this class + self._head_size = 0 + self._tail_size = 0 + self.total_bytes = 0 + self._head_closed = False + + @property + def truncated(self) -> bool: + return self.total_bytes > self._head_size + self._tail_size + + @property + def retained_bytes(self) -> int: + return sum(len(chunk) for chunk in self) + + def append(self, chunk: bytes) -> None: + self.total_bytes += len(chunk) + if self._head_limit is None or self._tail_limit is None: + self._head.append(chunk) + self._head_size += len(chunk) + return + chunk_rest: Final = chunk if self._head_closed else self._fill_head(chunk, self._head_limit) + if chunk_rest is None: + return + self._tail.append(chunk_rest[-self._tail_limit :] if self._tail_limit else b"") + self._tail_size += len(self._tail[-1]) + while self._tail_size > self._tail_limit: + overflow = self._tail_size - self._tail_limit + oldest = self._tail.popleft() + if len(oldest) > overflow: + self._tail.appendleft(oldest[overflow:]) + self._tail_size -= overflow + break + self._tail_size -= len(oldest) + + def _fill_head(self, chunk: bytes, head_limit: int) -> bytes | None: + """Keep what fits in the head; the rest of the chunk, or None when all of it fit.""" + room: Final = head_limit - self._head_size + if len(chunk) <= room: + self._head.append(chunk) + self._head_size += len(chunk) + return None + if room > 0: + self._head.append(chunk[:room]) + self._head_size += room + self._head_closed = True + return chunk[room:] + + def _frames(self) -> Sequence[bytes]: + if not self.truncated: + return [*self._head, *self._tail] + head: Final = b"".join(self._head) + tail: Final = b"".join(self._tail) + head_end: Final = _end_of_last_frame(head) + tail_start: Final = 0 if not self._tail_size else self._tail_start(tail) + return [part for part in (head[:head_end], tail[tail_start:]) if part] + + def _tail_start(self, tail: bytes) -> int: + # the tail begins wherever the eviction cut it, which is mid-frame unless nothing before it was dropped + return _end_of_first_frame(tail) if self._head_closed else 0 + + def __iter__(self) -> Iterator[bytes]: + return iter(self._frames()) + + def __len__(self) -> int: + return len(self._frames()) + + def __eq__(self, other: object) -> bool: + return list(self) == other + + def __bool__(self) -> bool: + return self.total_bytes > 0 + + @overload + def __getitem__(self, index: int) -> bytes: ... + + @overload + def __getitem__(self, index: slice) -> Sequence[bytes]: ... + + def __getitem__(self, index: int | slice) -> bytes | Sequence[bytes]: + return self._frames()[index] diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index c94dc16a8f5..1d9f52cfcae 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -19,6 +19,7 @@ from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType from litellm.types.utils import StandardPassThroughResponseObject, Usage +from .bounded_chunk_buffer import BoundedChunkBuffer from .llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, ) @@ -49,7 +50,7 @@ class RouteStreamingLogging(Protocol): request_body: dict, endpoint_type: EndpointType, start_time: datetime, - raw_bytes: list[bytes], + raw_bytes: Sequence[bytes], end_time: datetime, ) -> Coroutine[None, None, None]: ... @@ -170,10 +171,11 @@ class PassThroughStreamingHandler: resolved_route_streaming_logging: Final[RouteStreamingLogging] = ( route_streaming_logging or PassThroughStreamingHandler.route_streaming_logging_to_handler ) - raw_bytes: Final[list[bytes]] = [] + raw_bytes: Final = BoundedChunkBuffer(litellm.pass_through_max_logged_stream_bytes) resolved_request_body: Final[dict[str, object]] = request_body or {} def _build_logging_coroutine() -> Coroutine[None, None, None]: + PassThroughStreamingHandler._warn_if_truncated(raw_bytes, url_route) return resolved_route_streaming_logging( litellm_logging_obj=litellm_logging_obj, passthrough_success_handler_obj=passthrough_success_handler_obj, @@ -295,6 +297,7 @@ class PassThroughStreamingHandler: if stream_ended else PassThroughStreamingHandler._error_frame_closed_at(endpoint_type, raw_bytes) ) + PassThroughStreamingHandler._warn_if_truncated(raw_bytes, url_route) try: await PassThroughStreamingHandler._schedule_stream_end_logging( litellm_logging_obj=litellm_logging_obj, @@ -309,6 +312,17 @@ class PassThroughStreamingHandler: else: bind_budget_reservation_to_callbacks(litellm_logging_obj.litellm_params) + @staticmethod + def _warn_if_truncated(raw_bytes: BoundedChunkBuffer, url_route: str) -> None: + if raw_bytes.truncated: + verbose_proxy_logger.warning( + "Pass-through stream for %s was %s bytes; only the first and last %s bytes are kept for logging " + "(litellm.pass_through_max_logged_stream_bytes)", + url_route, + raw_bytes.total_bytes, + raw_bytes.retained_bytes, + ) + @staticmethod def _error_frame_closed_at(endpoint_type: EndpointType, raw_bytes: Sequence[bytes]) -> tuple[str, str, int] | None: from litellm.llms.anthropic.messages.native.streaming_iterator import ( diff --git a/tests/unit/proxy/pass_through_endpoints/test_bounded_chunk_buffer.py b/tests/unit/proxy/pass_through_endpoints/test_bounded_chunk_buffer.py new file mode 100644 index 00000000000..3b740c2d961 --- /dev/null +++ b/tests/unit/proxy/pass_through_endpoints/test_bounded_chunk_buffer.py @@ -0,0 +1,163 @@ +import json +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest + +import litellm +from litellm.proxy.pass_through_endpoints.bounded_chunk_buffer import BoundedChunkBuffer +from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler +from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType + + +def _frame(name: str, payload: dict) -> bytes: + return f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() + + +def test_buffer_under_the_limit_keeps_every_chunk_untouched(): + buffer = BoundedChunkBuffer(max_bytes=1024) + chunks = [b"event: a\ndata: 1\n\n", b"event: b\ndata: 2\n\n"] + + for chunk in chunks: + buffer.append(chunk) + + assert list(buffer) == chunks, f"truncated={buffer.truncated}" + assert not buffer.truncated + + +def test_buffer_without_a_limit_keeps_everything(): + buffer = BoundedChunkBuffer(max_bytes=None) + chunks = [b"x" * 1_000_000 for _ in range(5)] + + for chunk in chunks: + buffer.append(chunk) + + assert list(buffer) == chunks + assert buffer.total_bytes == 5_000_000 + + +def test_buffer_over_the_limit_keeps_the_first_and_last_events_whole(): + first = _frame("message_start", {"usage": {"input_tokens": 29}}) + middle = [_frame("content_block_delta", {"text": "x" * 40, "n": i}) for i in range(200)] + last = _frame("message_delta", {"usage": {"output_tokens": 77}}) + buffer = BoundedChunkBuffer(max_bytes=600) + + for chunk in [first, *middle, last]: + buffer.append(chunk) + + kept = b"".join(buffer) + assert kept.startswith(first), kept + assert kept.endswith(last), kept + assert buffer.truncated + assert buffer.retained_bytes <= 600 + assert buffer.total_bytes == len(first) + sum(map(len, middle)) + len(last) + + +def test_buffer_never_keeps_a_half_cut_event_when_chunks_split_frames(): + stream = b"".join(_frame("delta", {"i": i, "pad": "y" * 30}) for i in range(100)) + buffer = BoundedChunkBuffer(max_bytes=500) + + for start in range(0, len(stream), 7): + buffer.append(stream[start : start + 7]) + + for part in b"".join(buffer).split(b"\n\n"): + if part: + assert part.startswith(b"event: delta\ndata: {"), part + assert part.endswith(b"}"), part + + +def test_buffer_holds_one_chunk_larger_than_the_limit_to_its_head_and_tail(): + buffer = BoundedChunkBuffer(max_bytes=100) + + buffer.append(b"event: big\ndata: " + b"z" * 10_000 + b"\n\n") + + assert buffer.truncated + assert buffer.retained_bytes <= 100 + assert buffer.total_bytes == 10_000 + len(b"event: big\ndata: \n\n") + + +def test_buffer_memory_stays_bounded_while_a_huge_stream_passes(): + buffer = BoundedChunkBuffer(max_bytes=1000) + + for _ in range(10_000): + buffer.append(b"data: " + b"q" * 1000 + b"\n\n") + + assert buffer.total_bytes > 10_000_000 + assert buffer.retained_bytes <= 1000 + + +def test_buffer_trims_newline_delimited_streams_to_whole_lines(): + # no blank-line separators: frames are single lines, as in newline-delimited JSON + buffer = BoundedChunkBuffer(max_bytes=40) + lines = [json.dumps({"n": i}).encode() + b"\n" for i in range(20)] + + for line in lines: + buffer.append(line) + + kept = b"".join(buffer) + assert buffer.truncated + assert kept.startswith(lines[0]) + assert kept.endswith(lines[-1]) + assert all(json.loads(part) is not None for part in kept.splitlines()) + + +def test_buffer_supports_indexing_slicing_and_equality_with_a_list(): + buffer = BoundedChunkBuffer(max_bytes=1024) + chunks = [b"event: a\ndata: 1\n\n", b"event: b\ndata: 2\n\n"] + for chunk in chunks: + buffer.append(chunk) + + assert buffer[0] == chunks[0] + assert buffer[-1] == chunks[-1] + assert buffer[0:2] == chunks + assert buffer == chunks + assert buffer != chunks[:1] + assert buffer != "not a sequence" + + +def _response(chunks: list[bytes]) -> MagicMock: + response = MagicMock(spec=httpx.Response) + response.status_code = 200 + + async def _aiter_bytes(): + for chunk in chunks: + yield chunk + + response.aiter_bytes = _aiter_bytes + return response + + +@pytest.mark.asyncio +async def test_chunk_processor_relays_every_byte_but_logs_only_the_bounded_ends( + monkeypatch: pytest.MonkeyPatch, caplog +): + monkeypatch.setattr(litellm, "pass_through_max_logged_stream_bytes", 400) + first = _frame("message_start", {"usage": {"input_tokens": 29}}) + middle = [_frame("content_block_delta", {"text": "x" * 40, "n": i}) for i in range(100)] + last = _frame("message_delta", {"usage": {"output_tokens": 77}}) + chunks = [first, *middle, last] + logging_obj = MagicMock() + logging_obj._on_deferred_stream_complete = None + + with patch.object(PassThroughStreamingHandler, "route_streaming_logging_to_handler", new=AsyncMock()) as route: + relayed = [ + chunk + async for chunk in PassThroughStreamingHandler.chunk_processor( + response=_response(chunks), + request_body={"model": "m"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.GENERIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/codex/responses", + ) + ] + import asyncio + + await asyncio.sleep(0) + + assert relayed == chunks + logged = b"".join(route.call_args.kwargs["raw_bytes"]) + assert logged.startswith(first) and logged.endswith(last), logged + assert len(logged) <= 400