Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 64 additions & 7 deletions src/google/adk/models/lite_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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.

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
141 changes: 136 additions & 5 deletions tests/unittests/models/test_litellm.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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."""
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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",
Expand Down
Loading