From 6894243bcefa01e9caa1db4e42351d43b4f7716e Mon Sep 17 00:00:00 2001 From: Ehsan Barkhordar Date: Thu, 16 Jul 2026 16:37:22 +0000 Subject: [PATCH] fix(models/gemini): emit thought signature as its own reasoning delta GeminiModel packed "text" and "signature" into a single reasoningContent delta. handle_content_block_delta dispatches on key presence, so the "text" arm always won and the signature never reached the aggregated message. The signature is now emitted as its own delta, matching the convention every other provider already follows. --- strands-py/src/strands/models/gemini.py | 29 ++++++--- .../tests/strands/models/test_gemini.py | 62 +++++++++++++++++-- 2 files changed, 76 insertions(+), 15 deletions(-) diff --git a/strands-py/src/strands/models/gemini.py b/strands-py/src/strands/models/gemini.py index 7b1d7f4522..07cf3c7762 100644 --- a/strands-py/src/strands/models/gemini.py +++ b/strands-py/src/strands/models/gemini.py @@ -396,15 +396,17 @@ def _format_chunk(self, event: dict[str, Any]) -> StreamEvent: "delta": { "reasoningContent": { "text": event["data"].text, - **( - { - "signature": base64.b64encode(event["data"].thought_signature).decode( - "ascii" - ) - } - if event["data"].thought_signature - else {} - ), + }, + }, + }, + } + + case "reasoning_signature": + return { + "contentBlockDelta": { + "delta": { + "reasoningContent": { + "signature": base64.b64encode(event["data"].thought_signature).decode("ascii"), }, }, }, @@ -582,6 +584,15 @@ async def stream( }, ) + if data_type == "reasoning_content" and part.thought_signature: + yield self._format_chunk( + { + "chunk_type": "content_delta", + "data_type": "reasoning_signature", + "data": part, + }, + ) + if data_type is not None: yield self._format_chunk({"chunk_type": "content_stop", "data_type": data_type}) yield self._format_chunk( diff --git a/strands-py/tests/strands/models/test_gemini.py b/strands-py/tests/strands/models/test_gemini.py index 518c7e4ed3..7e9b1ed532 100644 --- a/strands-py/tests/strands/models/test_gemini.py +++ b/strands-py/tests/strands/models/test_gemini.py @@ -744,7 +744,8 @@ async def test_stream_response_reasoning(gemini_client, model, messages, agenera exp_chunks = [ {"messageStart": {"role": "assistant"}}, {"contentBlockStart": {"start": {}}}, - {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "YWJj", "text": "test reason"}}}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "test reason"}}}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "YWJj"}}}}, {"contentBlockStop": {}}, {"messageStop": {"stopReason": "end_turn"}}, {"metadata": {"usage": {"inputTokens": 1, "outputTokens": 2, "totalTokens": 3}, "metrics": {"latencyMs": 0}}}, @@ -752,6 +753,58 @@ async def test_stream_response_reasoning(gemini_client, model, messages, agenera assert tru_chunks == exp_chunks +@pytest.mark.asyncio +async def test_stream_response_reasoning_signature_survives_aggregation( + gemini_client, model, messages, agenerator, alist +): + """Test that a thought signature round-trips from the stream back into a request part. + + Guarantees that a signed thought part keeps its signature through stream aggregation, so the + signature Gemini requires on a subsequent turn is the one it originally issued. + """ + gemini_client.aio.models.generate_content_stream.return_value = agenerator( + [ + genai.types.GenerateContentResponse( + candidates=[ + genai.types.Candidate( + content=genai.types.Content( + parts=[ + genai.types.Part( + text="test reason", + thought=True, + thought_signature=b"abc", + ), + ], + ), + finish_reason="STOP", + ), + ], + usage_metadata=genai.types.GenerateContentResponseUsageMetadata( + prompt_token_count=1, + total_token_count=3, + ), + ), + ] + ) + + stream = strands.event_loop.streaming.process_stream(model.stream(messages)) + events = await alist(stream) + message = events[-1]["stop"][1] + + tru_reasoning = message["content"][0]["reasoningContent"]["reasoningText"] + exp_reasoning = {"text": "test reason", "signature": "YWJj"} + assert tru_reasoning == exp_reasoning + + # The reasoning text must still reach consumers as its own stream event. + tru_reasoning_text = [event["reasoningText"] for event in events if "reasoningText" in event] + exp_reasoning_text = ["test reason"] + assert tru_reasoning_text == exp_reasoning_text + + # Feeding the aggregated message back must reproduce the original signature bytes. + tru_part = model._format_request_content_part(message["content"][0], {}) + assert tru_part.thought_signature == b"abc" + + @pytest.mark.asyncio async def test_stream_response_reasoning_and_text(gemini_client, model, messages, agenerator, alist): """Test that both reasoning and text content are captured in separate blocks.""" @@ -803,11 +856,8 @@ async def test_stream_response_reasoning_and_text(gemini_client, model, messages exp_chunks = [ {"messageStart": {"role": "assistant"}}, {"contentBlockStart": {"start": {}}}, - { - "contentBlockDelta": { - "delta": {"reasoningContent": {"signature": "c2lnMQ==", "text": "thinking about math"}} - } - }, + {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "thinking about math"}}}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "c2lnMQ=="}}}}, {"contentBlockStop": {}}, {"contentBlockStart": {"start": {}}}, {"contentBlockDelta": {"delta": {"text": "2 + 2 = 4"}}},