diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 20cc5746667..1f9c017be53 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -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 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index dee3e2dfb4c..17a7dad09f0 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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 diff --git a/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py new file mode 100644 index 00000000000..893be55a700 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -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