Merge 16367a5969 into b195d15b93

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

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]

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 (

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