diff --git a/litellm/litellm_core_utils/json_fragment_accumulator.py b/litellm/litellm_core_utils/json_fragment_accumulator.py index d94b4510830..e829309b43b 100644 --- a/litellm/litellm_core_utils/json_fragment_accumulator.py +++ b/litellm/litellm_core_utils/json_fragment_accumulator.py @@ -1,25 +1,36 @@ import json from typing import Final, cast # noqa: TID251 # raw_decode returns tuple[Any, int]; no cast-free unpack +_JSON_WHITESPACE: Final = " \t\n\r" # matches json.decoder.WHITESPACE's character class + class JSONFragmentAccumulator: """ Buffers a JSON value that arrives piecemeal over a stream (SSE data split across TCP packets, one shard per network read, etc) without the O(n^2) - cost of repeated `buffer += fragment` string concatenation. + cost of repeated `buffer += fragment` string concatenation, and without + the O(n^2) cost of re-copying the unconsumed remainder on every peeled + value when one payload holds many concatenated JSON values. Fragments are appended to a list in O(1). The buffer is only rebuilt into a single string, and only decoded, when a caller asks for a value via `pop_next_value`, and `could_close_json` lets callers skip that rebuild - entirely for fragments that plainly cannot close a JSON value yet. + entirely for fragments that plainly cannot close a JSON value yet. Once + rebuilt, consumed values are dropped by advancing a cursor rather than + slicing a new string, so draining N concatenated values already sitting + in the buffer costs O(n) total, not O(n^2). """ def __init__(self) -> None: self._chunks: list[str] = [] # mutable-ok: O(1) append; string concat would copy the buffer each time + self._buffer: str = ( + "" # mutable-ok: lazily materialized join of _chunks, rebuilt only when _chunks is non-empty + ) + self._offset: int = 0 # mutable-ok: cursor past already-consumed values; avoids re-slicing on every pop self._could_close: bool = False # mutable-ok: cached heuristic; rescanning past fragments was itself O(n^2) def __bool__(self) -> bool: - return bool(self._chunks) + return bool(self._chunks) or self._offset < len(self._buffer) def append(self, fragment: str) -> None: self._chunks.append(fragment) # mutable-ok: see __init__ @@ -37,34 +48,52 @@ class JSONFragmentAccumulator: """ return self._could_close + def _materialize(self) -> None: + if not self._chunks: + return + unconsumed: Final = self._buffer[self._offset :] + self._buffer = unconsumed + "".join(self._chunks) # mutable-ok: merge pending fragments, once per append batch + self._offset = 0 # mutable-ok: see __init__ + self._chunks = [] # mutable-ok: see __init__ + def pop_next_value(self) -> tuple[bool, object]: """ Attempt to decode one complete JSON value from the front of the - buffer. On success, advances the buffer past that value (keeping any - unconsumed tail, e.g. a second concatenated value) and returns - (True, value). If the buffer is empty or holds no complete value - yet, it is left untouched and this returns (False, None). + buffer. On success, advances a cursor past that value (keeping any + unconsumed tail, e.g. a second concatenated value, in place rather + than copying it) and returns (True, value). If the buffer is empty + or holds no complete value yet, it is left untouched and this + returns (False, None). """ - full: Final = "".join(self._chunks).strip() - if not full: + self._materialize() + length: Final = len(self._buffer) + start = self._offset + while start < length and self._buffer[start] in _JSON_WHITESPACE: + start += 1 + if start >= length: + self._offset = start # mutable-ok: see __init__ return False, None decoder: Final = json.JSONDecoder() try: - raw_value: Final = decoder.raw_decode(full) + raw_value: Final = decoder.raw_decode(self._buffer, start) except json.JSONDecodeError: return False, None decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode returns tuple[Any, int] - remainder: Final = full[end_index:].strip() - self._chunks = [remainder] if remainder else [] # mutable-ok: replace buffer with the unconsumed tail - if not remainder: + self._offset = end_index # mutable-ok: see __init__ + if self._offset >= len(self._buffer): + self._buffer = "" # mutable-ok: see __init__ + self._offset = 0 # mutable-ok: see __init__ self._could_close = False # mutable-ok: buffer is empty, nothing can close return True, decoded def snapshot(self) -> str: - return "".join(self._chunks) + self._materialize() + return self._buffer[self._offset :] def set(self, value: str) -> None: """Replace the buffer's contents with a single fragment.""" - self._chunks = [value] if value else [] # mutable-ok: see __init__ + self._chunks = [] # mutable-ok: see __init__ + self._buffer = value # mutable-ok: see __init__ + self._offset = 0 # mutable-ok: see __init__ stripped: Final = value.rstrip() self._could_close = bool(stripped) and stripped[-1] in ("}", "]") # mutable-ok: see __init__ diff --git a/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py b/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py index d94f2a2fa48..af141115103 100644 --- a/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py +++ b/tests/test_litellm/litellm_core_utils/test_json_fragment_accumulator.py @@ -156,6 +156,30 @@ def test_accumulation_of_many_fragments_is_not_quadratic(): assert elapsed_ms < 50, f"1000-fragment append took {elapsed_ms:.1f} ms (expected < 50 ms); O(n^2) regression?" +def test_draining_many_concatenated_values_is_not_quadratic(): + """ + Regression guard: peeling N JSON values already sitting in one buffer, + one pop_next_value() call per value with no new fragments in between, + must be O(n) total. Re-copying the shrinking remainder on every pop + (slicing a new string instead of advancing a cursor) makes total drain + time scale with the square of the buffer size. + """ + accumulator = JSONFragmentAccumulator() + accumulator.append('{"a": 1}' * 80_000) + + start = time.perf_counter() + drained = 0 + while True: + found, _ = accumulator.pop_next_value() + if not found: + break + drained += 1 + elapsed_ms = (time.perf_counter() - start) * 1000 + + assert drained == 80_000 + assert elapsed_ms < 150, f"draining 80000 concatenated values took {elapsed_ms:.1f} ms (expected < 150 ms); O(n^2) regression?" + + def test_could_close_json_after_many_blank_fragments_is_not_quadratic(): """ Regression test: a hostile upstream can send malformed JSON that never