diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index c215e3f8395..0da04b5c25d 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1,6 +1,6 @@ import time import uuid -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Any, Final, cast import litellm @@ -11,7 +11,10 @@ from litellm.responses.litellm_completion_transformation.custom_tools import ( is_custom_tool_call, serialize_tool_call_arguments, ) -from litellm.responses.litellm_completion_transformation.reasoning_items import mint_reasoning_item_id +from litellm.responses.litellm_completion_transformation.reasoning_items import ( + encode_thinking_blocks, + mint_reasoning_item_id, +) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -655,6 +658,15 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): ) return response + def _encoded_thinking_blocks(self) -> str | None: + response: Final = self.create_litellm_model_response() + if response is None: + return None + thinking_blocks: Final[Sequence[Mapping[str, object]]] = ( + getattr(response.choices[0].message, "thinking_blocks", None) or () + ) + return encode_thinking_blocks(thinking_blocks) + @staticmethod def _snapshot_chunk_for_stream_chunk_builder( chunk: ModelResponseStream, @@ -845,6 +857,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): **{ "id": reasoning_item_id, "type": "reasoning", + "encrypted_content": self._encoded_thinking_blocks(), "summary": [ { "type": "summary_text", diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 041bcf1b6d7..3e800a19b5e 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -19,9 +19,14 @@ 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 ( BaseLiteLLMOpenAIResponseObject, + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, ResponsesAPIStreamEvents, ) from litellm.types.responses.main import build_web_search_call @@ -1171,3 +1176,102 @@ async def test_plain_text_stream_announces_exactly_one_message_item(sync_mode: b ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE, ): assert event.item_id == message_item_adds[0].item.id + + +@pytest.mark.parametrize( + "thinking_text,include_redacted", + [("echoecho", False), ("", False), (None, True), ("echoecho", True)], + ids=["visible", "signature-only", "redacted-only", "mixed"], +) +@pytest.mark.parametrize("end_with_tool", [False, True], ids=["text", "tool"]) +@pytest.mark.asyncio +async def test_reasoning_done_payload_replays_signed_blocks( + thinking_text: str | None, include_redacted: bool, end_with_tool: bool +) -> None: + fragments: Final[tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]] = ( + *( + ChatCompletionThinkingBlock(type="thinking", thinking=part) + for part in (("echo", "echo") if thinking_text else ()) + ), + *( + (ChatCompletionThinkingBlock(type="thinking", thinking="", signature="test-signature"),) + if thinking_text is not None + else () + ), + *( + (ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data="test-redacted-data"),) + if include_redacted + else () + ), + ) + chunks: Final = tuple( + ModelResponseStream( + id=CHAT_COMPLETION_ID, + model="test-model", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + reasoning_content=fragment.get("thinking", ""), + thinking_blocks=[fragment], + provider_specific_fields={"thinking_blocks": [fragment]}, + ), + ) + ], + ) + for fragment in fragments + ) + ending: Final = ( + _tool_call_chunk(finish_reason="tool_calls") if end_with_tool else _chunk("answer", finish_reason="stop") + ) + iterator: Final = _build_iterator([*chunks, ending]) + + events: Final = [event async for event in iterator] + done: Final = next( + event.item + for event in events + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and event.item.type == "reasoning" + ) + completed: Final = next( + event.response for event in events if event.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + ) + reasoning: Final = next(item for item in completed.output if item.type == "reasoning") + payload: Final = done.model_dump().get("encrypted_content") + expected: Final = [ + *( + [ChatCompletionThinkingBlock(type="thinking", thinking=thinking_text, signature="test-signature")] + if thinking_text is not None + else [] + ), + *( + [ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data="test-redacted-data")] + if include_redacted + else [] + ), + ] + + assert isinstance(payload, str) + assert json.loads(payload) == expected + assert payload == reasoning.encrypted_content + messages: Final = LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( + input_item=done.model_dump(exclude_none=True), replay_reasoning=True + ) + assert messages[0]["thinking_blocks"] == expected + + +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."}], + }