Python: Stream Bedrock responses with ConverseStream instead of returning them in one update

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,40 @@ 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(
+            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:
             return None

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,56 @@ 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 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"
+    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",

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)