From 6b49c10fe5ea2da034cc2cb31b2bcbcc95f5331d Mon Sep 17 00:00:00 2001 From: Leela0o5 Date: Thu, 1 Oct 2026 15:44:53 +0530 Subject: [PATCH 1/2] Python: fix(bedrock): keep reasoning content so extended thinking works with tool calls --- .../agent_framework_bedrock/_chat_client.py | 35 +++++++- .../bedrock/tests/test_bedrock_client.py | 90 +++++++++++++++++++ 2 files changed, 124 insertions(+), 1 deletion(-) diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index ea67ecdb3d4..65319ed6492 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -570,7 +570,11 @@ def _convert_message_to_content_blocks(self, message: Message) -> list[dict[str, blocks: list[dict[str, Any]] = [] for content in message.contents: block = self._convert_content_to_bedrock_block(content) - if block is None or ("image" in block and message.role != "user"): + if ( + block is None + or ("image" in block and message.role != "user") + or ("reasoningContent" in block and message.role != "assistant") + ): logger.debug("Skipping unsupported content type for Bedrock: %s", type(content)) continue blocks.append(block) @@ -587,6 +591,12 @@ def _convert_content_to_bedrock_block(self, content: Content) -> dict[str, Any] logger.warning("Skipping %s image: Bedrock accepts gif, jpeg, png or webp.", content.media_type) return None return {"image": {"format": image_format, "source": {"bytes": _get_data_bytes(content)}}} + case "text_reasoning" if content.text is not None: + # Extended thinking with tool use requires the signed reasoning to be sent back unchanged. + reasoning_text: dict[str, Any] = {"text": content.text} + if content.protected_data: + reasoning_text["signature"] = content.protected_data + return {"reasoningContent": {"reasoningText": reasoning_text}} case "function_call": arguments = content.parse_arguments() or {} return { @@ -733,6 +743,17 @@ def _process_converse_stream_event( arguments=tool_use_delta.get("input", ""), ) ) + elif reasoning_delta := delta.get("reasoningContent"): + # Reasoning text and its signature arrive in separate deltas; they are merged into one + # text_reasoning content when the updates are combined into the final response. + reasoning_text = reasoning_delta.get("text") + signature = reasoning_delta.get("signature") + if reasoning_text is not None or signature is not None: + contents.append( + Content.from_text_reasoning( + text=reasoning_text, protected_data=signature, raw_representation=delta + ) + ) elif message_stop := event.get("messageStop"): finish_reason = self._map_finish_reason(message_stop.get("stopReason")) elif (metadata := event.get("metadata")) and (usage_details := self._parse_usage(metadata.get("usage"))): @@ -771,6 +792,18 @@ def _parse_message_contents(self, content_blocks: Sequence[dict[str, Any]]) -> l Content.from_text(text=json.dumps(json_value, ensure_ascii=False), raw_representation=block) ) continue + if isinstance(reasoning := block.get("reasoningContent"), Mapping): + if isinstance(reasoning_text := reasoning.get("reasoningText"), Mapping): + contents.append( + Content.from_text_reasoning( + text=reasoning_text.get("text") or "", + protected_data=reasoning_text.get("signature"), + raw_representation=block, + ) + ) + else: + logger.debug("Ignoring Bedrock reasoning block without reasoning text: %s", block) + continue tool_use_value = block.get("toolUse") tool_use = ( tool_use_value diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index 5e25cad9692..946d7e694ec 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -148,6 +148,96 @@ async def test_stream_yields_updates_as_converse_stream_events_arrive() -> None: assert stub.stream.closed +def test_process_converse_response_parses_reasoning_content() -> None: + client = _make_client() + + response = client._process_converse_response({ + "output": { + "message": { + "role": "assistant", + "content": [ + {"reasoningContent": {"reasoningText": {"text": "Need the weather tool.", "signature": "sig-1"}}}, + {"toolUse": {"toolUseId": "call-1", "name": "get_weather", "input": {"city": "Paris"}}}, + ], + } + }, + "stopReason": "tool_use", + }) + + reasoning, function_call = response.messages[0].contents + assert reasoning.type == "text_reasoning" + assert reasoning.text == "Need the weather tool." + assert reasoning.protected_data == "sig-1" + assert function_call.type == "function_call" + + +async def test_stream_merges_reasoning_text_and_signature() -> None: + stub = _StubBedrockStreamRuntime([ + {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "Let me "}}, "contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "think."}}, "contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig-1"}}, "contentBlockIndex": 0}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"text": "Done."}, "contentBlockIndex": 1}}, + {"messageStop": {"stopReason": "end_turn"}}, + ]) + client = BedrockChatClient( + model="anthropic.claude-sonnet-4", + region="us-east-1", + client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType] + ) + + stream = client._inner_get_response( + messages=[Message(role="user", contents=[Content.from_text(text="hi")])], options={}, stream=True + ) + assert isinstance(stream, ResponseStream) + async for _ in stream: + pass + response = await stream.get_final_response() + + reasoning = [content for content in response.messages[0].contents if content.type == "text_reasoning"] + assert len(reasoning) == 1 + assert reasoning[0].text == "Let me think." + assert reasoning[0].protected_data == "sig-1" + assert response.text == "Done." + + +def test_prepare_bedrock_messages_sends_reasoning_back_with_tool_use() -> None: + """Extended thinking with tool use fails unless the signed reasoning is replayed before the toolUse block.""" + client = _make_client() + messages = [ + Message(role="user", contents=[Content.from_text(text="Weather in Paris?")]), + Message( + role="assistant", + contents=[ + Content.from_text_reasoning(text="Need the weather tool.", protected_data="sig-1"), + Content.from_function_call(call_id="call-1", name="get_weather", arguments={"city": "Paris"}), + ], + ), + Message(role="tool", contents=[Content.from_function_result(call_id="call-1", result="88F")]), + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + assert conversation[1]["content"][0] == { + "reasoningContent": {"reasoningText": {"text": "Need the weather tool.", "signature": "sig-1"}} + } + assert "toolUse" in conversation[1]["content"][1] + + +def test_prepare_bedrock_messages_skips_reasoning_outside_assistant_messages() -> None: + client = _make_client() + messages = [ + Message( + role="user", + contents=[Content.from_text(text="hi"), Content.from_text_reasoning(text="not from the model")], + ) + ] + + _, conversation = client._prepare_bedrock_messages(messages) + + assert conversation[0]["content"] == [{"text": "hi"}] + + def test_build_request_requires_non_system_messages() -> None: client = BedrockChatClient( model="amazon.titan-text", From 367a89cd6a29545cd39a956540339c060b4f3040 Mon Sep 17 00:00:00 2001 From: Leela0o5 Date: Sun, 4 Oct 2026 10:15:10 +0530 Subject: [PATCH 2/2] Python: fix(bedrock): replay redacted reasoning and keep streamed reasoning blocks separate --- .../agent_framework_bedrock/_chat_client.py | 85 +++++++++++-- .../bedrock/tests/test_bedrock_client.py | 118 ++++++++++++++++++ 2 files changed, 192 insertions(+), 11 deletions(-) diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 65319ed6492..2d822fcfe6a 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -4,12 +4,14 @@ from __future__ import annotations import asyncio +import base64 import copy import json import logging import sys from collections import deque from collections.abc import AsyncIterable, Awaitable, Mapping, MutableMapping, Sequence +from dataclasses import dataclass from typing import Any, ClassVar, Generic, Literal, TypedDict from uuid import uuid4 @@ -58,6 +60,18 @@ logger = logging.getLogger("agent_framework.bedrock") +# Marks text_reasoning content whose protected_data is base64 of a Bedrock ``redactedContent`` blob, +# as opposed to the signature of a ``reasoningText`` block. +_REDACTED_REASONING_KEY = "bedrock.redacted_reasoning" + + +@dataclass +class _StreamedReasoningBlock: + """State for one reasoning content block while its ConverseStream deltas arrive.""" + + id: str + redacted: bytes = b"" + __all__ = [ "BedrockChatClient", @@ -382,9 +396,12 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: events = response["stream"] event_iterator = iter(events) tool_call_ids: dict[int, str] = {} + reasoning_blocks: dict[int, _StreamedReasoningBlock] = {} try: while (event := await asyncio.to_thread(next, event_iterator, None)) is not None: - if update := self._process_converse_stream_event(event, request["modelId"], tool_call_ids): + if update := self._process_converse_stream_event( + event, request["modelId"], tool_call_ids, reasoning_blocks + ): yield update finally: events.close() @@ -591,8 +608,17 @@ def _convert_content_to_bedrock_block(self, content: Content) -> dict[str, Any] logger.warning("Skipping %s image: Bedrock accepts gif, jpeg, png or webp.", content.media_type) return None return {"image": {"format": image_format, "source": {"bytes": _get_data_bytes(content)}}} - case "text_reasoning" if content.text is not None: - # Extended thinking with tool use requires the signed reasoning to be sent back unchanged. + case "text_reasoning": + # Extended thinking with tool use requires reasoning blocks to be sent back unchanged. + if content.additional_properties.get(_REDACTED_REASONING_KEY): + try: + redacted = base64.b64decode(content.protected_data or "", validate=True) + except ValueError: + logger.warning("Skipping redacted reasoning for Bedrock: protected_data is not valid base64") + return None + return {"reasoningContent": {"redactedContent": redacted}} if redacted else None + if content.text is None: + return None reasoning_text: dict[str, Any] = {"text": content.text} if content.protected_data: reasoning_text["signature"] = content.protected_data @@ -695,6 +721,19 @@ def _prepare_tools(self, tools: list[FunctionTool | MutableMapping[str, Any]] | def _generate_tool_call_id() -> str: return f"tool-call-{uuid4().hex}" + @staticmethod + def _generate_reasoning_id() -> str: + return f"reasoning-{uuid4().hex}" + + @staticmethod + def _redacted_reasoning_content(redacted: bytes, *, id: str, raw_representation: Any) -> Content: + return Content.from_text_reasoning( + id=id, + protected_data=base64.b64encode(redacted).decode("utf-8"), + additional_properties={_REDACTED_REASONING_KEY: True}, + raw_representation=raw_representation, + ) + def _process_converse_response( self, response: dict[str, Any], options: Mapping[str, Any] | None = None ) -> ChatResponse: @@ -720,7 +759,11 @@ def _process_converse_response( ) def _process_converse_stream_event( - self, event: Mapping[str, Any], model: str, tool_call_ids: dict[int, str] + self, + event: Mapping[str, Any], + model: str, + tool_call_ids: dict[int, str], + reasoning_blocks: dict[int, _StreamedReasoningBlock], ) -> ChatResponseUpdate | None: """Convert a single Bedrock ConverseStream event to a ChatResponseUpdate.""" contents: list[Content] = [] @@ -744,14 +787,26 @@ def _process_converse_stream_event( ) ) elif reasoning_delta := delta.get("reasoningContent"): - # Reasoning text and its signature arrive in separate deltas; they are merged into one - # text_reasoning content when the updates are combined into the final response. - reasoning_text = reasoning_delta.get("text") - signature = reasoning_delta.get("signature") - if reasoning_text is not None or signature is not None: + # The deltas of one reasoning block share an id, so they merge into a single text_reasoning + # content in the final response while consecutive blocks keep their own text and signature. + block = reasoning_blocks.setdefault( + block_delta.get("contentBlockIndex", 0), + _StreamedReasoningBlock(id=self._generate_reasoning_id()), + ) + if isinstance(redacted := reasoning_delta.get("redactedContent"), (bytes, bytearray)): + # Merging keeps only the latest protected_data, so send everything received so far. + block.redacted += bytes(redacted) + contents.append( + self._redacted_reasoning_content(block.redacted, id=block.id, raw_representation=delta) + ) + elif "text" in reasoning_delta or "signature" in reasoning_delta: contents.append( Content.from_text_reasoning( - text=reasoning_text, protected_data=signature, raw_representation=delta + id=block.id, + # "" rather than None, so a signed block without text is still replayed. + text=reasoning_delta.get("text") or "", + protected_data=reasoning_delta.get("signature"), + raw_representation=delta, ) ) elif message_stop := event.get("messageStop"): @@ -793,16 +848,24 @@ def _parse_message_contents(self, content_blocks: Sequence[dict[str, Any]]) -> l ) continue if isinstance(reasoning := block.get("reasoningContent"), Mapping): + # Each block gets its own id so consecutive reasoning blocks are never merged into one. if isinstance(reasoning_text := reasoning.get("reasoningText"), Mapping): contents.append( Content.from_text_reasoning( + id=self._generate_reasoning_id(), text=reasoning_text.get("text") or "", protected_data=reasoning_text.get("signature"), raw_representation=block, ) ) + elif isinstance(redacted := reasoning.get("redactedContent"), (bytes, bytearray)): + contents.append( + self._redacted_reasoning_content( + bytes(redacted), id=self._generate_reasoning_id(), raw_representation=block + ) + ) else: - logger.debug("Ignoring Bedrock reasoning block without reasoning text: %s", block) + logger.debug("Ignoring unsupported Bedrock reasoning block: %s", block) continue tool_use_value = block.get("toolUse") tool_use = ( diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index 946d7e694ec..f37ff135b6a 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -238,6 +238,124 @@ def test_prepare_bedrock_messages_skips_reasoning_outside_assistant_messages() - assert conversation[0]["content"] == [{"text": "hi"}] +async def test_stream_keeps_consecutive_reasoning_blocks_separate() -> None: + """Each streamed reasoning block must keep its own text and signature when replayed.""" + stub = _StubBedrockStreamRuntime([ + {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "First "}}, "contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "thought."}}, "contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig-1"}}, "contentBlockIndex": 0}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "Second thought."}}, "contentBlockIndex": 1}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig-2"}}, "contentBlockIndex": 1}}, + {"contentBlockStop": {"contentBlockIndex": 1}}, + { + "contentBlockStart": { + "start": {"toolUse": {"toolUseId": "call-1", "name": "get_weather"}}, + "contentBlockIndex": 2, + } + }, + {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"city": "Paris"}'}}, "contentBlockIndex": 2}}, + {"messageStop": {"stopReason": "tool_use"}}, + ]) + client = BedrockChatClient( + model="anthropic.claude-sonnet-4", + region="us-east-1", + client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType] + ) + + stream = client._inner_get_response( + messages=[Message(role="user", contents=[Content.from_text(text="Weather in Paris?")])], + options={}, + stream=True, + ) + assert isinstance(stream, ResponseStream) + async for _ in stream: + pass + response = await stream.get_final_response() + + reasoning = [content for content in response.messages[0].contents if content.type == "text_reasoning"] + assert [(content.text, content.protected_data) for content in reasoning] == [ + ("First thought.", "sig-1"), + ("Second thought.", "sig-2"), + ] + + _, conversation = client._prepare_bedrock_messages(response.messages) + assert conversation[0]["content"][:2] == [ + {"reasoningContent": {"reasoningText": {"text": "First thought.", "signature": "sig-1"}}}, + {"reasoningContent": {"reasoningText": {"text": "Second thought.", "signature": "sig-2"}}}, + ] + assert "toolUse" in conversation[0]["content"][2] + + +def test_redacted_reasoning_content_is_replayed_unchanged() -> None: + client = _make_client() + redacted = b"\x00\x01encrypted-reasoning\xff" + + response = client._process_converse_response({ + "output": { + "message": { + "role": "assistant", + "content": [ + {"reasoningContent": {"redactedContent": redacted}}, + {"toolUse": {"toolUseId": "call-1", "name": "get_weather", "input": {"city": "Paris"}}}, + ], + } + }, + "stopReason": "tool_use", + }) + + reasoning = response.messages[0].contents[0] + assert reasoning.type == "text_reasoning" + assert reasoning.text is None + + _, conversation = client._prepare_bedrock_messages(response.messages) + assert conversation[0]["content"][0] == {"reasoningContent": {"redactedContent": redacted}} + assert "toolUse" in conversation[0]["content"][1] + + +async def test_stream_redacted_reasoning_content_is_replayed_unchanged() -> None: + stub = _StubBedrockStreamRuntime([ + { + "contentBlockDelta": { + "delta": {"reasoningContent": {"redactedContent": b"\x00\x01part"}}, + "contentBlockIndex": 0, + } + }, + { + "contentBlockDelta": { + "delta": {"reasoningContent": {"redactedContent": b"-two\xff"}}, + "contentBlockIndex": 0, + } + }, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"text": "Visible."}}, "contentBlockIndex": 1}}, + {"contentBlockDelta": {"delta": {"reasoningContent": {"signature": "sig-1"}}, "contentBlockIndex": 1}}, + {"contentBlockStop": {"contentBlockIndex": 1}}, + {"contentBlockDelta": {"delta": {"text": "Done."}, "contentBlockIndex": 2}}, + {"messageStop": {"stopReason": "end_turn"}}, + ]) + client = BedrockChatClient( + model="anthropic.claude-3-7-sonnet", + region="us-east-1", + client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType] + ) + + stream = client._inner_get_response( + messages=[Message(role="user", contents=[Content.from_text(text="hi")])], options={}, stream=True + ) + assert isinstance(stream, ResponseStream) + async for _ in stream: + pass + response = await stream.get_final_response() + + _, conversation = client._prepare_bedrock_messages(response.messages) + assert conversation[0]["content"] == [ + {"reasoningContent": {"redactedContent": b"\x00\x01part-two\xff"}}, + {"reasoningContent": {"reasoningText": {"text": "Visible.", "signature": "sig-1"}}}, + {"text": "Done."}, + ] + + def test_build_request_requires_non_system_messages() -> None: client = BedrockChatClient( model="amazon.titan-text",