scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(streaming): word-sliced cache replay for stream=true cache hits

Browse source
This commit is contained in:
michelligabriele 2026-04-15 21:26:58 +02:00
parent 72a461ba4a
commit 199c7a3429
No known key found for this signature in database
3 changed files with 263 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

108
litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py
View file

@ -1,5 +1,6 @@
import asyncio
import json
import re
import time
import traceback
from typing import Dict, Iterable, List, Literal, Optional, Tuple, Union, cast
@ -118,7 +119,29 @@ def convert_tool_call_to_json_mode(
return None, None
async def convert_to_streaming_response_async(response_object: Optional[dict] = None):
# Whitespace-preserving word splitter used by the cache-hit replay generators.
# Each match is any leading whitespace plus a non-whitespace run plus any
# trailing whitespace, so concatenating the matches losslessly reconstructs
# the original string (including content that starts with whitespace).
_REPLAY_CONTENT_SLICE_RE = re.compile(r"\s*\S+\s*", re.UNICODE)
def _split_assembled_content_for_replay(content: Optional[str]) -> List[str]:
"""
Slice an assembled cached completion's ``content`` into word-shaped pieces
for cadence-preserving streaming replay. The split is lossless:
``"".join(_split_assembled_content_for_replay(s)) == s`` for every
non-empty ``s``. Returns ``[]`` for ``None`` / empty / all-whitespace
content.
"""
if not content:
return []
return _REPLAY_CONTENT_SLICE_RE.findall(content)
async def convert_to_streaming_response_async( # noqa: PLR0915
response_object: Optional[dict] = None,
):
"""
Asynchronously converts a response object to a streaming response.
@ -202,8 +225,42 @@ async def convert_to_streaming_response_async(response_object: Optional[dict] =
if "model" in response_object:
model_response_object.model = response_object["model"]
yield model_response_object
await asyncio.sleep(0)
# Replay cached content with per-word cadence so stream=true cache hits
# don't arrive as a single SSE frame. Multi-choice (n>1) responses and
# unsplittable content (None/empty/whitespace-free) keep the original
# single-yield behavior.
slices: List[str] = []
if len(model_response_object.choices) == 1:
slices = _split_assembled_content_for_replay(
model_response_object.choices[0].delta.content
)
if len(slices) <= 1:
yield model_response_object
await asyncio.sleep(0)
return
# Detach usage from the base object so we can re-attach it only to the
# final slice chunk.
original_usage = getattr(model_response_object, "usage", None)
if original_usage is not None:
try:
delattr(model_response_object, "usage")
except AttributeError:
pass
original_finish_reason = model_response_object.choices[0].finish_reason
last_idx = len(slices) - 1
for i, piece in enumerate(slices):
slice_chunk = model_response_object.model_copy(deep=True)
slice_chunk.choices[0].delta.content = piece
if i > 0:
slice_chunk.choices[0].delta.role = None
slice_chunk.choices[0].finish_reason = (
original_finish_reason if i == last_idx else None
)
if i == last_idx and original_usage is not None:
setattr(slice_chunk, "usage", original_usage)
yield slice_chunk
await asyncio.sleep(0)
def convert_to_streaming_response(response_object: Optional[dict] = None):
@ -251,7 +308,38 @@ def convert_to_streaming_response(response_object: Optional[dict] = None):
if "model" in response_object:
model_response_object.model = response_object["model"]
yield model_response_object
# Replay cached content with per-word cadence on sync cache-hit paths
# (S3Cache, sync completion()). See convert_to_streaming_response_async
# for the full rationale — this mirrors its tail.
slices: List[str] = []
if len(model_response_object.choices) == 1:
slices = _split_assembled_content_for_replay(
model_response_object.choices[0].delta.content
)
if len(slices) <= 1:
yield model_response_object
return
original_usage = getattr(model_response_object, "usage", None)
if original_usage is not None:
try:
delattr(model_response_object, "usage")
except AttributeError:
pass
original_finish_reason = model_response_object.choices[0].finish_reason
last_idx = len(slices) - 1
for i, piece in enumerate(slices):
slice_chunk = model_response_object.model_copy(deep=True)
slice_chunk.choices[0].delta.content = piece
if i > 0:
slice_chunk.choices[0].delta.role = None
slice_chunk.choices[0].finish_reason = (
original_finish_reason if i == last_idx else None
)
if i == last_idx and original_usage is not None:
setattr(slice_chunk, "usage", original_usage)
yield slice_chunk
from collections import defaultdict
@ -596,9 +684,9 @@ def convert_to_model_response_object( # noqa: PLR0915
provider_specific_fields["thinking_blocks"] = thinking_blocks
if reasoning_content:
provider_specific_fields[
"reasoning_content"
] = reasoning_content
provider_specific_fields["reasoning_content"] = (
reasoning_content
)
message = Message(
content=content,
@ -787,9 +875,9 @@ def convert_to_model_response_object( # noqa: PLR0915
# tracking without exposing it in the response body. Must be set
# after hidden_params assignment to avoid being overwritten.
if "_audio_transcription_duration" in response_object:
model_response_object._hidden_params[
"audio_transcription_duration"
] = response_object["_audio_transcription_duration"]
model_response_object._hidden_params["audio_transcription_duration"] = (
response_object["_audio_transcription_duration"]
)
if _response_headers is not None:
model_response_object._response_headers = _response_headers

5
litellm/litellm_core_utils/streaming_handler.py
View file

@ -1369,10 +1369,11 @@ class CustomStreamWrapper:
self.received_finish_reason = response_obj["finish_reason"]
elif self.custom_llm_provider == "cached_response":
chunk = cast(ModelResponseStream, chunk)
chunk_finish_reason = chunk.choices[0].finish_reason
response_obj = {
"text": chunk.choices[0].delta.content,
"is_finished": True,
"finish_reason": chunk.choices[0].finish_reason,
"is_finished": chunk_finish_reason is not None,
"finish_reason": chunk_finish_reason,
"original_chunk": chunk,
"tool_calls": (
chunk.choices[0].delta.tool_calls

162
tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py Normal file
View file

@ -0,0 +1,162 @@
"""
Tests for the cache-hit replay generators in
``litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response``.
These generators are used by ``LLMCachingHandler._convert_cached_stream_response``
to replay a cached non-streaming ``ModelResponse`` as a stream when the
incoming request has ``stream=True``. The fix in this test file ensures the
replay yields multiple word-shaped chunks instead of a single one-shot
content frame, restoring per-token cadence on cache hits (case
2026-04-13-pramod-streaming-buffered-subsequent-requests).
"""
import pytest
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
_split_assembled_content_for_replay,
convert_to_streaming_response,
convert_to_streaming_response_async,
)
from litellm.types.utils import ModelResponseStream
def _async_payload(content="Hello world! How are you?"):
return {
"id": "chatcmpl-test-async",
"object": "chat.completion",
"created": 1700000000,
"model": "gpt-4o-mini",
"system_fingerprint": "fp_test",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": content},
"finish_reason": "stop",
}
],
"usage": {
"prompt_tokens": 5,
"completion_tokens": 7,
"total_tokens": 12,
},
}
async def _collect_async(payload):
return [chunk async for chunk in convert_to_streaming_response_async(payload)]
# ---------- helper ----------
@pytest.mark.parametrize(
"text",
[
"Hello world!",
"Sure! Here's a list of 25 fruits:\n\n1. Apple\n2. Banana\n3. Orange\n",
" leading whitespace matters",
"Hi",
"你好世界",
],
)
def test_split_is_lossless(text):
assert "".join(_split_assembled_content_for_replay(text)) == text
def test_split_returns_empty_for_none_and_empty():
assert _split_assembled_content_for_replay(None) == []
assert _split_assembled_content_for_replay("") == []
# ---------- async generator ----------
@pytest.mark.asyncio
async def test_async_yields_multiple_content_chunks_with_lossless_join():
text = "Sure! Here's a list of fruits: apple, banana, orange."
chunks = await _collect_async(_async_payload(content=text))
assert len(chunks) > 1
assert all(isinstance(c, ModelResponseStream) for c in chunks)
reassembled = "".join((c.choices[0].delta.content or "") for c in chunks)
assert reassembled == text
@pytest.mark.asyncio
async def test_async_finish_reason_only_on_last_chunk():
chunks = await _collect_async(_async_payload())
finish_reasons = [c.choices[0].finish_reason for c in chunks]
assert finish_reasons[-1] == "stop"
assert all(fr is None for fr in finish_reasons[:-1])
@pytest.mark.asyncio
async def test_async_role_only_on_first_chunk():
chunks = await _collect_async(_async_payload())
assert chunks[0].choices[0].delta.role == "assistant"
for c in chunks[1:]:
assert c.choices[0].delta.role is None
@pytest.mark.asyncio
async def test_async_usage_attached_to_last_chunk_only():
chunks = await _collect_async(_async_payload())
usage_frames = [c for c in chunks if getattr(c, "usage", None) is not None]
assert len(usage_frames) == 1
assert usage_frames[0] is chunks[-1]
assert usage_frames[0].usage.completion_tokens == 7
assert usage_frames[0].usage.prompt_tokens == 5
assert usage_frames[0].usage.total_tokens == 12
assert usage_frames[0].choices[0].finish_reason == "stop"
@pytest.mark.asyncio
async def test_async_empty_content_yields_single_finish_frame():
# tool-calls-only-style response: short-circuit to the single-yield path.
payload = _async_payload(content=None)
payload["usage"] = None
chunks = await _collect_async(payload)
assert len(chunks) == 1
assert chunks[0].choices[0].delta.content is None
assert chunks[0].choices[0].finish_reason == "stop"
assert chunks[0].choices[0].delta.role == "assistant"
assert getattr(chunks[0], "usage", None) is None
@pytest.mark.asyncio
async def test_async_metadata_propagated_to_every_chunk():
chunks = await _collect_async(_async_payload())
for c in chunks:
assert c.id == "chatcmpl-test-async"
assert c.model == "gpt-4o-mini"
assert c.system_fingerprint == "fp_test"
assert c.created == 1700000000
# ---------- sync generator (parity smoke test) ----------
def test_sync_multi_chunk_and_lossless_join():
text = "Sure! Here's a list of fruits: apple, banana, orange."
payload = {
"id": "chatcmpl-test-sync",
"object": "chat.completion",
"created": 1700000000,
"model": "gpt-4o-mini",
"system_fingerprint": "fp_test",
"choices": [
{
"index": 0,
"message": {"role": "assistant", "content": text},
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12},
}
chunks = list(convert_to_streaming_response(payload))
assert len(chunks) > 1
reassembled = "".join((c.choices[0].delta.content or "") for c in chunks)
assert reassembled == text
# Sync path must also honor the "usage on last chunk" invariant.
assert chunks[-1].choices[0].finish_reason == "stop"
assert getattr(chunks[-1], "usage", None) is not None
assert chunks[-1].usage.total_tokens == 12