scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(streaming): preserve interleaved thinking/redacted blocks

Browse source
This commit is contained in:
Emerson Gomes 2026-02-08 05:04:47 -06:00
parent 4f96a3b126
commit f7d03f8a43
2 changed files with 104 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

58
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -1,6 +1,6 @@
import base64
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
from litellm.types.llms.openai import (
ChatCompletionAssistantContentValue,
@ -326,10 +326,22 @@ class ChunkProcessor:
thinking_blocks: List[
Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]
] = []
combined_thinking_text: Optional[str] = None
data: Optional[str] = None
signature: Optional[str] = None
type: Literal["thinking", "redacted_thinking"] = "thinking"
current_thinking_text_parts: List[str] = []
current_signature: Optional[str] = None
def _flush_thinking_block() -> None:
nonlocal current_thinking_text_parts, current_signature
if len(current_thinking_text_parts) > 0 and current_signature:
thinking_blocks.append(
ChatCompletionThinkingBlock(
type="thinking",
thinking="".join(current_thinking_text_parts),
signature=current_signature,
)
)
current_thinking_text_parts = []
current_signature = None
for chunk in chunks:
choices = chunk["choices"]
for choice in choices:
@ -339,33 +351,25 @@ class ChunkProcessor:
for thinking_block in thinking:
thinking_type = thinking_block.get("type", None)
if thinking_type and thinking_type == "redacted_thinking":
type = "redacted_thinking"
data = thinking_block.get("data", None)
_flush_thinking_block()
redacted_data = thinking_block.get("data", None)
if redacted_data:
thinking_blocks.append(
ChatCompletionRedactedThinkingBlock(
type="redacted_thinking",
data=redacted_data,
)
)
else:
type = "thinking"
thinking_text = thinking_block.get("thinking", None)
if thinking_text:
if combined_thinking_text is None:
combined_thinking_text = ""
combined_thinking_text += thinking_text
current_thinking_text_parts.append(thinking_text)
signature = thinking_block.get("signature", None)
if signature:
current_signature = signature
_flush_thinking_block()
if combined_thinking_text and type == "thinking" and signature:
thinking_blocks.append(
ChatCompletionThinkingBlock(
type=type,
thinking=combined_thinking_text,
signature=signature,
)
)
elif data and type == "redacted_thinking":
thinking_blocks.append(
ChatCompletionRedactedThinkingBlock(
type=type,
data=data,
)
)
_flush_thinking_block()
if len(thinking_blocks) > 0:
return thinking_blocks

74
tests/test_litellm/litellm_core_utils/test_streaming_chunk_builder_utils.py
View file

@ -158,6 +158,78 @@ def test_get_combined_tool_content():
]
def test_get_combined_thinking_content_preserves_interleaved_blocks():
base_chunk = {
"id": "chatcmpl-123",
"object": "chat.completion.chunk",
"created": 1234567890,
"model": "claude-sonnet-4-20250514",
}
def make_chunk(**delta_kwargs):
return ModelResponseStream(
**{
**base_chunk,
"choices": [
{
"index": 0,
"delta": delta_kwargs,
"finish_reason": None,
}
],
}
)
chunks = [
make_chunk(role="assistant", content=None),
make_chunk(
thinking_blocks=[
{"type": "thinking", "thinking": "Step 1 analysis...", "signature": None}
]
),
make_chunk(
thinking_blocks=[
{"type": "thinking", "thinking": None, "signature": "sig_block1"}
]
),
make_chunk(
thinking_blocks=[
{
"type": "redacted_thinking",
"data": "EuoBCoYBGAIi...encrypted...",
}
]
),
make_chunk(
thinking_blocks=[
{"type": "thinking", "thinking": "Step 2 analysis...", "signature": None}
]
),
make_chunk(
thinking_blocks=[
{"type": "thinking", "thinking": None, "signature": "sig_block2"}
]
),
]
thinking_chunks = [
chunk for chunk in chunks if chunk["choices"][0]["delta"].get("thinking_blocks")
]
processor = ChunkProcessor(chunks=chunks)
result = processor.get_combined_thinking_content(thinking_chunks)
assert result is not None
assert len(result) == 3
assert result[0]["type"] == "thinking"
assert result[0]["thinking"] == "Step 1 analysis..."
assert result[0]["signature"] == "sig_block1"
assert result[1]["type"] == "redacted_thinking"
assert result[1]["data"] == "EuoBCoYBGAIi...encrypted..."
assert result[2]["type"] == "thinking"
assert result[2]["thinking"] == "Step 2 analysis..."
assert result[2]["signature"] == "sig_block2"
def test_cache_read_input_tokens_retained():
chunk1 = ModelResponseStream(
id="chatcmpl-95aabb85-c39f-443d-ae96-0370c404d70c",
@ -441,4 +513,4 @@ def test_stream_chunk_builder_anthropic_web_search():
assert usage.prompt_tokens == 50
assert usage.completion_tokens == 27
assert usage.total_tokens == 77
assert usage.server_tool_use['web_search_requests'] == 2
assert usage.server_tool_use['web_search_requests'] == 2