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

Merge 236df8e316 into 98d46ee59d

Browse source
This commit is contained in:
Jin 2026-09-12 13:02:01 +08:00 • committed by GitHub
commit 902419c45a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 253 additions and 60 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

104
litellm/litellm_core_utils/streaming_chunk_builder_utils.py
View file

@ -1,9 +1,9 @@
import base64
import time
from collections.abc import Callable, Iterator, Mapping, Sequence
from itertools import groupby
from itertools import accumulate, chain, groupby, tee
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, TypeAlias, TypedDict, Union, cast
from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeAlias, TypedDict, Union, cast
from typing_extensions import ReadOnly, Required
@ -58,6 +58,7 @@ class _ThinkingBlockFragment(TypedDict, total=False):
class _ThinkingDelta(TypedDict, total=False):
thinking_blocks: Sequence[_ThinkingBlockFragment]
provider_specific_fields: ReadOnly[Mapping[str, object] | None]
class _ThinkingChoice(TypedDict, total=False):
@ -68,6 +69,11 @@ class _ThinkingChunk(TypedDict):
choices: Sequence[_ThinkingChoice]
class _ThinkingStreamFragment(NamedTuple):
block: _ThinkingBlockFragment
is_snapshot: bool
class _ContentChoice(TypedDict, total=False):
delta: Mapping[str, str | None]
@ -661,60 +667,56 @@ class ChunkProcessor:
def get_combined_thinking_content(
self, chunks: Sequence["_ThinkingChunk"]
) -> list[Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"]] | None:
fragments, boundary_fragments = tee(self._iter_thinking_fragments(chunks))
# Count completed blocks before each fragment, keeping signatures with their preceding text.
closed_blocks: Final = accumulate(
(
int(fragment.block.get("type") == "redacted_thinking" or bool(fragment.block.get("signature")))
for fragment in boundary_fragments
),
initial=0,
)
grouped: Final = groupby(zip(closed_blocks, fragments, strict=False), key=lambda entry: entry[0])
groups: Final = (tuple(fragment for _, fragment in group) for _, group in grouped)
blocks: Final = tuple(block for group in groups if (block := self._assemble_thinking_block(group)) is not None)
return list(blocks) if blocks else None # mutable-ok: Message.thinking_blocks requires a list
@staticmethod
def _iter_thinking_fragments(chunks: Sequence["_ThinkingChunk"]) -> Iterator[_ThinkingStreamFragment]:
for choice in chain.from_iterable(chunk["choices"] for chunk in chunks):
if (delta := choice.get("delta")) is None or not isinstance(blocks := delta.get("thinking_blocks"), list):
continue
for block in blocks:
yield _ThinkingStreamFragment(
block,
isinstance(provider_fields := delta.get("provider_specific_fields"), Mapping)
and provider_fields.get("thinking_blocks") == blocks,
)
@staticmethod
def _assemble_thinking_block(
fragments: Sequence[_ThinkingStreamFragment],
) -> Union["ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock", None]:
from litellm.types.llms.openai import (
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
)
thinking_blocks: Final[list[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock]] = []
current_thinking_text_parts: list[str] = []
current_signature: str | None = 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:
delta = choice.get("delta", {})
thinking = delta.get("thinking_blocks", None)
if thinking and isinstance(thinking, list):
for thinking_block in thinking:
thinking_type = thinking_block.get("type", None)
if thinking_type and thinking_type == "redacted_thinking":
_flush_thinking_block()
redacted_data = thinking_block.get("data", None)
if redacted_data:
thinking_blocks.append(
ChatCompletionRedactedThinkingBlock(
type="redacted_thinking",
data=redacted_data,
)
)
else:
thinking_text = thinking_block.get("thinking", None)
if thinking_text:
current_thinking_text_parts.append(thinking_text)
signature = thinking_block.get("signature", None)
if signature:
current_signature = signature
_flush_thinking_block()
_flush_thinking_block()
if len(thinking_blocks) > 0:
return thinking_blocks
return None
last: Final = fragments[-1]
if last.block.get("type") == "redacted_thinking":
return (
ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data=data)
if (data := last.block.get("data"))
else None
)
if not (signature := last.block.get("signature")):
return None
text: Final = (
last.block.get("thinking") or ""
if last.is_snapshot
else "".join(fragment.block.get("thinking") or "" for fragment in fragments)
)
return ChatCompletionThinkingBlock(type="thinking", thinking=text, signature=signature)
def get_combined_reasoning_content(self, chunks: Sequence["_ContentChunk"]) -> ChatCompletionAssistantContentValue:
return self.get_combined_content(chunks, delta_key="reasoning_content")

