diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py index 3a6c36624d..8192ea0925 100644 --- a/src/google/adk/models/lite_llm.py +++ b/src/google/adk/models/lite_llm.py @@ -398,15 +398,47 @@ def _iter_reasoning_texts(reasoning_value: Any) -> Iterable[str]: def _is_thinking_blocks_format(reasoning_value: Any) -> bool: - """Returns True if reasoning_value is Anthropic thinking_blocks format. + """Returns True if reasoning_value is structured thinking_blocks format. - Anthropic thinking_blocks is a list of dicts, each with 'type', 'thinking', - and 'signature' keys. + Anthropic blocks include per-block signatures. Gemini blocks can use the same + type/thinking shape while carrying signatures separately on the message. """ if not isinstance(reasoning_value, list) or not reasoning_value: return False first = reasoning_value[0] - return isinstance(first, dict) and "signature" in first + return isinstance(first, dict) and ( + "thinking" in first or "signature" in first + ) + + +def _with_parallel_thought_signatures( + thinking_blocks: Any, message: Message | Delta +) -> Any: + """Adds Gemini's parallel thought signatures to thinking_blocks.""" + if not isinstance(thinking_blocks, list) or not thinking_blocks: + return thinking_blocks + first = thinking_blocks[0] + if not isinstance(first, dict) or "signature" in first: + return thinking_blocks + + provider_fields = message.get("provider_specific_fields") or {} + if not isinstance(provider_fields, dict): + return thinking_blocks + signatures = provider_fields.get("thought_signatures") or [] + if not isinstance(signatures, list) or not signatures: + return thinking_blocks + + merged = [] + for index, block in enumerate(thinking_blocks): + if ( + isinstance(block, dict) + and index < len(signatures) + and signatures[index] + ): + merged.append({**block, "signature": signatures[index]}) + else: + merged.append(block) + return merged def _convert_reasoning_value_to_parts(reasoning_value: Any) -> List[types.Part]: @@ -455,7 +487,7 @@ def _extract_reasoning_value(message: Message | Delta | None) -> Any: # This must be preserved to maintain thinking across tool call boundaries. thinking_blocks = message.get("thinking_blocks") if thinking_blocks is not None: - return thinking_blocks + return _with_parallel_thought_signatures(thinking_blocks, message) reasoning_content = message.get("reasoning_content") if reasoning_content is not None: return reasoning_content diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py index c195076349..b46792946d 100644 --- a/tests/unittests/models/test_litellm.py +++ b/tests/unittests/models/test_litellm.py @@ -4797,6 +4797,26 @@ def test_extract_reasoning_value_prefers_thinking_blocks(): assert result is thinking_blocks +def test_extract_reasoning_value_gemini_thinking_blocks_zips_signatures(): + """Gemini keeps thought signatures beside thinking_blocks.""" + message = { + "role": "assistant", + "content": "Answer", + "thinking_blocks": [ + {"type": "thinking", "thinking": "step 1"}, + {"type": "thinking", "thinking": "step 2"}, + ], + "provider_specific_fields": { + "thought_signatures": ["sig-1", "sig-2"], + }, + } + result = _extract_reasoning_value(message) + assert result == [ + {"type": "thinking", "thinking": "step 1", "signature": "sig-1"}, + {"type": "thinking", "thinking": "step 2", "signature": "sig-2"}, + ] + + def test_extract_reasoning_value_falls_back_without_thinking_blocks(): """When thinking_blocks is absent, falls back to reasoning_content.""" message = { @@ -4845,6 +4865,18 @@ def test_convert_reasoning_value_to_parts_skips_empty_thinking(): assert parts[0].text == "real thought" +def test_convert_reasoning_value_to_parts_gemini_blocks_without_signature(): + """Gemini thinking_blocks should still produce thought parts.""" + thinking_blocks = [ + {"type": "thinking", "thinking": "gemini thought"}, + ] + parts = _convert_reasoning_value_to_parts(thinking_blocks) + assert len(parts) == 1 + assert parts[0].text == "gemini thought" + assert parts[0].thought is True + assert parts[0].thought_signature is None + + def test_convert_reasoning_value_to_parts_flat_string_unchanged(): """Flat string reasoning still produces thought parts without signature.""" parts = _convert_reasoning_value_to_parts("simple reasoning text") @@ -4854,6 +4886,27 @@ def test_convert_reasoning_value_to_parts_flat_string_unchanged(): assert parts[0].thought_signature is None +def test_message_to_generate_content_response_gemini_thinking_blocks(): + """Gemini thinking_blocks are surfaced before visible text.""" + message = { + "role": "assistant", + "content": "I am a large language model.", + "thinking_blocks": [ + {"type": "thinking", "thinking": "Identity check"}, + ], + "provider_specific_fields": { + "thought_signatures": ["sig-gemini"], + }, + } + response = _message_to_generate_content_response(message) + parts = response.content.parts + assert len(parts) == 2 + assert parts[0].thought is True + assert parts[0].text == "Identity check" + assert parts[0].thought_signature == b"sig-gemini" + assert parts[1].text == "I am a large language model." + + @pytest.mark.asyncio async def test_content_to_message_param_anthropic_outputs_thinking_blocks(): """For Anthropic models, thinking_blocks are output instead of reasoning_content."""