fix(streaming): keep reasoning_items when reassembling a streamed response

stream_chunk_builder dropped delta.reasoning_items on both of its paths. The
simple-text fast path returned before any aggregation ran, and the full path
had no reasoning_items branch, so the encrypted reasoning state a provider
sends back never reached the assembled assistant message.

That message is what gets cached by _assemble_complete_response_from_streaming_chunks
and what the Responses API bridge reads through _get_reasoning_items to rebuild
its input, so a streamed reasoning turn lost its reasoning state and could not
be round-tripped.

Add reasoning_items to the fast-path bail-out condition so such a stream is no
longer classified as simple text, and merge the items from every chunk in the
full path, matching how annotations and images are already handled.
This commit is contained in:
Vineeth Sai 2026-08-20 10:49:11 -07:00
parent d2d158f271
commit 49e549d9e9
2 changed files with 111 additions and 0 deletions

View file

@ -8589,6 +8589,7 @@ def stream_chunk_builder(
or delta.get("audio") is not None
or delta.get("images") is not None
or delta.get("provider_specific_fields") is not None
or delta.get("reasoning_items") is not None
):
is_simple_text_stream = False
break
@ -8741,6 +8742,22 @@ def stream_chunk_builder(
all_images.extend(chunk["choices"][0]["delta"]["images"])
response["choices"][0]["message"]["images"] = all_images
# Reasoning items carry the provider's encrypted reasoning state, which the
# Responses API bridge reads back off the assembled assistant message.
reasoning_item_chunks: Final = [
chunk
for chunk in chunks
if len(chunk["choices"]) > 0
and "reasoning_items" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["reasoning_items"] is not None
]
if len(reasoning_item_chunks) > 0:
all_reasoning_items: Final[list] = []
for chunk in reasoning_item_chunks:
all_reasoning_items.extend(chunk["choices"][0]["delta"]["reasoning_items"])
response["choices"][0]["message"]["reasoning_items"] = all_reasoning_items
# Combine provider_specific_fields from streaming chunks (e.g., web_search_results, citations)
# See: https://github.com/BerriAI/litellm/issues/17737
provider_specific_chunks: Final = [

View file

@ -0,0 +1,94 @@
"""
Tests for stream_chunk_builder reasoning item reassembly.
Previously, stream_chunk_builder dropped delta.reasoning_items entirely: the
simple-text fast path returned before any aggregation ran, and the full path
had no reasoning_items branch. The encrypted reasoning state a provider sends
back was therefore lost from the assembled message, which is what gets cached
and what the Responses API bridge reads to rebuild its input.
"""
from litellm import stream_chunk_builder
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
MESSAGES = [{"role": "user", "content": "hi"}]
def _chunk(**delta_kwargs) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-test",
created=1700000000,
model="test-model",
object="chat.completion.chunk",
choices=[
StreamingChoices(
finish_reason=None,
index=0,
delta=Delta(**delta_kwargs),
)
],
)
def test_reasoning_items_survive_a_text_only_stream():
"""
A stream carrying only content and reasoning_items must keep the reasoning
items. This is the fast path, which used to return before aggregating.
"""
reasoning_item = {
"id": "rs_abc123",
"type": "reasoning",
"encrypted_content": "ENCRYPTED-REASONING-BLOB",
"summary": [{"type": "summary_text", "text": "Thinking"}],
}
chunks = [
_chunk(content="Hello", role="assistant"),
_chunk(content="", reasoning_items=[reasoning_item]),
]
response = stream_chunk_builder(chunks, messages=MESSAGES)
assert response is not None
message = response.choices[0].message
assert message.content == "Hello"
assert getattr(message, "reasoning_items", None) == [reasoning_item]
def test_reasoning_items_are_merged_across_chunks():
"""
Reasoning items arriving in more than one chunk are concatenated, matching
how annotations and images are merged. thinking_blocks here also forces the
full aggregation path rather than the text-only fast path.
"""
item_a = {"id": "rs_a", "type": "reasoning", "encrypted_content": "BLOB-A"}
item_b = {"id": "rs_b", "type": "reasoning", "encrypted_content": "BLOB-B"}
chunks = [
_chunk(
content="Part one. ",
role="assistant",
thinking_blocks=[{"type": "thinking", "thinking": "step 1", "signature": "sig"}],
reasoning_items=[item_a],
),
_chunk(content="Part two.", reasoning_items=[item_b]),
]
response = stream_chunk_builder(chunks, messages=MESSAGES)
assert response is not None
message = response.choices[0].message
assert message.content == "Part one. Part two."
assert getattr(message, "reasoning_items", None) == [item_a, item_b]
def test_no_reasoning_items_leaves_the_message_alone():
"""A stream without reasoning items must not grow a reasoning_items field."""
chunks = [_chunk(content="Hello", role="assistant"), _chunk(content=" there")]
response = stream_chunk_builder(chunks, messages=MESSAGES)
assert response is not None
message = response.choices[0].message
assert message.content == "Hello there"
assert getattr(message, "reasoning_items", None) is None