scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-30 06:43:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 249

fix(caching): keep reasoning, citations, and audio on streamed cache hits (#32096)

Browse source 
* fix(proxy): preserve reasoning fields on streaming cache hits

* fix(proxy): validate cached streaming delta fields

Signed-off-by: jesco-absolut <team@srswti.com>

* fix(caching): replay cached audio, images, reasoning and citations in stream order

A cached answer replayed as a stream now ends with its own empty finish chunk, so a chunk that
carries only audio or images keeps its payload and role. Reasoning fields lead on a chunk with no
text, and annotations arrive once on the last text slice, matching a live stream's order.

* fix(caching): replay a cached tool call after the text, like a live stream

Only the role, reasoning fields and provider fields lead a replayed answer. Tool calls now arrive on
their own chunk after the last text slice, and logprobs ride the first text chunk, so a cached
answer with reasoning and a tool call keeps a live stream's reasoning, text, tool call order.

* fix(caching): replay the reasoning before a cached tool call that has no text

A cached answer that thinks and then only calls a tool now replays its reasoning on its own chunk before the
tool call, the order a live stream sends, instead of one chunk carrying both.

* fix(caching): keep blank-only text when replaying a cached answer with reasoning

A cached answer that thinks, sends only whitespace as text, and then calls a tool now replays that whitespace
between the reasoning and the tool call, instead of dropping it.

---------

Signed-off-by: jesco-absolut <team@srswti.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
jesco 2026-10-29 13:08:12 -04:00 • committed by GitHub
parent a659e8be70
commit 3057c47d6d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 596 additions and 150 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -3,7 +3,7 @@ import json
import re
import time
import traceback
from collections.abc import Mapping, Sequence
from collections.abc import Generator, Iterator, Mapping, Sequence
from typing import Final, Literal, cast
from pydantic import ConfigDict, TypeAdapter, ValidationError, with_config
@ -20,7 +20,6 @@ from litellm.llms.ollama.common_utils import OllamaResponseIdNormalizer
from litellm.types.llms.databricks import DatabricksTool
from litellm.types.llms.openai import (
ChatCompletionAnnotation,
ChatCompletionReasoningItem,
ChatCompletionThinkingBlock,
ImageURLListItem,
OpenAIModerationResponse,
@ -171,31 +170,89 @@ def _split_assembled_content_for_replay(content: str | None) -> list[str]:
return _REPLAY_CONTENT_SLICE_RE.findall(content)
def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None:
# Rebuild the delta as content-only so every accumulate-able field (role,
# tool_calls, reasoning_content, thinking_blocks, audio, images,
# annotations, ...) is dropped on later slices instead of an enumerated
# subset; repeating any of them makes downstream handlers that accumulate
# streamed deltas collect it once per slice, and a field added to Delta
# later can't silently re-introduce the duplication.
choice.delta = Delta(content=choice.delta.content)
choice.logprobs = None
if hasattr(choice, "enhancements"):
del choice.enhancements
_REPLAY_REASONING_FIELDS: Final = frozenset({"reasoning_content", "thinking_blocks", "reasoning_items"})
_REPLAY_LEAD_FIELDS: Final = _REPLAY_REASONING_FIELDS | {"role", "provider_specific_fields"}
_REPLAY_TOOL_FIELDS: Final = frozenset({"tool_calls", "function_call"})
_REPLAY_TEXT_FIELDS: Final = frozenset({"content", "annotations"})
def _first_replay_delta(
message: Mapping[str, object], reasoning_items: Sequence[ChatCompletionReasoningItem] | None
) -> Delta:
return Delta(
content=message.get("content"),
role=message["role"],
function_call=message.get("function_call"),
tool_calls=message.get("tool_calls"),
reasoning_items=None if reasoning_items is None else list(reasoning_items),
provider_specific_fields=_cached_provider_specific_fields(
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message)
),
def _fields_in(fields: Mapping[str, object], keys: frozenset[str]) -> dict[str, object]:
return {key: value for key, value in fields.items() if key in keys}
def _replay_deltas(delta: Delta) -> tuple[Delta, ...]:
fields: Final[dict[str, object]] = delta.model_dump(exclude_none=True)
lead: Final = _fields_in(fields, _REPLAY_LEAD_FIELDS) if _REPLAY_REASONING_FIELDS.intersection(fields) else {}
pieces: Final = _split_assembled_content_for_replay(delta.content) or ([delta.content] if delta.content else [])
if not pieces and not lead:
return (delta,)
if not pieces:
held_back: Final = lead.keys() | {"content"}
rest: Final = {key: value for key, value in fields.items() if key not in held_back}
return (Delta.model_validate(lead),) + ((Delta.model_validate(rest),) if rest else ())
tools: Final = _fields_in(fields, _REPLAY_TOOL_FIELDS)
annotations: Final = _fields_in(fields, frozenset({"annotations"}))
moved: Final = lead.keys() | _REPLAY_TOOL_FIELDS | _REPLAY_TEXT_FIELDS
first_slice: Final = {key: value for key, value in fields.items() if key not in moved}
last_index: Final = len(pieces) - 1
return (
((Delta.model_validate(lead),) if lead else ())
+ tuple(
Delta.model_validate(
{
**(first_slice if index == 0 else {}),
"content": piece,
**(annotations if index == last_index else {}),
}
)
for index, piece in enumerate(pieces)
)
+ ((Delta.model_validate(tools),) if tools else ())
)
def _replay_frame(
response: ModelResponseStream, delta: Delta, finish_reason: str | None, keeps_choice_metadata: bool
) -> ModelResponseStream:
frame: Final = response.model_copy(deep=True)
choice: Final = frame.choices[0]
choice.delta = delta
choice.finish_reason = finish_reason
if not keeps_choice_metadata:
choice.logprobs = None
if hasattr(choice, "enhancements"):
del choice.enhancements
return frame
def _replay_frames(response: ModelResponseStream) -> Iterator[ModelResponseStream]:
if len(response.choices) != 1:
yield response
return
usage: Final = getattr(response, "usage", None)
if usage is not None:
delattr(response, "usage")
choice: Final = response.choices[0]
deltas: Final = _replay_deltas(choice.delta)
first_text: Final = next((index for index, delta in enumerate(deltas) if delta.content), 0)
for index, delta in enumerate(deltas):
yield _replay_frame(response, delta, finish_reason=None, keeps_choice_metadata=index == first_text)
finish_frame: Final = _replay_frame(
response, Delta(), finish_reason=choice.finish_reason, keeps_choice_metadata=False
)
if usage is not None:
setattr(finish_frame, "usage", usage)
yield finish_frame
def _first_replay_delta(message: Mapping[str, object]) -> Delta:
return Delta.model_validate(
{
**message,
"provider_specific_fields": _cached_provider_specific_fields(
_CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message)
),
}
)
@ -260,7 +317,7 @@ async def convert_to_streaming_response_async(
t["index"] = index
pydantic_tool_calls.append(ChatCompletionDeltaToolCall(**t))
choice["message"]["tool_calls"] = pydantic_tool_calls
delta = _first_replay_delta(choice["message"], choice["message"].get("reasoning_items"))
delta = _first_replay_delta(choice["message"])
finish_reason = choice.get("finish_reason", None)
if finish_reason is None:
@ -292,41 +349,14 @@ async def convert_to_streaming_response_async(
if "model" in response_object:
model_response_object.model = response_object["model"]
# 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. A non-None usage always lives in __pydantic_extra__
# here (set via setattr above), so delattr cannot fail.
original_usage: Final = getattr(model_response_object, "usage", None)
if original_usage is not None:
delattr(model_response_object, "usage")
original_finish_reason: Final = model_response_object.choices[0].finish_reason
last_idx: Final = 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:
_clear_later_replay_slice_metadata(slice_chunk.choices[0])
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
for frame in _replay_frames(model_response_object):
yield frame
await asyncio.sleep(0)
def convert_to_streaming_response(
response_object: dict | None = None,
):
) -> Generator[ModelResponseStream, None, None]:
# used for yielding Cache hits when stream == True
if response_object is None:
raise Exception("Error in response object format")
@ -382,30 +412,7 @@ def convert_to_streaming_response(
if "model" in response_object:
model_response_object.model = response_object["model"]
# 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: Final = getattr(model_response_object, "usage", None)
if original_usage is not None:
delattr(model_response_object, "usage")
original_finish_reason: Final = model_response_object.choices[0].finish_reason
last_idx: Final = 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:
_clear_later_replay_slice_metadata(slice_chunk.choices[0])
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
yield from _replay_frames(model_response_object)
from collections import defaultdict

8
tests/integration/caching/response_cache_case.py
View file

@ -275,11 +275,11 @@ def prompt() -> str:
return f"cache probe {uuid.uuid4().hex}"
def redis_store() -> Redis:
def redis_store() -> Redis[bytes]:
return Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
def entry_keys(store: Redis, needle: str) -> Iterator[str]:
def entry_keys(store: Redis[bytes], needle: str) -> Iterator[str]:
for raw_key in store.scan_iter(count=1000):
key: Final = raw_key.decode()
if not CACHE_KEY.match(key) or store.type(key) != b"string":
@ -289,12 +289,12 @@ def entry_keys(store: Redis, needle: str) -> Iterator[str]:
yield key
def entry_key(store: Redis, needle: str) -> str | None:
def entry_key(store: Redis[bytes], needle: str) -> str | None:
"""The response cache key (a bare sha256) whose stored value names ``needle``, or ``None``."""
return next(entry_keys(store, needle), None)
def await_entry(store: Redis, needle: str) -> str:
def await_entry(store: Redis[bytes], needle: str) -> str:
key: Final = eventually(lambda: entry_key(store, needle), lambda found: found is not None)
assert key is not None
return key

208
tests/integration/caching/test_response_cache_stream_replay.py Normal file
View file

@ -0,0 +1,208 @@
"""A chat stream served from the response cache carries the reasoning, citations, audio and tool calls it first had.
Every cell scripts one chat stream for its own deployment, streams it once through the OpenAI SDK so the proxy
assembles and caches it, then streams the identical request again until an answer arrives without an upstream call
and reads the replayed deltas the way a client would.
"""
from __future__ import annotations
import json
from collections.abc import Iterator
from typing import Final
import pytest
from openai import OpenAI
from pydantic import JsonValue
from tests.integration._support.client import Gateway
from tests.integration.caching.response_cache_case import (
CHAT_CHUNK,
Upstream,
await_entry,
prompt,
redis_store,
scripted,
)
from tests.integration.cost_calculation.cost_tracking_case import SseResponse
REPLAY_ATTEMPTS: Final = 8
CITATION: Final = (
'{"type":"url_citation","url_citation":{"url":"https://example.com/apples","title":"Apples",'
'"start_index":0,"end_index":19}}'
)
STOP: Final = CHAT_CHUNK + '[{"index":0,"delta":{},"finish_reason":"stop"}]}'
REASONING_STREAM: Final = SseResponse(
content_type="text/event-stream",
frames=(
CHAT_CHUNK + '[{"index":0,"delta":{"role":"assistant","reasoning_content":"Count the apples "},'
'"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"reasoning_content":"in each basket"},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"content":"There are "},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"content":"14 apples"},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"annotations":[' + CITATION + ']},"finish_reason":null}]}',
STOP,
"data: [DONE]",
),
)
AUDIO_STREAM: Final = SseResponse(
content_type="text/event-stream",
frames=(
CHAT_CHUNK + '[{"index":0,"delta":{"role":"assistant","audio":{"id":"audio_$UNIQUE_ID","transcript":"Hello"}},'
'"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"audio":{"data":"UklGRg=="}},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"audio":{"expires_at":1700000600}},"finish_reason":null}]}',
STOP,
"data: [DONE]",
),
)
TOOL_CALL: Final = (
'{"index":0,"id":"call_$UNIQUE_ID","type":"function","function":{"name":"count_apples","arguments":"{}"}}'
)
TOOL_STREAM: Final = SseResponse(
content_type="text/event-stream",
frames=(
CHAT_CHUNK + '[{"index":0,"delta":{"role":"assistant","reasoning_content":"I need the basket totals"},'
'"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"content":"Let me count "},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"content":"the apples"},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"tool_calls":[' + TOOL_CALL + ']},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}',
"data: [DONE]",
),
)
TOOL_ONLY_STREAM: Final = SseResponse(
content_type="text/event-stream",
frames=(
CHAT_CHUNK + '[{"index":0,"delta":{"role":"assistant","reasoning_content":"I need the basket totals"},'
'"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"tool_calls":[' + TOOL_CALL + ']},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}',
"data: [DONE]",
),
)
WHITESPACE_TOOL_STREAM: Final = SseResponse(
content_type="text/event-stream",
frames=(
CHAT_CHUNK + '[{"index":0,"delta":{"role":"assistant","reasoning_content":"I need the basket totals"},'
'"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"content":"\\n\\n"},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{"tool_calls":[' + TOOL_CALL + ']},"finish_reason":null}]}',
CHAT_CHUNK + '[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}',
"data: [DONE]",
),
)
_Stream = tuple[frozenset[str], tuple[dict[str, JsonValue], ...]]
def _stream(gateway: Gateway, model: str, text: str) -> _Stream:
"""The chunk ids and every choice's delta with its finish reason, after the whole stream was read."""
with (
OpenAI(
base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1", api_key=gateway.key, max_retries=0
) as client,
client.chat.completions.create(
model=model, messages=[{"role": "user", "content": text}], stream=True
) as stream,
):
chunks: Final = tuple(stream)
deltas: Final = tuple(
{**choice.delta.model_dump(exclude_none=True), "finish_reason": choice.finish_reason}
for chunk in chunks
for choice in chunk.choices
) # comprehension-ok: flattens chunk choices
return frozenset(chunk.id for chunk in chunks), deltas
def _attempt(gateway: Gateway, upstream: Upstream, identity: str, model: str, text: str) -> _Stream | None:
before: Final = upstream.calls(identity)
answer: Final = _stream(gateway, model, text)
return answer if upstream.calls(identity) == before else None
def _attempts(gateway: Gateway, upstream: Upstream, identity: str, model: str, text: str) -> Iterator[_Stream | None]:
return (_attempt(gateway, upstream, identity, model, text) for _ in range(REPLAY_ATTEMPTS))
def _replayed(gateway: Gateway, script: SseResponse) -> tuple[dict[str, JsonValue], ...]:
"""The deltas of the first identical stream served without an upstream call, after one stream filled the cache.
A streamed request can reach the upstream once per proxy worker before a replay is served, since a worker keys a
stream only once its router resolves the model the other worker registered.
"""
upstream: Final = Upstream(gateway.upstream_url)
text: Final = prompt()
with gateway.scenario() as scenario:
handle: Final = scripted(scenario, script)
model: Final = scenario.model(api_base=handle.api_base())
first_ids, _ = _stream(gateway, model, text)
assert upstream.calls(handle.scenario_id) == 1
await_entry(redis_store(), next(iter(first_ids)))
replay: Final = next(
(answer for answer in _attempts(gateway, upstream, handle.scenario_id, model, text) if answer), None
)
if replay is None:
pytest.fail(f"no stream was replayed from the cache in {REPLAY_ATTEMPTS} attempts")
return replay[1]
def _joined(deltas: tuple[dict[str, JsonValue], ...], field: str) -> str:
return "".join(str(delta[field]) for delta in deltas if delta.get(field))
def test_a_replayed_chat_stream_sends_the_reasoning_before_the_text_and_the_citation_after_it(
gateway: Gateway,
) -> None:
deltas: Final = _replayed(gateway, REASONING_STREAM)
reasoning: Final = [index for index, delta in enumerate(deltas) if delta.get("reasoning_content")]
text: Final = [index for index, delta in enumerate(deltas) if delta.get("content")]
cited: Final = [index for index, delta in enumerate(deltas) if delta.get("annotations")]
assert _joined(deltas, "reasoning_content") == "Count the apples in each basket", deltas
assert _joined(deltas, "content") == "There are 14 apples", deltas
assert reasoning and text and max(reasoning) < min(text), deltas
assert cited == [max(text)] and deltas[max(text)]["annotations"] == [json.loads(CITATION)], deltas
assert deltas[-1]["finish_reason"] == "stop", deltas
def test_a_replayed_chat_stream_without_text_sends_its_role_and_audio(gateway: Gateway) -> None:
deltas: Final = _replayed(gateway, AUDIO_STREAM)
audio: Final = [delta["audio"] for delta in deltas if delta.get("audio")]
assert deltas[0].get("role") == "assistant", deltas
assert len(audio) == 1, deltas
assert isinstance(audio[0], dict) and audio[0]["data"] == "UklGRg==" and audio[0]["transcript"] == "Hello", deltas
assert deltas[-1]["finish_reason"] == "stop", deltas
def test_a_replayed_chat_stream_sends_the_tool_call_after_the_reasoning_and_the_text(gateway: Gateway) -> None:
deltas: Final = _replayed(gateway, TOOL_STREAM)
reasoning: Final = [index for index, delta in enumerate(deltas) if delta.get("reasoning_content")]
text: Final = [index for index, delta in enumerate(deltas) if delta.get("content")]
called: Final = [index for index, delta in enumerate(deltas) if delta.get("tool_calls")]
assert _joined(deltas, "content") == "Let me count the apples", deltas
assert reasoning and text and max(reasoning) < min(text), deltas
assert called == [max(text) + 1], deltas
assert deltas[-1]["finish_reason"] == "tool_calls", deltas
def test_a_replayed_chat_stream_without_text_sends_the_reasoning_before_the_tool_call(gateway: Gateway) -> None:
deltas: Final = _replayed(gateway, TOOL_ONLY_STREAM)
reasoning: Final = [index for index, delta in enumerate(deltas) if delta.get("reasoning_content")]
called: Final = [index for index, delta in enumerate(deltas) if delta.get("tool_calls")]
assert _joined(deltas, "reasoning_content") == "I need the basket totals", deltas
assert reasoning and called and max(reasoning) < min(called), deltas
assert deltas[-1]["finish_reason"] == "tool_calls", deltas
def test_a_replayed_chat_stream_keeps_whitespace_only_text_between_the_reasoning_and_the_tool_call(
gateway: Gateway,
) -> None:
deltas: Final = _replayed(gateway, WHITESPACE_TOOL_STREAM)
reasoning: Final = [index for index, delta in enumerate(deltas) if delta.get("reasoning_content")]
text: Final = [index for index, delta in enumerate(deltas) if delta.get("content")]
called: Final = [index for index, delta in enumerate(deltas) if delta.get("tool_calls")]
assert _joined(deltas, "content") == "\n\n", deltas
assert reasoning and text and called and max(reasoning) < min(text) and max(text) < min(called), deltas
assert deltas[-1]["finish_reason"] == "tool_calls", deltas

132
tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py
View file

@ -2287,17 +2287,19 @@ class TestConvertToStreamingResponse:
}
chunks = list(convert_to_streaming_response(response_object=response_object))
assert len(chunks) == 1
chunk = chunks[0]
assert chunk.id == "chatcmpl-123"
assert chunk.model == "gpt-4"
assert chunk.created == 1700000000
assert chunk.system_fingerprint == "fp_abc"
assert chunk.choices[0].delta.content == "Hello!"
assert chunk.choices[0].delta.role == "assistant"
assert chunk.choices[0].finish_reason == "stop"
assert chunk.usage.prompt_tokens == 5
assert chunk.usage.completion_tokens == 2
assert len(chunks) == 2
assert all(chunk.id == "chatcmpl-123" for chunk in chunks)
assert all(chunk.model == "gpt-4" for chunk in chunks)
assert all(chunk.created == 1700000000 for chunk in chunks)
assert all(chunk.system_fingerprint == "fp_abc" for chunk in chunks)
content, finish = chunks
assert content.choices[0].delta.content == "Hello!"
assert content.choices[0].delta.role == "assistant"
assert content.choices[0].finish_reason is None
assert finish.choices[0].delta.content is None
assert finish.choices[0].finish_reason == "stop"
assert finish.usage.prompt_tokens == 5
assert finish.usage.completion_tokens == 2
def test_finish_details_fallback(self):
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
@ -2315,7 +2317,7 @@ class TestConvertToStreamingResponse:
}
chunks = list(convert_to_streaming_response(response_object=response_object))
assert chunks[0].choices[0].finish_reason == "length"
assert chunks[-1].choices[0].finish_reason == "length"
def test_tool_calls_in_streaming(self):
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
@ -2347,17 +2349,20 @@ class TestConvertToStreamingResponse:
}
async def run():
chunks = []
async for chunk in convert_to_streaming_response_async(
response_object=response_object
):
chunks.append(chunk)
return chunks
return [
chunk
async for chunk in convert_to_streaming_response_async(
response_object=response_object
)
]
chunks = asyncio.run(run())
assert len(chunks) == 1
assert len(chunks) == 2
assert chunks[0].choices[0].delta.tool_calls[0].id == "call_1"
assert chunks[0].choices[0].delta.tool_calls[0].function.name == "get_weather"
assert chunks[0].choices[0].finish_reason is None
assert chunks[1].choices[0].delta.tool_calls is None
assert chunks[1].choices[0].finish_reason == "tool_calls"
class TestConvertToStreamingResponseAsync:
def test_none_raises(self):
@ -2395,17 +2400,17 @@ class TestConvertToStreamingResponseAsync:
}
async def run():
chunks = []
async for chunk in convert_to_streaming_response_async(
response_object=response_object
):
chunks.append(chunk)
return chunks
return [
chunk
async for chunk in convert_to_streaming_response_async(
response_object=response_object
)
]
chunks = asyncio.run(run())
# Cached replay is sliced into word-shaped chunks to preserve
# streaming cadence; joining the slices reconstructs the content.
assert len(chunks) == 2
assert len(chunks) == 3
assert all(c.id == "msg_async_1" for c in chunks)
assert all(c.model == "claude-3" for c in chunks)
assert "".join(c.choices[0].delta.content or "" for c in chunks) == "Hi there"
@ -2413,6 +2418,81 @@ class TestConvertToStreamingResponseAsync:
assert chunks[-1].choices[0].finish_reason == "stop"
assert chunks[-1].usage.prompt_tokens == 3
def test_preserves_reasoning_fields(self):
import asyncio
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
convert_to_streaming_response_async,
)
thinking_blocks = [
{
"type": "thinking",
"thinking": "cached reasoning",
"signature": "sig-cache",
}
]
annotations = [
{
"type": "url_citation",
"url_citation": {
"start_index": 0,
"end_index": 5,
"title": "Cached source",
"url": "https://example.com/cached-source",
},
}
]
response_object = {
"id": "msg_async_reasoning_cache",
"model": "claude-3",
"created": 1700000000,
"choices": [
{
"finish_reason": "stop",
"index": 0,
"message": {
"content": "Final answer",
"role": "assistant",
"reasoning_content": "cached reasoning",
"thinking_blocks": thinking_blocks,
"annotations": annotations,
},
}
],
}
async def run():
return [
chunk
async for chunk in convert_to_streaming_response_async(
response_object=response_object
)
]
chunks = asyncio.run(run())
assert (
"".join(c.choices[0].delta.content or "" for c in chunks) == "Final answer"
)
assert chunks[0].choices[0].delta.content is None
assert chunks[0].choices[0].delta.reasoning_content == "cached reasoning"
assert chunks[0].choices[0].delta.thinking_blocks == thinking_blocks
last_content = max(
i for i, c in enumerate(chunks) if c.choices[0].delta.content
)
assert chunks[last_content].choices[0].delta.annotations == annotations
assert [
i for i, c in enumerate(chunks) if c.choices[0].delta.get("annotations")
] == [last_content]
assert not any(
hasattr(c.choices[0].delta, "reasoning_content") for c in chunks[1:]
)
assert not any(
hasattr(c.choices[0].delta, "thinking_blocks") for c in chunks[1:]
)
class TestHandleInvalidParallelToolCalls:
def test_none_input(self):
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (

231
tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py
View file

@ -9,9 +9,11 @@ replay yields multiple word-shaped chunks instead of a single one-shot
content frame, restoring per-token cadence on cache hits.
"""
from typing import Final
from datetime import datetime
from typing import Final, Literal
import pytest
from pydantic import JsonValue
import litellm
from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import (
@ -19,6 +21,7 @@ from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response impo
convert_to_streaming_response,
convert_to_streaming_response_async,
)
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.utils import ModelResponseStream
@ -267,23 +270,20 @@ async def test_async_usage_attached_to_last_chunk_only():
@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.
async def test_async_empty_content_yields_the_role_then_a_separate_finish_frame():
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
assert [(c.choices[0].delta.role, c.choices[0].finish_reason) for c in chunks] == [
("assistant", None),
(None, "stop"),
]
assert all(c.choices[0].delta.content is None for c in chunks)
assert all(getattr(c, "usage", None) is None for c in chunks)
@pytest.mark.asyncio
async def test_async_tool_calls_and_function_call_only_on_first_chunk():
# A cached response combining multi-word content with tool_calls must not
# repeat the tool_calls on every slice — downstream handlers accumulate
# tool-call deltas and would collect them N times.
async def test_async_tool_calls_arrive_once_after_the_last_content_slice():
payload = _async_payload(content="I'll look that up for you")
payload["choices"][0]["message"]["tool_calls"] = [
{
@ -294,13 +294,13 @@ async def test_async_tool_calls_and_function_call_only_on_first_chunk():
}
]
chunks = await _collect_async(payload)
assert len(chunks) > 1
first_tool_calls = chunks[0].choices[0].delta.tool_calls
assert first_tool_calls is not None and len(first_tool_calls) == 1
assert first_tool_calls[0].function.name == "lookup"
for c in chunks[1:]:
assert c.choices[0].delta.tool_calls is None
assert c.choices[0].delta.function_call is None
called: Final = [i for i, c in enumerate(chunks) if c.choices[0].delta.tool_calls]
last_content: Final = max(i for i, c in enumerate(chunks) if c.choices[0].delta.content)
assert called == [last_content + 1] and last_content > 0
assert chunks[called[0]].choices[0].delta.tool_calls[0].function.name == "lookup"
assert chunks[called[0]].choices[0].delta.content is None
assert all(c.choices[0].delta.function_call is None for c in chunks)
assert chunks[-1].choices[0].finish_reason == "stop" and len(chunks) == called[0] + 2
@pytest.mark.asyncio
@ -385,7 +385,7 @@ def test_sync_multi_chunk_and_lossless_join():
assert chunks[-1].usage.total_tokens == 12
def test_sync_delta_and_choice_metadata_only_on_first_chunk():
def test_sync_reasoning_leads_and_choice_metadata_rides_the_first_content_slice():
thinking_blocks = [{"type": "thinking", "thinking": "cached thinking", "signature": "sig"}]
payload = {
"id": "chatcmpl-test-sync",
@ -408,25 +408,22 @@ def test_sync_delta_and_choice_metadata_only_on_first_chunk():
],
}
chunks = list(convert_to_streaming_response(payload))
assert len(chunks) > 1
first_choice = chunks[0].choices[0]
assert getattr(first_choice.delta, "reasoning_content", None) == "cached reasoning"
assert getattr(first_choice.delta, "thinking_blocks", None) == thinking_blocks
assert first_choice.logprobs is not None
assert first_choice.enhancements == {"source": "cache"}
for c in chunks[1:]:
lead: Final = chunks[0].choices[0]
assert getattr(lead.delta, "reasoning_content", None) == "cached reasoning"
assert getattr(lead.delta, "thinking_blocks", None) == thinking_blocks
assert lead.delta.content is None
first_text: Final = chunks[1].choices[0]
assert first_text.delta.content and first_text.logprobs is not None
assert first_text.enhancements == {"source": "cache"}
for index, c in enumerate(chunks):
choice = c.choices[0]
assert getattr(choice.delta, "reasoning_content", None) is None
assert getattr(choice.delta, "thinking_blocks", None) is None
assert choice.logprobs is None
assert getattr(choice, "enhancements", None) is None
assert index == 0 or getattr(choice.delta, "reasoning_content", None) is None
assert index == 0 or getattr(choice.delta, "thinking_blocks", None) is None
assert index == 1 or choice.logprobs is None
assert index == 1 or getattr(choice, "enhancements", None) is None
def test_sync_non_enumerated_delta_fields_only_on_first_chunk():
# annotations (and any other Delta field beyond the role/tool_call/reasoning
# set) must also stay on the first slice. Rebuilding later slices as a bare
# content delta drops the whole class, so this holds without enumerating
# every field by hand.
def test_sync_annotations_arrive_once_with_the_last_content_slice():
annotations = [
{
"type": "url_citation",
@ -456,7 +453,161 @@ def test_sync_non_enumerated_delta_fields_only_on_first_chunk():
],
}
chunks = list(convert_to_streaming_response(payload))
assert len(chunks) > 1
assert getattr(chunks[0].choices[0].delta, "annotations", None) == annotations
for c in chunks[1:]:
assert getattr(c.choices[0].delta, "annotations", None) is None
annotated: Final = [i for i, c in enumerate(chunks) if getattr(c.choices[0].delta, "annotations", None)]
last_content: Final = max(i for i, c in enumerate(chunks) if c.choices[0].delta.content)
assert annotated == [last_content] and last_content > 0
assert chunks[last_content].choices[0].delta.annotations == annotations
async def _client_stream(
message: dict[str, JsonValue], replay: Literal["async", "sync"]
) -> list[dict[str, JsonValue]]:
payload: Final = {
**_async_payload(),
"choices": [{"index": 0, "message": {"role": "assistant", **message}, "finish_reason": "stop"}],
}
logging_obj: Final = litellm.Logging(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=True,
call_type="acompletion",
start_time=datetime.now(),
litellm_call_id="cached-replay",
function_id="cached-replay",
)
source: Final = (
convert_to_streaming_response_async(response_object=payload)
if replay == "async"
else convert_to_streaming_response(response_object=payload)
)
wrapper: Final = CustomStreamWrapper(
completion_stream=source, model="gpt-4o-mini", custom_llm_provider="cached_response", logging_obj=logging_obj
)
chunks: Final = [chunk async for chunk in wrapper] if replay == "async" else list(wrapper)
return [
{**chunk.choices[0].delta.model_dump(exclude_none=True), "finish": chunk.choices[0].finish_reason}
for chunk in chunks
]
@pytest.mark.asyncio
@pytest.mark.parametrize("replay", ["async", "sync"])
@pytest.mark.parametrize(
"field, value",
[
("audio", {"id": "audio_1", "data": "UklGRg==", "expires_at": 1700000600, "transcript": "Hello"}),
("images", [{"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}, "index": 0}]),
],
)
async def test_a_cached_answer_without_text_streams_its_role_and_payload_to_the_client(
replay: Literal["async", "sync"], field: str, value: JsonValue
) -> None:
deltas: Final = await _client_stream({"content": None, field: value}, replay)
assert deltas[0]["role"] == "assistant"
assert [delta[field] for delta in deltas if field in delta] == [value]
assert deltas[-1]["finish"] == "stop"
@pytest.mark.asyncio
@pytest.mark.parametrize("replay", ["async", "sync"])
async def test_a_cached_answer_streams_reasoning_before_the_text_and_citations_after_it(
replay: Literal["async", "sync"],
) -> None:
annotations: Final[JsonValue] = [
{
"type": "url_citation",
"url_citation": {"url": "https://example.com", "title": "Example", "start_index": 0, "end_index": 23},
}
]
deltas: Final = await _client_stream(
{
"content": "The answer is 14 apples",
"reasoning_content": "Count the apples in each basket",
"annotations": annotations,
},
replay,
)
reasoning: Final = [i for i, delta in enumerate(deltas) if "reasoning_content" in delta]
text: Final = [i for i, delta in enumerate(deltas) if delta.get("content")]
cited: Final = [i for i, delta in enumerate(deltas) if "annotations" in delta]
assert "".join(str(deltas[i]["reasoning_content"]) for i in reasoning) == "Count the apples in each basket"
assert "".join(str(deltas[i]["content"]) for i in text) == "The answer is 14 apples"
assert max(reasoning) < min(text)
assert cited == [max(text)] and deltas[max(text)]["annotations"] == annotations
assert deltas[-1]["finish"] == "stop"
@pytest.mark.asyncio
@pytest.mark.parametrize("replay", ["async", "sync"])
async def test_a_cached_answer_streams_its_tool_call_after_the_reasoning_and_the_text(
replay: Literal["async", "sync"],
) -> None:
tool_call: Final[JsonValue] = {
"id": "call_1",
"type": "function",
"function": {"name": "count_apples", "arguments": '{"basket": "all"}'},
"index": 0,
}
deltas: Final = await _client_stream(
{
"content": "Let me count the apples",
"reasoning_content": "I need the basket totals",
"tool_calls": [tool_call],
},
replay,
)
reasoning: Final = [i for i, delta in enumerate(deltas) if "reasoning_content" in delta]
text: Final = [i for i, delta in enumerate(deltas) if delta.get("content")]
called: Final = [i for i, delta in enumerate(deltas) if delta.get("tool_calls")]
assert "".join(str(deltas[i]["content"]) for i in text) == "Let me count the apples"
assert reasoning and max(reasoning) < min(text)
assert called == [max(text) + 1]
assert deltas[called[0]]["tool_calls"] == [tool_call]
assert deltas[-1]["finish"] == "stop"
@pytest.mark.asyncio
@pytest.mark.parametrize("replay", ["async", "sync"])
async def test_a_cached_tool_turn_without_text_streams_the_reasoning_before_the_tool_call(
replay: Literal["async", "sync"],
) -> None:
tool_call: Final[JsonValue] = {
"id": "call_1",
"type": "function",
"function": {"name": "count_apples", "arguments": '{"basket": "all"}'},
"index": 0,
}
deltas: Final = await _client_stream(
{"content": None, "reasoning_content": "I need the basket totals", "tool_calls": [tool_call]}, replay
)
reasoning: Final = [i for i, delta in enumerate(deltas) if "reasoning_content" in delta]
called: Final = [i for i, delta in enumerate(deltas) if delta.get("tool_calls")]
assert deltas[0]["role"] == "assistant"
assert "".join(str(deltas[i]["reasoning_content"]) for i in reasoning) == "I need the basket totals"
assert reasoning and called and max(reasoning) < min(called)
assert [deltas[i]["tool_calls"] for i in called] == [[tool_call]]
assert deltas[-1]["finish"] == "stop"
@pytest.mark.asyncio
@pytest.mark.parametrize("replay", ["async", "sync"])
@pytest.mark.parametrize("whitespace", ["\n\n", " "], ids=["newlines", "space"])
async def test_a_cached_tool_turn_keeps_its_whitespace_only_text_between_the_reasoning_and_the_tool_call(
replay: Literal["async", "sync"], whitespace: str
) -> None:
tool_call: Final[JsonValue] = {
"id": "call_1",
"type": "function",
"function": {"name": "count_apples", "arguments": '{"basket": "all"}'},
"index": 0,
}
deltas: Final = await _client_stream(
{"content": whitespace, "reasoning_content": "I need the basket totals", "tool_calls": [tool_call]}, replay
)
reasoning: Final = [i for i, delta in enumerate(deltas) if "reasoning_content" in delta]
text: Final = [i for i, delta in enumerate(deltas) if delta.get("content")]
called: Final = [i for i, delta in enumerate(deltas) if delta.get("tool_calls")]
assert "".join(str(deltas[i]["content"]) for i in text) == whitespace
assert reasoning and text and called and max(reasoning) < min(text) and max(text) < min(called)
assert [deltas[i]["tool_calls"] for i in called] == [[tool_call]]
assert deltas[-1]["finish"] == "stop"