From 78d1957a0d877ab976d44599cafd304632adc043 Mon Sep 17 00:00:00 2001 From: Gurjot Singh Date: Tue, 7 Jul 2026 23:50:52 -0700 Subject: [PATCH] fix(litellm): preserve signature-only blocks for Anthropic streaming thinking aggregation Key changes: * **Streaming Signature Preservation:** When Anthropic thinking mode is enabled and responses are streamed via LiteLLM/OpenAI format, Anthropic splits thinking blocks across deltas: text-only chunks first (`signature: None`), followed by a final signature-only chunk at `block_stop` with empty text (`thinking: ""`). This change ensures signature-only deltas are preserved rather than discarded. * **Signal Detection & Part Conversion:** Updated `_has_meaningful_signal` to recognize `thinking_blocks` so signature-only deltas are not dropped as empty noise. Updated `_convert_reasoning_value_to_parts` to preserve blocks when `signature` is present, storing the cryptographic signature on `Part.thought_signature`. * **Outbound Thought Aggregation:** Added `_aggregate_streaming_thought_parts` helper in `lite_llm.py` and integrated it into `_content_to_message_param` (Anthropic branch). It joins fragmented text chunks and attaches the cryptographic signature from whichever part carries it, emitting a clean `thinking_blocks` entry for multi-turn session history. * **Evidence & Verification:** Added unit test `test_aggregate_streaming_thought_parts` and updated `test_content_to_message_param_anthropic_aggregates_streaming_split_thinking` to verify that streaming deltas (split text chunks + signature chunk `ErEDClsIDBACGAIfull`) are re-joined into a valid Anthropic `thinking_block` with base64 signature (`RXJFRENsc0lEQkFDR0FJZnVsbA==`). Verified 321 unit tests and 957 internal TAP targets pass 100%. Note for callers: This fixes continuity of Anthropic extended thinking across streaming responses that include tool use or multi-turn chats. Previously, after a tool call or first streaming turn, thinking blocks were rejected by Anthropic API with a 400 Bad Request error (`missing thinking signature`) because the signature was discarded during streaming. Fixes: https://github.com/google/adk-python/issues/4801 Closes: https://github.com/google/adk-python/pull/5437 Co-authored-by: Yi Liu COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/5437 from gurjot-05:fix/anthropic-thinking-streaming-continuity 0ce04f16dac92d579a78c427e07765a6b65072ea PiperOrigin-RevId: 944305912 --- src/google/adk/models/lite_llm.py | 51 ++++++++- tests/unittests/models/test_litellm.py | 141 ++++++++++++++++++++++++- 2 files changed, 184 insertions(+), 8 deletions(-) diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 2e4cf65f32..ec1f5d663d 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -543,9 +543,11 @@ def _convert_reasoning_value_to_parts(reasoning_value: Any) -> List[types.Part]: continue if block_type == "thinking": thinking_text = block.get("thinking", "") - if thinking_text: + signature = block.get("signature") + # Anthropic streams a signature in a final chunk with empty text. + # Preserve signature-only blocks so the signature survives aggregation. + if thinking_text or signature: part = types.Part(text=thinking_text, thought=True) - signature = block.get("signature") if signature: decoded_signature = _decode_thought_signature(signature) part.thought_signature = decoded_signature or str( @@ -565,6 +567,43 @@ def _convert_reasoning_value_to_parts(reasoning_value: Any) -> List[types.Part]: ] +def _aggregate_streaming_thought_parts( + thought_parts: Iterable[types.Part], +) -> List[types.Part]: + """Aggregates fragmented streaming thought parts into clean individual parts. + + During streaming, Anthropic splits a thinking block across many deltas: + text-only chunks followed by a signature-only chunk at block_stop. This helper + joins the text chunks and attaches the signature, producing clean individual + thought parts for session history and outbound requests. + """ + parts_list = list(thought_parts) + if not parts_list: + return [] + aggregated: List[types.Part] = [] + current_texts: List[str] = [] + for part in parts_list: + if part.text: + current_texts.append(part.text) + if part.thought_signature: + aggregated.append( + types.Part( + text="".join(current_texts), + thought=True, + thought_signature=part.thought_signature, + ) + ) + current_texts = [] + if current_texts: + aggregated.append( + types.Part( + text="".join(current_texts), + thought=True, + ) + ) + return aggregated + + def _extract_reasoning_value(message: Message | Delta | None) -> Any: """Fetches the reasoning payload from a LiteLLM message. @@ -1023,9 +1062,14 @@ async def _content_to_message_param( # For Anthropic models, rebuild thinking_blocks with signatures so that # thinking is preserved across tool call boundaries. Without this, # Anthropic silently drops thinking after the first turn. + # + # Streaming splits one Anthropic thinking block across many deltas: + # text-only chunks followed by a signature-only chunk at block_stop. + # Aggregate them back into one thinking block for outbound. if model and _is_anthropic_model(model) and reasoning_parts: + aggregated_parts = _aggregate_streaming_thought_parts(reasoning_parts) thinking_blocks = [] - for part in reasoning_parts: + for part in aggregated_parts: if part.text and part.thought_signature: sig = part.thought_signature if isinstance(sig, bytes): @@ -1877,6 +1921,7 @@ def _has_meaningful_signal(message: Message | Delta | None) -> bool: or message.get("function_call") or message.get("reasoning_content") or message.get("reasoning") + or message.get("thinking_blocks") ) if isinstance(response, ModelResponseStream): diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index 6f4341f55f..b887f598f2 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -27,6 +27,7 @@ from unittest.mock import patch import warnings +from google.adk.models.lite_llm import _aggregate_streaming_thought_parts from google.adk.models.lite_llm import _append_fallback_user_content_if_missing from google.adk.models.lite_llm import _content_to_message_param from google.adk.models.lite_llm import _convert_reasoning_value_to_parts @@ -5396,15 +5397,48 @@ def test_convert_reasoning_value_to_parts_skips_redacted_blocks(): assert parts[0].text == "visible" -def test_convert_reasoning_value_to_parts_skips_empty_thinking(): - """Blocks with empty thinking text are excluded.""" +def test_convert_reasoning_value_to_parts_preserves_signature_only_blocks(): + """Signature-only blocks (empty text) are preserved for streaming aggregation. + + Anthropic emits the block_stop signature as a delta with empty thinking text. + Dropping it would lose the signature, breaking multi-turn thinking continuity. + Blocks with neither text nor signature are still skipped. + """ thinking_blocks = [ {"type": "thinking", "thinking": "", "signature": "c2lnMQ=="}, {"type": "thinking", "thinking": "real thought", "signature": "c2lnMg=="}, + { + "type": "thinking", + "thinking": "", + "signature": "", + }, # fully empty: drop ] parts = _convert_reasoning_value_to_parts(thinking_blocks) - assert len(parts) == 1 - assert parts[0].text == "real thought" + assert len(parts) == 2 + assert parts[0].text == "" + assert parts[0].thought is True + assert parts[0].thought_signature == b"sig1" + assert parts[1].text == "real thought" + assert parts[1].thought_signature == b"sig2" + + +def test_aggregate_streaming_thought_parts(): + """Tests aggregating fragmented streaming thought parts and multiple blocks.""" + parts = [ + types.Part(text="First block ", thought=True), + types.Part(text="text.", thought=True), + types.Part(text="", thought=True, thought_signature=b"sig1"), + types.Part(text="Second block", thought=True, thought_signature=b"sig2"), + types.Part(text="Trailing without sig", thought=True), + ] + aggregated = _aggregate_streaming_thought_parts(parts) + assert len(aggregated) == 3 + assert aggregated[0].text == "First block text." + assert aggregated[0].thought_signature == b"sig1" + assert aggregated[1].text == "Second block" + assert aggregated[1].thought_signature == b"sig2" + assert aggregated[2].text == "Trailing without sig" + assert aggregated[2].thought_signature is None def test_convert_reasoning_value_to_parts_flat_string_unchanged(): @@ -5478,6 +5512,33 @@ async def test_content_to_message_param_anthropic_model_round_trip_preserves_sig assert result.get("reasoning_content") is None +@pytest.mark.asyncio +async def test_content_to_message_param_anthropic_split_thinking_and_signature(): + """Combines separate thinking and signature parts into a single thinking_block.""" + content = types.Content( + role="model", + parts=[ + types.Part(text="deep thought", thought=True), + types.Part( + text="", thought=True, thought_signature=b"sig_round_trip" + ), + types.Part(text="Hello!"), + ], + ) + result = await _content_to_message_param( + content, model="anthropic/claude-4-sonnet" + ) + assert result["role"] == "assistant" + assert "thinking_blocks" in result + assert result.get("reasoning_content") is None + blocks = result["thinking_blocks"] + assert len(blocks) == 1 + assert blocks[0]["type"] == "thinking" + assert blocks[0]["thinking"] == "deep thought" + assert blocks[0]["signature"] == "c2lnX3JvdW5kX3RyaXA=" + assert result["content"] == "Hello!" + + @pytest.mark.asyncio async def test_content_to_message_param_non_anthropic_uses_reasoning_content(): """For non-Anthropic models, reasoning_content is used as before.""" @@ -5577,7 +5638,7 @@ def test_convert_reasoning_value_to_parts_empty_thinking_does_not_fall_through() "type": "thinking", "thinking": "", "text": "leaked", - "signature": "c2ln", + "signature": "", }, ] parts = _convert_reasoning_value_to_parts(thinking_blocks) @@ -5678,6 +5739,76 @@ async def test_content_to_message_param_anthropic_provider_embeds_thinking_block assert result.get("reasoning_content") is None +@pytest.mark.asyncio +async def test_content_to_message_param_anthropic_aggregates_streaming_split_thinking(): + """Streaming splits one Anthropic thinking block across many parts: + text-only chunks followed by a signature-only chunk at block_stop. + _content_to_message_param must re-join them into one thinking_block. + """ + content = types.Content( + role="model", + parts=[ + # Text-only chunks from streaming deltas (no signature) + types.Part(text="The user wants ", thought=True), + types.Part(text="GST research ", thought=True), + types.Part(text="on secondment.", thought=True), + # Final signature-only chunk (empty text, signature carries the whole block) + types.Part( + text="", thought=True, thought_signature=b"ErEDClsIDBACGAIfull" + ), + # Non-thought response content + types.Part.from_function_call(name="create_plan", args={"q": "test"}), + ], + ) + result = await _content_to_message_param( + content, model="anthropic/claude-4-sonnet" + ) + # One aggregated thinking block with combined text and the block's signature + blocks = result["thinking_blocks"] + assert len(blocks) == 1 + assert blocks[0]["type"] == "thinking" + assert blocks[0]["thinking"] == "The user wants GST research on secondment." + assert blocks[0]["signature"] == "RXJFRENsc0lEQkFDR0FJZnVsbA==" + # Legacy reasoning_content is not set when the Anthropic branch takes + assert result.get("reasoning_content") is None + + +def test_model_response_to_chunk_preserves_signature_only_delta(): + """Anthropic streams a final thinking delta where content and + reasoning_content are empty but thinking_blocks carries the signature. + _has_meaningful_signal must recognize thinking_blocks as signal so the + signature survives into a ReasoningChunk. + """ + stream = ModelResponseStream( + id="x", + created=0, + model="claude", + choices=[ + StreamingChoices( + index=0, + delta=Delta( + role=None, + content="", + reasoning_content="", + thinking_blocks=[{ + "type": "thinking", + "thinking": "", + "signature": "SignatureOnlyChunk", + }], + ), + ) + ], + ) + chunks = list(_model_response_to_chunk(stream)) + reasoning_chunks = [c for c, _ in chunks if isinstance(c, ReasoningChunk)] + assert len(reasoning_chunks) == 1 + parts = reasoning_chunks[0].parts + assert len(parts) == 1 + assert parts[0].text == "" + assert parts[0].thought is True + assert parts[0].thought_signature == b"SignatureOnlyChunk" + + @pytest.mark.asyncio @pytest.mark.parametrize( "log_level,should_call",