diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 7b90803902..455c271ee5 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -420,7 +420,7 @@ async def _stream() -> AsyncIterable[ChatResponseUpdate]: async def _get_response() -> ChatResponse: raw_response = await asyncio.to_thread(self._invoke_converse, request) - return self._process_converse_response(raw_response, options) + return self._process_converse_response(raw_response, options, model=request["modelId"]) return _get_response() @@ -758,9 +758,12 @@ def _redacted_reasoning_content( ) def _process_converse_response( - self, response: dict[str, Any], options: Mapping[str, Any] | None = None + self, response: dict[str, Any], options: Mapping[str, Any] | None = None, *, model: str | None = None ) -> ChatResponse: - """Convert Bedrock Converse API response to ChatResponse.""" + """Convert Bedrock Converse API response to ChatResponse. + + Converse does not return the model id, so ``model`` should be the ``modelId`` the request was sent with. + """ output = response.get("output") or {} message = output.get("message") or {} content_blocks = message.get("content") or [] @@ -770,7 +773,7 @@ def _process_converse_response( usage_details = self._parse_usage(usage_source) finish_reason = self._map_finish_reason(output.get("completionReason") or response.get("stopReason")) response_id = response.get("responseId") or message.get("id") - model = response.get("modelId") or output.get("modelId") or self.model + model = response.get("modelId") or output.get("modelId") or model or self.model return ChatResponse( response_id=response_id, messages=[chat_message], diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index bd698efac7..2f79832820 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -103,6 +103,73 @@ async def test_get_response_invokes_bedrock_runtime() -> None: assert response.usage_details and response.usage_details["input_token_count"] == 10 +class _StubConverseWithoutModelId(_StubBedrockRuntime): + """Returns the real Converse response shape, which does not echo the model id.""" + + def converse(self, **kwargs: Any) -> dict[str, Any]: + self.calls.append(kwargs) + return { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}, + } + + +@pytest.mark.parametrize("client_model", ["amazon.titan-text", None]) +async def test_get_response_reports_the_requested_model(client_model: str | None) -> None: + """A per-call model override must be the model reported on the response, like the streaming path.""" + stub = _StubConverseWithoutModelId() + client = BedrockChatClient( + model=client_model, + region="us-west-2", + client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType] + ) + + response = await client.get_response( + messages=[Message(role="user", contents=[Content.from_text(text="hello")])], + options={"model": "us.openai.gpt-6-sol"}, + ) + + assert stub.calls[0]["modelId"] == "us.openai.gpt-6-sol" + assert response.model == "us.openai.gpt-6-sol" + + +class _StubConverseAndStreamWithoutModelId(_StubConverseWithoutModelId): + def converse_stream(self, **kwargs: Any) -> dict[str, Any]: + self.calls.append(kwargs) + return { + "stream": _StubEventStream([ + {"messageStart": {"role": "assistant"}}, + {"contentBlockDelta": {"delta": {"text": "hi"}, "contentBlockIndex": 0}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + ]) + } + + +@pytest.mark.parametrize("stream", [False, True]) +async def test_streaming_and_non_streaming_report_the_same_model(stream: bool) -> None: + """Both paths must report the per-call model override, not the client default.""" + stub = _StubConverseAndStreamWithoutModelId() + client = BedrockChatClient( + model="amazon.titan-text", + region="us-west-2", + client=stub, # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] # pyright: ignore[reportArgumentType] + ) + messages = [Message(role="user", contents=[Content.from_text(text="hello")])] + options = {"model": "us.openai.gpt-6-sol"} + + if stream: + response_stream = client.get_response(messages=messages, options=options, stream=True) + response = await response_stream.get_final_response() + else: + response = await client.get_response(messages=messages, options=options) + + assert stub.calls[0]["modelId"] == "us.openai.gpt-6-sol" + assert response.text == "hi" + assert response.model == "us.openai.gpt-6-sol" + + 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([