diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 2e4cf65f32..c996971618 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,17 +1062,24 @@ 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): - sig = base64.b64encode(sig).decode("utf-8") + sig_str = base64.b64encode(sig).decode("utf-8") + else: + sig_str = str(sig) thinking_blocks.append({ "type": "thinking", "thinking": part.text, - "signature": sig, + "signature": sig_str, }) if thinking_blocks: msg = ChatCompletionAssistantMessage( @@ -1877,6 +1923,11 @@ def _has_meaningful_signal(message: Message | Delta | None) -> bool: or message.get("function_call") or message.get("reasoning_content") or message.get("reasoning") + # Anthropic streams the thinking block's signature in a final delta + # where content/reasoning_content are empty and only thinking_blocks + # carries the signature. Without this, the delta is discarded before + # _extract_reasoning_value can preserve the signature. + or message.get("thinking_blocks") ) if isinstance(response, ModelResponseStream): @@ -2717,7 +2768,10 @@ def _finalize_tool_call_response( ) mapped_finish_reason = _map_finish_reason(finish_reason) llm_response.finish_reason = mapped_finish_reason - if mapped_finish_reason != types.FinishReason.STOP: + if ( + mapped_finish_reason is not None + and mapped_finish_reason != types.FinishReason.STOP + ): llm_response.error_code = mapped_finish_reason llm_response.error_message = _finish_reason_to_error_message( mapped_finish_reason @@ -2738,7 +2792,10 @@ def _finalize_text_response( ) mapped_finish_reason = _map_finish_reason(finish_reason) llm_response.finish_reason = mapped_finish_reason - if mapped_finish_reason != types.FinishReason.STOP: + if ( + mapped_finish_reason is not None + and mapped_finish_reason != types.FinishReason.STOP + ): llm_response.error_code = mapped_finish_reason llm_response.error_message = _finish_reason_to_error_message( mapped_finish_reason 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",