mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
6684bb343b
* perf(streaming): add shared JSONFragmentAccumulator for Vertex and Anthropic Vertex's handle_accumulated_json_chunk and Anthropic's _handle_accumulated_json_chunk each independently accumulated SSE fragments into a JSON envelope with self.accumulated_json += fragment. Because the attribute holds a live reference, CPython copies the whole prior buffer on every fragment, making buffer assembly O(n^2) in total payload size. Anthropic additionally had no completeness heuristic at all and called json.loads on the whole buffer after every fragment, and could wedge forever on two concatenated envelopes. Add JSONFragmentAccumulator in litellm_core_utils/: fragments append to a list in O(1), a could_close_json heuristic lets callers skip the join+parse entirely until a value could plausibly be complete, and pop_next_value peels one JSON value off the front of the buffer at a time using json.JSONDecoder().raw_decode, keeping any unconsumed remainder instead of failing on concatenated values. Migrate both providers onto it; Anthropic's __next__/__anext__ end-of-stream handlers now delegate to _handle_accumulated_json_chunk(is_final=True) instead of duplicating the parse-and-reset logic inline. Fixes #31861. * test(streaming): close diff-coverage gaps in JSONFragmentAccumulator migration Codecov flagged 11 uncovered lines in the migration: the accumulated_json setter, __next__/__anext__'s end-of-stream drain branches in both providers, and the pop_next_value "not found" path when a buffer's newest fragment ends in "}" but is genuinely incomplete (an inner object closed, the outer one didn't). Add targeted tests for each. * fix(streaming): make JSONFragmentAccumulator's completeness heuristic O(1) could_close_json rescanned every preceding blank fragment on each call, so a hostile upstream that sends malformed JSON (never closing) followed by many blank keepalive fragments could drive that scan, and the join+parse it gates, to O(n^2) total. Track the last non-blank fragment's trailing byte incrementally in append/pop_next_value/set instead of rescanning the buffer. Reported by automated review on PR #36610. * fix(streaming): make JSONFragmentAccumulator.pop_next_value O(1) per call pop_next_value previously rebuilt the full remaining string and sliced a new remainder on every call, so draining N concatenated JSON values already sitting in one buffer cost O(n^2) total. Replace the rebuild-and- slice with a materialize-once cursor: pending fragments are joined into the buffer only when new ones have arrived since the last pop, and consumed values are dropped by advancing an offset instead of copying the remaining string. * test(streaming): make JSONFragmentAccumulator drain regression test CI-stable The 80k-value drain test used an absolute ms budget that flaked on a busier CI runner (233.5ms vs a 150ms budget calibrated on a quiet machine). Replace it with a doubling-ratio check: draining twice as many concatenated values should take roughly 2x as long for O(n), not ~4x for O(n^2), and that ratio holds regardless of machine speed. * test(streaming): suppress TQ002 on the append-laziness spy test A test-quality gate (TQ002: don't assert only that a mock was called) landed upstream since this branch's last rebase and now flags test_append_never_calls_raw_decode. The test verifies append() defers all decoding to pop_next_value, which has no caller-observable proxy other than spying on the stdlib call it must avoid making. * test(vertex): spy on raw_decode instead of json.loads in accumulator regression tests Post-migration to the shared JSONFragmentAccumulator, Vertex's decode path goes through json.JSONDecoder.raw_decode, not json.loads. The two O(n^2)/partial-fragment regression tests still patched json.loads, which that path never calls, so both passed unconditionally regardless of whether the underlying implementation regressed. Verified by simulating an eager-reparse regression: both tests now fail against it and pass against the correct implementation. * fix(streaming): widen JSONFragmentAccumulator's whitespace skip to match str.strip() pop_next_value's whitespace skip only matched json.decoder.WHITESPACE's ASCII set, narrower than str.strip() (Unicode-aware) which the O(1)-cursor rewrite replaced. A non-ASCII separator like U+00A0 between two concatenated JSON values on one SSE line made raw_decode fail on it, and the buffer never advanced past that byte again, permanently stranding everything after it for the rest of the stream. Use str.isspace() to match str.strip()'s tolerance. --------- Co-authored-by: Deepanshu <deepanshu.lulla@alpha-sense.com> |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| test_context_circulation.py | ||
| test_function_call_args_serialization.py | ||
| test_gemini_image_url_missing_field.py | ||
| test_gemini_streaming_tool_call_finish_reason.py | ||
| test_thought_signature_in_tool_call_id.py | ||
| test_tool_call_followed_by_text_assistant.py | ||
| test_transformation.py | ||
| test_vertex_ai_gemini_transformation.py | ||
| test_vertex_and_google_ai_studio_gemini.py | ||
| test_vertex_gemini_unbound_local_error.py | ||