mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
fix(proxy): recognize crlf sse frame boundaries in passthrough reassembly
This commit is contained in:
parent
46fb1cd514
commit
938396ef90
3 changed files with 24 additions and 4 deletions
|
|
@ -3066,7 +3066,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
elif isinstance(chunk, (bytes, bytearray)):
|
||||
try:
|
||||
s: Final = chunk.decode("utf-8")
|
||||
if s.endswith("\n\n"):
|
||||
if s.endswith(("\n\n", "\r\n\r\n")):
|
||||
maybe_mod = ProxyBaseLLMRequestProcessing._inject_cost_into_sse_frame_str(s, model_name)
|
||||
if maybe_mod is not None:
|
||||
return maybe_mod.encode("utf-8")
|
||||
|
|
|
|||
|
|
@ -141,10 +141,12 @@ class PassThroughStreamingHandler:
|
|||
|
||||
@staticmethod
|
||||
def _split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]:
|
||||
frame_boundary: Final = pending.rfind(b"\n\n")
|
||||
if frame_boundary == -1:
|
||||
lf_boundary_end: Final = pending.rfind(b"\n\n") + 2
|
||||
crlf_boundary_end: Final = pending.rfind(b"\r\n\r\n") + 4
|
||||
boundary_end: Final = max(lf_boundary_end if lf_boundary_end >= 2 else 0, crlf_boundary_end if crlf_boundary_end >= 4 else 0)
|
||||
if boundary_end == 0:
|
||||
return b"", pending
|
||||
return pending[: frame_boundary + 2], pending[frame_boundary + 2 :]
|
||||
return pending[:boundary_end], pending[boundary_end:]
|
||||
|
||||
@staticmethod
|
||||
async def _route_streaming_logging_to_handler(
|
||||
|
|
|
|||
|
|
@ -443,6 +443,24 @@ async def test_chunk_processor_injects_cost_into_usage_frame_fragmented_across_c
|
|||
assert reassembled.endswith("data: [DONE]\n\n")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_streams_crlf_delimited_frames_live_and_injects_cost(monkeypatch):
|
||||
"""Regression: CRLF-delimited SSE frames must flow as they complete instead of
|
||||
buffering until EOF, and the usage frame must still get cost injected."""
|
||||
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
|
||||
chunks = [chunk.replace(b"\n\n", b"\r\n\r\n") for chunk in _openai_passthrough_stream_chunks()]
|
||||
|
||||
received = await _collect_openai_passthrough_chunks(chunks, EndpointType.OPENAI)
|
||||
|
||||
assert len(received) == len(chunks)
|
||||
assert received[0] == chunks[0]
|
||||
reassembled = b"".join(received).decode("utf-8")
|
||||
usage_lines = [ln for ln in reassembled.replace("\r\n", "\n").split("\n") if '"total_tokens"' in ln]
|
||||
assert len(usage_lines) == 1
|
||||
final_payload = json.loads(usage_lines[0].split("data:", 1)[1].strip())
|
||||
assert final_payload["usage"]["cost"] > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_chunk_processor_flag_off_leaves_openai_passthrough_stream_byte_identical(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue