From b73066cb404a358617a694e76972c2d4dd0ed32e Mon Sep 17 00:00:00 2001 From: kimnamu Date: Wed, 23 Sep 2026 18:22:37 +0900 Subject: [PATCH 1/2] Python: Stream Bedrock responses with ConverseStream instead of returning them in one update stream=True called the non-streaming Converse API and yielded a single update after the whole response had arrived, so time to first text was the full generation time. Call converse_stream and yield an update per event: text deltas, tool use start and input deltas (merged by call_id), messageStop finish reason and metadata usage. The event stream is closed when iteration ends. The non-streaming path is unchanged. Co-Authored-By: Claude --- .../agent_framework_bedrock/_chat_client.py | 65 ++++++++++++------ .../bedrock/tests/test_bedrock_client.py | 68 ++++++++++++++++++- .../tests/test_bedrock_structured_output.py | 9 +++ 3 files changed, 122 insertions(+), 20 deletions(-) diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 6d2d222d1dd..04f34327870 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -332,10 +332,11 @@ def _create_session(settings: BedrockSettings) -> Boto3Session: session_kwargs["aws_session_token"] = session_token.get_secret_value() return Boto3Session(**session_kwargs) - def _invoke_converse(self, request: Mapping[str, Any]) -> dict[str, Any]: + def _invoke_converse(self, request: Mapping[str, Any], stream: bool = False) -> dict[str, Any]: mark_feature_used(FeatureIndex.BEDROCK) try: - response = self._bedrock_client.converse(**request) + invoke = self._bedrock_client.converse_stream if stream else self._bedrock_client.converse + response = invoke(**request) if not isinstance(response, Mapping): raise ChatClientInvalidResponseException("Bedrock converse response must be a mapping.") return response @@ -369,24 +370,18 @@ def _inner_get_response( request = self._prepare_options(messages, options, **kwargs) if stream: - # Streaming mode - simulate streaming by yielding a single update + # Streaming mode - yield an update per ConverseStream event async def _stream() -> AsyncIterable[ChatResponseUpdate]: - response = await asyncio.to_thread(self._invoke_converse, request) - parsed_response = self._process_converse_response(response, options) - contents = list(parsed_response.messages[0].contents if parsed_response.messages else []) - if parsed_response.usage_details: - contents.append(Content.from_usage(usage_details=parsed_response.usage_details)) - raw_finish_reason = ( - parsed_response.finish_reason if isinstance(parsed_response.finish_reason, str) else None - ) - finish_reason = self._map_finish_reason(raw_finish_reason) - yield ChatResponseUpdate( - response_id=parsed_response.response_id, - contents=contents, - model=parsed_response.model, - finish_reason=finish_reason, - raw_representation=parsed_response.raw_representation, - ) + response = await asyncio.to_thread(self._invoke_converse, request, True) + events = response["stream"] + event_iterator = iter(events) + tool_call_ids: dict[int, str] = {} + 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): + yield update + finally: + events.close() return self._build_response_stream(_stream(), response_format=options.get("response_format")) @@ -672,6 +667,38 @@ def _process_converse_response( raw_representation=response, ) + def _process_converse_stream_event( + self, event: Mapping[str, Any], model: str, tool_call_ids: dict[int, str] + ) -> ChatResponseUpdate | None: + """Convert a single Bedrock ConverseStream event to a ChatResponseUpdate.""" + contents: list[Content] = [] + finish_reason = None + if block_start := event.get("contentBlockStart"): + if tool_use := block_start.get("start", {}).get("toolUse"): + tool_call_ids[block_start.get("contentBlockIndex", 0)] = tool_use["toolUseId"] + contents.append( + Content.from_function_call(call_id=tool_use["toolUseId"], name=tool_use["name"], arguments="") + ) + elif block_delta := event.get("contentBlockDelta"): + delta = block_delta.get("delta", {}) + if text_value := delta.get("text"): + contents.append(Content.from_text(text=text_value, raw_representation=delta)) + elif tool_use_delta := delta.get("toolUse"): + contents.append( + Content.from_function_call( + call_id=tool_call_ids.get(block_delta.get("contentBlockIndex", 0), ""), + name="", + arguments=tool_use_delta.get("input", ""), + ) + ) + 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"))): + contents.append(Content.from_usage(usage_details=usage_details)) + if not contents and finish_reason is None: + return None + return ChatResponseUpdate(contents=contents, model=model, finish_reason=finish_reason, raw_representation=event) + def _parse_usage(self, usage: dict[str, Any] | None) -> UsageDetails | None: if not usage: return None diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index ee17483e84d..916e4f2a831 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -9,7 +9,7 @@ from unittest.mock import MagicMock, patch import pytest -from agent_framework import Agent, Content, FunctionTool, Message +from agent_framework import Agent, Content, FunctionTool, Message, ResponseStream from agent_framework._settings import SecretString from boto3.session import Session as Boto3Session from botocore.client import BaseClient @@ -40,6 +40,23 @@ def converse(self, **kwargs: Any) -> dict[str, Any]: } +class _StubEventStream(list[dict[str, Any]]): + closed = False + + def close(self) -> None: + self.closed = True + + +class _StubBedrockStreamRuntime(_StubBedrockRuntime): + def __init__(self, events: list[dict[str, Any]]) -> None: + super().__init__() + self.stream = _StubEventStream(events) + + def converse_stream(self, **kwargs: Any) -> dict[str, Any]: + self.calls.append(kwargs) + return {"stream": self.stream} + + def _make_client() -> BedrockChatClient: """Create a BedrockChatClient with a stub runtime for unit tests.""" return BedrockChatClient( @@ -80,6 +97,55 @@ async def test_get_response_invokes_bedrock_runtime() -> None: assert response.usage_details and response.usage_details["input_token_count"] == 10 +async def test_stream_yields_updates_as_converse_stream_events_arrive() -> None: + """stream=True should use ConverseStream and yield text, tool call, finish and usage updates per event.""" + stub = _StubBedrockStreamRuntime([ + {"messageStart": {"role": "assistant"}}, + {"contentBlockDelta": {"delta": {"text": "Checking"}, "contentBlockIndex": 0}}, + {"contentBlockDelta": {"delta": {"text": " the weather."}, "contentBlockIndex": 0}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + { + "contentBlockStart": { + "start": {"toolUse": {"toolUseId": "call-1", "name": "get_weather"}}, + "contentBlockIndex": 1, + } + }, + {"contentBlockDelta": {"delta": {"toolUse": {"input": '{"city":'}}, "contentBlockIndex": 1}}, + {"contentBlockDelta": {"delta": {"toolUse": {"input": ' "Rome"}'}}, "contentBlockIndex": 1}}, + {"contentBlockStop": {"contentBlockIndex": 1}}, + {"messageStop": {"stopReason": "tool_use"}}, + { + "metadata": { + "usage": {"inputTokens": 47, "outputTokens": 18, "totalTokens": 65}, + "metrics": {"latencyMs": 9}, + } + }, + ]) + client = BedrockChatClient( + model="us.openai.gpt-6-sol", + 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 Rome?")])], options={}, stream=True + ) + assert isinstance(stream, ResponseStream) + updates = [update async for update in stream] + response = await stream.get_final_response() + + assert [update.text for update in updates if update.text] == ["Checking", " the weather."] + assert response.text == "Checking the weather." + function_call = next(content for content in response.messages[0].contents if content.type == "function_call") + assert function_call.call_id == "call-1" + assert function_call.name == "get_weather" + assert function_call.parse_arguments() == {"city": "Rome"} + assert response.finish_reason == "tool_calls" + assert response.usage_details and response.usage_details["output_token_count"] == 18 + assert response.model == "us.openai.gpt-6-sol" + assert stub.stream.closed + + def test_build_request_requires_non_system_messages() -> None: client = BedrockChatClient( model="amazon.titan-text", diff --git a/python/packages/bedrock/tests/test_bedrock_structured_output.py b/python/packages/bedrock/tests/test_bedrock_structured_output.py index 1c4f2aeca65..9494e0bb51c 100644 --- a/python/packages/bedrock/tests/test_bedrock_structured_output.py +++ b/python/packages/bedrock/tests/test_bedrock_structured_output.py @@ -64,6 +64,15 @@ def converse(self, **kwargs: Any) -> dict[str, Any]: }, } + def converse_stream(self, **kwargs: Any) -> dict[str, Any]: + self.calls.append(kwargs) + stream = MagicMock() + stream.__iter__.return_value = iter([ + {"contentBlockDelta": {"delta": {"text": self._response_text}, "contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + ]) + return {"stream": stream} + def _make_client(response_text: str = "Bedrock says hi") -> tuple[BedrockChatClient, _StubBedrockRuntime]: stub = _StubBedrockRuntime(response_text) From 3885687bd7177b484f43287e46455214dea64e9a Mon Sep 17 00:00:00 2001 From: kimnamu Date: Thu, 24 Sep 2026 17:08:25 +0900 Subject: [PATCH 2/2] Python: Set role on Bedrock streaming updates Copilot's review of #8685 noted that the ConverseStream updates carry no role. The other chat clients tag their streaming updates with role="assistant" (OpenAI, Mistral, Ollama), and the non-streaming Bedrock path already builds Message(role="assistant"). Without it, Agent.run(stream=True) yielded AgentResponseUpdate.role=None for Bedrock while the other providers yield "assistant". Set role="assistant" on every streamed update. The aggregated ChatResponse already defaulted to assistant and is unchanged. Co-Authored-By: Claude --- .../packages/bedrock/agent_framework_bedrock/_chat_client.py | 4 +++- python/packages/bedrock/tests/test_bedrock_client.py | 1 + 2 files changed, 4 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 04f34327870..7e0613f2027 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -697,7 +697,9 @@ def _process_converse_stream_event( contents.append(Content.from_usage(usage_details=usage_details)) if not contents and finish_reason is None: return None - return ChatResponseUpdate(contents=contents, model=model, finish_reason=finish_reason, raw_representation=event) + return ChatResponseUpdate( + role="assistant", contents=contents, model=model, finish_reason=finish_reason, raw_representation=event + ) def _parse_usage(self, usage: dict[str, Any] | None) -> UsageDetails | None: if not usage: diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index 916e4f2a831..d74ba8d19ba 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -135,6 +135,7 @@ async def test_stream_yields_updates_as_converse_stream_events_arrive() -> None: response = await stream.get_final_response() assert [update.text for update in updates if update.text] == ["Checking", " the weather."] + assert all(update.role == "assistant" for update in updates) assert response.text == "Checking the weather." function_call = next(content for content in response.messages[0].contents if content.type == "function_call") assert function_call.call_id == "call-1"