From 498219c41cf1d8b5263dfdfa264a97662875cfd0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 17:32:23 -0700 Subject: [PATCH] fix(responses): trim thinking blocks on a reasoning turn that starts with web searches When Claude searches before writing any text, the reasoning item is followed directly by web_search_call items. Their text was merged onto the reasoning message, so it no longer looked reasoning-only and kept every thinking block side by side. Every replayed assistant message now keeps only its newest block. --- .../transformation.py | 10 ++-- .../test_reasoning_input_item_preservation.py | 50 ++++++++++--------- 2 files changed, 33 insertions(+), 27 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 3d2f54f78c1..63c0b415421 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1633,7 +1633,8 @@ class LiteLLMCompletionResponsesConfig: immediately following assistant message. Signed ``thinking_blocks`` decoded from ``encrypted_content`` travel the - same way, but only the newest one is replayed. A reasoning item does not + same way, but each assistant message keeps only the newest one, including + a reasoning message that hosted search text was merged onto. A reasoning item does not record where each block sat between the turn's server tool steps and text, and Claude rejects a latest assistant message whose thinking blocks sit side by side when it wrote them apart. @@ -1692,8 +1693,9 @@ class LiteLLMCompletionResponsesConfig: cast(dict[str, object], msg)["reasoning_content"] = combined # cast-ok: mutable reasoning carrier else: setattr(msg, "reasoning_content", combined) # noqa: B010 # attribute name is fixed, not dynamic - if pending_blocks: - replayed: Final = list((pending_blocks + (_thinking_blocks(msg) or ()))[-1:]) + own_blocks: Final = _thinking_blocks(msg) or () + if pending_blocks or len(own_blocks) > 1: + replayed: Final = list((pending_blocks + own_blocks)[-1:]) if isinstance(msg, dict): cast(dict[str, object], msg)["thinking_blocks"] = replayed # cast-ok: mutable reasoning carrier else: @@ -1736,7 +1738,7 @@ class LiteLLMCompletionResponsesConfig: pending.append((_reasoning_text(msg), _thinking_blocks(msg))) continue - if pending and _role(msg) == "assistant": + if _role(msg) == "assistant": _apply_pending(msg, pending) pending = [] elif pending: diff --git a/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py b/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py index 837b79b7550..f29062f8bcd 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_reasoning_input_item_preservation.py @@ -14,6 +14,7 @@ JSON array of thinking blocks on the response side. """ import json +from itertools import chain import pytest @@ -253,43 +254,46 @@ class TestEncryptedReasoningRoundTrip: assert len(messages) == 1 assert [block["thinking"] for block in messages[0]["thinking_blocks"]] == ["newer"] - def test_web_search_turn_replays_only_its_last_thinking_block(self): + @pytest.mark.parametrize("answer_first", [True, False], ids=["answer_before_searches", "searches_before_answer"]) + def test_web_search_turn_replays_only_its_last_thinking_block(self, answer_first): """Claude rejects a follow-up whose replay puts blocks it wrote apart side by side.""" blocks = [ {"type": "thinking", "thinking": "", "signature": "sig-before-search-1"}, {"type": "thinking", "thinking": "", "signature": "sig-before-search-2"}, {"type": "thinking", "thinking": "", "signature": "sig-before-answer"}, ] + answer = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "Paul Haggis.", "annotations": []}], + } + searches = [ + { + "type": "web_search_call", + "id": "ws_1", + "status": "completed", + "action": {"type": "search", "query": "2006 Best Picture"}, + }, + { + "type": "web_search_call", + "id": "ws_2", + "status": "completed", + "action": {"type": "search", "query": "Crash 2004 film director"}, + }, + ] messages = _transform_input( [ {"role": "user", "content": "Who directed the 2006 Best Picture winner?"}, {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": json.dumps(blocks)}, - { - "type": "message", - "id": "msg_1", - "role": "assistant", - "status": "completed", - "content": [{"type": "output_text", "text": "Paul Haggis.", "annotations": []}], - }, - { - "type": "web_search_call", - "id": "ws_1", - "status": "completed", - "action": {"type": "search", "query": "2006 Best Picture"}, - }, - { - "type": "web_search_call", - "id": "ws_2", - "status": "completed", - "action": {"type": "search", "query": "Crash 2004 film director"}, - }, + *([answer, *searches] if answer_first else [*searches, answer]), {"role": "user", "content": "Summarize that in one sentence."}, ] ) assistants = [message for message in messages if message["role"] == "assistant"] - assert len(assistants) == 1 - assert assistants[0]["thinking_blocks"] == [blocks[-1]] - assert "Paul Haggis." in json.dumps(assistants[0]["content"]) + assert list(chain.from_iterable(message.get("thinking_blocks") or [] for message in assistants)) == [blocks[-1]] + assert "Paul Haggis." in json.dumps([message["content"] for message in assistants]) assert messages[-1]["role"] == "user" def test_reasoning_before_a_user_turn_replays_one_block_in_one_message(self):