feat(pointfive): add gzipped ndjson batch encoding

litellm/integrations/pointfive/payload.py

@@ -0,0 +1,53 @@
+"""Turns buffered log records into the gzipped NDJSON objects that get uploaded."""
+
+import gzip
+from collections.abc import Iterator, Mapping, Sequence
+from itertools import accumulate, groupby, islice
+from typing import Final
+
+from litellm.litellm_core_utils.asyncify import asyncify
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+
+_NEWLINE_BYTES: Final = 1
+
+
+def serialize_records(records: Sequence[Mapping[str, object]]) -> tuple[str, ...]:
+    """Serialize each record to one JSON line."""
+    return tuple(safe_dumps(record) for record in records)
+
+
+def _encoded_size(line: str) -> int:
+    return len(line.encode("utf-8")) + _NEWLINE_BYTES
+
+
+def _object_indices(sizes: Sequence[int], max_bytes: int) -> Iterator[int]:
+    """Number each line with the object it belongs to, opening a new one on overflow."""
+
+    def advance(state: tuple[int, int], size: int) -> tuple[int, int]:
+        index, used = state
+        return (index + 1, size) if used and used + size > max_bytes else (index, used + size)
+
+    return (index for index, _ in islice(accumulate(sizes, advance, initial=(0, 0)), 1, None))
+
+
+def chunk_lines(lines: Sequence[str], max_bytes: int) -> tuple[tuple[str, ...], ...]:
+    """
+    Group serialized lines into objects of at most ``max_bytes`` uncompressed.
+
+    A line above the bound on its own still becomes its own object. A record cannot be
+    split, and holding it back would stall every record queued behind it.
+    """
+    sizes: Final = tuple(_encoded_size(line) for line in lines)
+    numbered: Final = zip(_object_indices(sizes, max_bytes), lines, strict=True)
+    return tuple(tuple(line for _, line in group) for _, group in groupby(numbered, lambda pair: pair[0]))
+
+
+async def encode_lines(lines: Sequence[str]) -> bytes:
+    """
+    Join lines as NDJSON and gzip them off the event loop.
+
+    An object can be several megabytes, and compressing that inline would block the
+    proxy for as long as it takes.
+    """
+    compress: Final = asyncify(gzip.compress)
+    return await compress("\n".join(lines).encode("utf-8"))

tests/test_litellm/integrations/pointfive/test_payload.py

@@ -0,0 +1,80 @@
+import gzip
+import json
+
+import pytest
+
+from litellm.integrations.pointfive.payload import chunk_lines, encode_lines, serialize_records
+
+UNBOUNDED = 10_000_000
+
+
+def test_each_record_becomes_one_json_line():
+    lines = serialize_records([{"id": "a"}, {"id": "b"}, {"id": "c"}])
+
+    assert len(lines) == 3
+    assert [json.loads(line)["id"] for line in lines] == ["a", "b", "c"]
+
+
+def test_non_serializable_values_do_not_raise():
+    """An odd payload must not kill the flush."""
+    lines = serialize_records([{"id": "a", "when": object()}])
+
+    assert json.loads(lines[0])["id"] == "a"
+
+
+def test_records_that_fit_stay_in_one_object():
+    lines = serialize_records([{"id": f"r{i}"} for i in range(50)])
+
+    assert chunk_lines(lines, UNBOUNDED) == (lines,)
+
+
+def test_objects_are_capped_by_uncompressed_size():
+    lines = serialize_records([{"id": f"r{i}", "blob": "x" * 100} for i in range(10)])
+    line_bytes = len(lines[0].encode("utf-8")) + 1
+
+    chunks = chunk_lines(lines, line_bytes * 3)
+
+    assert [len(chunk) for chunk in chunks] == [3, 3, 3, 1]
+
+
+def test_oversized_single_record_is_sent_alone_not_stalled():
+    """A record too big for the cap must still go out, or it blocks everything behind it."""
+    lines = serialize_records([{"id": "small"}, {"id": "huge", "blob": "x" * 5000}, {"id": "small2"}])
+
+    chunks = chunk_lines(lines, 200)
+
+    assert sum(len(chunk) for chunk in chunks) == 3
+    huge = [chunk for chunk in chunks if any("huge" in line for line in chunk)]
+    assert len(huge) == 1
+    assert len(huge[0]) == 1
+
+
+def test_no_records_produces_no_objects():
+    assert chunk_lines((), UNBOUNDED) == ()
+
+
+def test_every_record_appears_exactly_once():
+    lines = serialize_records([{"id": f"r{i}"} for i in range(37)])
+
+    chunks = chunk_lines(lines, len(lines[0]) * 4)
+
+    assert [line for chunk in chunks for line in chunk] == list(lines)
+
+
+@pytest.mark.asyncio
+async def test_encode_lines_round_trips_through_gzip():
+    lines = serialize_records([{"id": f"r{i}"} for i in range(5)])
+
+    encoded = await encode_lines(lines)
+
+    assert encoded[:2] == b"\x1f\x8b"
+    assert gzip.decompress(encoded).decode("utf-8") == "\n".join(lines)
+
+
+@pytest.mark.asyncio
+async def test_encode_lines_compresses_repetitive_records():
+    lines = serialize_records([{"id": f"r{i}", "model": "gpt-4o", "cost": 0.01} for i in range(200)])
+
+    encoded = await encode_lines(lines)
+
+    assert len(encoded) < len(gzip.decompress(encoded)) / 2