11
litellm/responses/litellm_completion_transformation/streaming_iterator.py
View file

@ -547,6 +547,16 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
)
return response
def _encoded_thinking_blocks(self) -> str | None:
response: Final = (
self.litellm_model_response
if isinstance(self.litellm_model_response, ModelResponse)
else self.create_litellm_model_response()
)
if response is None:
return None
return LiteLLMCompletionResponsesConfig.encode_thinking_blocks(response.choices[0].message)
@staticmethod
def _snapshot_chunk_for_stream_chunk_builder(
chunk: ModelResponseStream,
@ -746,6 +756,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
**{
"id": reasoning_item_id,
"type": "reasoning",
"encrypted_content": self._encoded_thinking_blocks(),
"summary": [
{
"type": "summary_text",

6
litellm/responses/litellm_completion_transformation/transformation.py
View file

@ -1384,7 +1384,7 @@ class LiteLLMCompletionResponsesConfig:
input_item: Mapping[str, object],
) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...] | None:
"""
Decode ``encrypted_content`` written by ``_encode_thinking_blocks`` back
Decode ``encrypted_content`` written by ``encode_thinking_blocks`` back
into the signed thinking blocks it serialized.
LiteLLM writes this field itself for providers whose reasoning is signed
@ -2403,7 +2403,7 @@ class LiteLLMCompletionResponsesConfig:
return output_items
@staticmethod
def _encode_thinking_blocks(message: Message) -> str | None:
def encode_thinking_blocks(message: Message) -> str | None:
thinking_blocks: Final[Sequence[Mapping[str, object]]] = getattr(message, "thinking_blocks", None) or ()
preserved: Final = tuple(block for block in thinking_blocks if block.get("signature") or block.get("data"))
return json.dumps(preserved, separators=(",", ":")) if preserved else None
@ -2417,7 +2417,7 @@ class LiteLLMCompletionResponsesConfig:
if hasattr(choice, "message") and choice.message:
message = choice.message
reasoning_content: str = getattr(message, "reasoning_content", None) or ""
encrypted_content = LiteLLMCompletionResponsesConfig._encode_thinking_blocks(message)
encrypted_content = LiteLLMCompletionResponsesConfig.encode_thinking_blocks(message)
if reasoning_content or encrypted_content:
# Only check the first choice for reasoning content
return [

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

@ -8,6 +8,7 @@ import pytest
from litellm import ChatCompletionUsageBlock, stream_chunk_builder
from litellm.types.utils import GenericStreamingChunk
from litellm.litellm_core_utils.streaming_chunk_builder_utils import ChunkProcessor
from litellm.types.llms.openai import ChatCompletionThinkingBlock
from litellm.types.utils import (
ChatCompletionDeltaToolCall,
ChatCompletionMessageToolCall,
@ -217,11 +218,8 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks():
),
]
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)
result = processor.get_combined_thinking_content(chunks)
assert result is not None
assert len(result) == 3
@ -235,6 +233,54 @@ def test_get_combined_thinking_content_preserves_interleaved_blocks():
assert result[2]["signature"] == "sig_block2"
@pytest.mark.parametrize("snapshot", [True, False], ids=["provider-snapshot", "genuine-final-delta"])
def test_stream_chunk_builder_distinguishes_thinking_snapshots_from_repeated_deltas(snapshot: bool) -> None:
signed: Final = ChatCompletionThinkingBlock(type="thinking", thinking="echo", signature="test-signature")
deltas: Final = (
Delta(thinking_blocks=[ChatCompletionThinkingBlock(type="thinking", thinking="echo")]),
Delta(thinking_blocks=[signed], provider_specific_fields={"thinking_blocks": [signed]} if snapshot else None),
)
chunks: Final = [
ModelResponseStream(
id="chatcmpl-thinking",
model="claude-opus-5",
choices=[StreamingChoices(index=0, delta=delta, finish_reason="stop" if index == 1 else None)],
)
for index, delta in enumerate(deltas)
]
response: Final = stream_chunk_builder(chunks=chunks)
assert response is not None
assert response.choices[0].message.thinking_blocks == [
{"type": "thinking", "thinking": "echo" if snapshot else "echoecho", "signature": "test-signature"}
]
def test_incomplete_thinking_stream_preserves_summary_without_signed_blocks() -> None:
chunk: Final = ModelResponseStream(
id="chatcmpl-incomplete-thinking",
model="claude-opus-5",
choices=[
StreamingChoices(
index=0,
finish_reason="length",
delta=Delta(
reasoning_content="Unfinished reasoning",
thinking_blocks=[ChatCompletionThinkingBlock(type="thinking", thinking="Unfinished reasoning")],
),
)
],
)
response: Final = stream_chunk_builder(chunks=[chunk])
assert response is not None
assert response.choices[0].finish_reason == "length"
assert response.choices[0].message.reasoning_content == "Unfinished reasoning"
assert response.choices[0].message.thinking_blocks is None
def test_cache_read_input_tokens_retained():
chunk1 = ModelResponseStream(
id="chatcmpl-95aabb85-c39f-443d-ae96-0370c404d70c",

2
tests/test_litellm/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py
View file

@ -167,7 +167,7 @@ class TestEncryptedReasoningRoundTrip:
{"type": "redacted_thinking", "data": "redacted-payload"},
]
message = Message(role="assistant", content="answer", thinking_blocks=blocks)
encoded = LiteLLMCompletionResponsesConfig._encode_thinking_blocks(message)
encoded = LiteLLMCompletionResponsesConfig.encode_thinking_blocks(message)
decoded = LiteLLMCompletionResponsesConfig._decode_thinking_blocks_from_input_item(
{"type": "reasoning", "encrypted_content": encoded}
)

136
tests/test_litellm/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py
View file

@ -19,8 +19,16 @@ import pytest
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import ResponsesAPIStreamEvents
from litellm.types.llms.openai import (
ChatCompletionRedactedThinkingBlock,
ChatCompletionThinkingBlock,
ResponseCompletedEvent,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import (
Delta,
ModelResponse,
@ -705,3 +713,129 @@ def test_streamed_unrecognized_tool_choice_is_echoed_as_auto() -> None:
]
assert [event.response.tool_choice for event in response_events] == ["auto", "auto", "auto"]
def _signed_thinking_chunks(cumulative: bool) -> tuple[ModelResponseStream, ...]:
blocks: Final = [
ChatCompletionThinkingBlock(type="thinking", thinking="One plus one equals two.", signature="test-signature"),
ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data="test-redacted-data"),
]
deltas: Final = (
Delta(
reasoning_content="One plus ",
thinking_blocks=[ChatCompletionThinkingBlock(type="thinking", thinking="One plus ")],
),
Delta(
reasoning_content="one equals two.",
thinking_blocks=[ChatCompletionThinkingBlock(type="thinking", thinking="one equals two.")],
),
Delta(
thinking_blocks=blocks
if cumulative
else [
ChatCompletionThinkingBlock(type="thinking", thinking="", signature="test-signature"),
ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data="test-redacted-data"),
],
provider_specific_fields={"thinking_blocks": blocks} if cumulative else None,
),
Delta(content="2"),
)
return tuple(
ModelResponseStream(
id=CHAT_COMPLETION_ID,
model="test-model",
choices=[StreamingChoices(index=0, delta=delta, finish_reason="stop" if index == 3 else None)],
)
for index, delta in enumerate(deltas)
)
@pytest.mark.parametrize("cumulative", [True, False], ids=["cumulative-provider-blocks", "delta-blocks"])
@pytest.mark.parametrize("asynchronous", [True, False], ids=["async", "sync"])
async def test_completed_response_replays_signed_thinking_unchanged(cumulative: bool, asynchronous: bool) -> None:
iterator: Final = _build_iterator(_signed_thinking_chunks(cumulative))
events: Final = [event async for event in iterator] if asynchronous else list(iterator)
completed: Final = next(event for event in events if isinstance(event, ResponseCompletedEvent))
reasoning: Final = next(item for item in completed.response.output if item.type == "reasoning")
assert reasoning.encrypted_content is not None
assert json.loads(reasoning.encrypted_content) == [
{"type": "thinking", "thinking": "One plus one equals two.", "signature": "test-signature"},
{"type": "redacted_thinking", "data": "test-redacted-data"},
]
messages: Final = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message(
input_item=reasoning.model_dump(exclude_none=True), replay_reasoning=True
)
assert len(messages) == 1
assert messages[0]["thinking_blocks"] == [
{"type": "thinking", "thinking": "One plus one equals two.", "signature": "test-signature"},
{"type": "redacted_thinking", "data": "test-redacted-data"},
]
@pytest.mark.parametrize("cumulative", [True, False], ids=["cumulative-provider-blocks", "delta-blocks"])
async def test_reasoning_done_preserves_the_replay_payload(cumulative: bool) -> None:
iterator: Final = _build_iterator(_signed_thinking_chunks(cumulative))
events: Final = [event async for event in iterator]
done: Final = next(
event for event in events if event.type == "response.output_item.done" and event.item.type == "reasoning"
)
completed: Final = next(event for event in events if isinstance(event, ResponseCompletedEvent))
reasoning: Final = next(item for item in completed.response.output if item.type == "reasoning")
payload: Final = done.item.model_dump().get("encrypted_content")
assert payload is not None
assert json.loads(payload) == [
{"type": "thinking", "thinking": "One plus one equals two.", "signature": "test-signature"},
{"type": "redacted_thinking", "data": "test-redacted-data"},
]
assert payload == reasoning.encrypted_content
@pytest.mark.parametrize("asynchronous", [True, False], ids=["async", "sync"])
async def test_streamed_signature_only_thinking_is_replayable(asynchronous: bool) -> None:
block: Final = ChatCompletionThinkingBlock(type="thinking", thinking="", signature="opaque-signature")
chunks: Final = [
ModelResponseStream(
id=CHAT_COMPLETION_ID,
model="test-model",
choices=[
StreamingChoices(
index=0,
delta=Delta(thinking_blocks=[block], provider_specific_fields={"thinking_blocks": [block]}),
)
],
),
_tool_call_chunk(finish_reason="tool_calls"),
]
iterator: Final = _build_iterator(chunks)
events: Final = [event async for event in iterator] if asynchronous else list(iterator)
completed: Final = next(event for event in events if isinstance(event, ResponseCompletedEvent))
reasoning: Final = next(item for item in completed.response.output if item.type == "reasoning")
assert json.loads(reasoning.encrypted_content) == [block]
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=[item.model_dump(exclude_none=True) for item in completed.response.output],
responses_api_request={},
replay_reasoning=True,
)
tool_message: Final = next(message for message in messages if message.get("tool_calls"))
assert tool_message["thinking_blocks"] == [block]
def test_reasoning_done_without_a_response_snapshot_preserves_summary() -> None:
iterator: Final = _build_iterator([])
event: Final = iterator.create_reasoning_output_item_done_event(
reasoning_item_id="rs_pending",
reasoning_content="The response snapshot is not available yet.",
sequence_number=7,
)
assert event.type == "response.output_item.done"
assert event.sequence_number == 7
assert event.item.model_dump(exclude_none=True) == {
"id": "rs_pending",
"type": "reasoning",
"summary": [{"type": "summary_text", "text": "The response snapshot is not available yet."}],
}