From 8942c9874710fb1cf6efa79de7a6200253ba89e7 Mon Sep 17 00:00:00 2001 From: Aleksei Iancheruk Date: Fri, 21 Nov 2025 14:16:41 +0100 Subject: [PATCH 1/4] feat(bedrock): add guardrail_last_turn_only option --- src/strands/models/bedrock.py | 53 ++++++++++++++++++++++++- tests/strands/models/test_bedrock.py | 44 +++++++++++++++++++++ tests_integ/test_bedrock_guardrails.py | 54 ++++++++++++++++++++++++++ 3 files changed, 149 insertions(+), 2 deletions(-) diff --git a/src/strands/models/bedrock.py b/src/strands/models/bedrock.py index 08d8f400c..3225254f1 100644 --- a/src/strands/models/bedrock.py +++ b/src/strands/models/bedrock.py @@ -82,6 +82,8 @@ class BedrockConfig(TypedDict, total=False): guardrail_redact_input_message: If a Bedrock Input guardrail triggers, replace the input with this message. guardrail_redact_output: Flag to redact output if guardrail is triggered. Defaults to False. guardrail_redact_output_message: If a Bedrock Output guardrail triggers, replace output with this message. + guardrail_last_turn_only: Flag to send only the last turn to guardrails instead of full conversation. + Defaults to False. max_tokens: Maximum number of tokens to generate in the response model_id: The Bedrock model ID (e.g., "us.anthropic.claude-sonnet-4-20250514-v1:0") include_tool_result_status: Flag to include status field in tool results. @@ -105,6 +107,7 @@ class BedrockConfig(TypedDict, total=False): guardrail_redact_input_message: Optional[str] guardrail_redact_output: Optional[bool] guardrail_redact_output_message: Optional[str] + guardrail_last_turn_only: Optional[bool] max_tokens: Optional[int] model_id: str include_tool_result_status: Optional[Literal["auto"] | bool] @@ -206,9 +209,19 @@ def _format_request( Returns: A Bedrock converse stream request. """ + # Filter messages for guardrails if guardrail_last_turn_only is enabled + messages_for_request = messages + if ( + self.config.get("guardrail_last_turn_only", False) + and self.config.get("guardrail_id") + and self.config.get("guardrail_version") + ): + messages_for_request = self._get_last_turn_messages(messages) + if not tool_specs: has_tool_content = any( - any("toolUse" in block or "toolResult" in block for block in msg.get("content", [])) for msg in messages + any("toolUse" in block or "toolResult" in block for block in msg.get("content", [])) + for msg in messages_for_request ) if has_tool_content: tool_specs = [noop_tool.tool_spec] @@ -224,7 +237,7 @@ def _format_request( return { "modelId": self.config["model_id"], - "messages": self._format_bedrock_messages(messages), + "messages": self._format_bedrock_messages(messages_for_request), "system": system_blocks, **( { @@ -295,6 +308,42 @@ def _format_request( ), } + def _get_last_turn_messages(self, messages: Messages) -> Messages: + """Get the last turn messages for guardrail evaluation. + + Returns the latest user message and the previous assistant message (if it exists). + This reduces the conversation context sent to guardrails when guardrail_last_turn_only is True. + + Args: + messages: Full conversation messages. + + Returns: + Messages containing only the last turn (user + previous assistant if exists). + """ + if not messages: + return [] + + # Find the last user message + last_user_index = -1 + for i in range(len(messages) - 1, -1, -1): + if messages[i]["role"] == "user": + last_user_index = i + break + + if last_user_index == -1: + # No user message found, return empty + return [] + + # Include the previous assistant message if it exists + result_messages: Messages = [] + if last_user_index > 0 and messages[last_user_index - 1]["role"] == "assistant": + result_messages.append(messages[last_user_index - 1]) + + # Add the last user message + result_messages.append(messages[last_user_index]) + + return result_messages + def _format_bedrock_messages(self, messages: Messages) -> list[dict[str, Any]]: """Format messages for Bedrock API compatibility. diff --git a/tests/strands/models/test_bedrock.py b/tests/strands/models/test_bedrock.py index 33be44b1b..279f712fe 100644 --- a/tests/strands/models/test_bedrock.py +++ b/tests/strands/models/test_bedrock.py @@ -2052,6 +2052,50 @@ def test_format_request_filters_output_schema(model, messages, model_id): assert tool_spec["inputSchema"] == {"type": "object", "properties": {}} +def test_get_last_turn_messages(model): + """Test _get_last_turn_messages helper method.""" + # Test empty messages + assert model._get_last_turn_messages([]) == [] + + # Test single user message + messages = [{"role": "user", "content": [{"text": "Hello"}]}] + result = model._get_last_turn_messages(messages) + assert len(result) == 1 + assert result[0]["role"] == "user" + + # Test user-assistant pair + messages = [ + {"role": "user", "content": [{"text": "Hello"}]}, + {"role": "assistant", "content": [{"text": "Hi"}]}, + {"role": "user", "content": [{"text": "How are you?"}]}, + ] + result = model._get_last_turn_messages(messages) + assert len(result) == 2 + assert result[0]["role"] == "assistant" + assert result[1]["role"] == "user" + assert result[1]["content"][0]["text"] == "How are you?" + + +def test_format_request_with_guardrail_last_turn_only(model, model_id): + """Test _format_request uses filtered messages when guardrail_last_turn_only=True.""" + model.update_config(guardrail_id="test-guardrail", guardrail_version="DRAFT", guardrail_last_turn_only=True) + + messages = [ + {"role": "user", "content": [{"text": "First message"}]}, + {"role": "assistant", "content": [{"text": "First response"}]}, + {"role": "user", "content": [{"text": "Latest message"}]}, + ] + + request = model._format_request(messages) + + # Should only include the last turn (assistant + user) + formatted_messages = request["messages"] + assert len(formatted_messages) == 2 + assert formatted_messages[0]["role"] == "assistant" + assert formatted_messages[1]["role"] == "user" + assert formatted_messages[1]["content"][0]["text"] == "Latest message" + + @pytest.mark.asyncio async def test_stream_backward_compatibility_system_prompt(bedrock_client, model, messages, alist): """Test that system_prompt is converted to system_prompt_content when system_prompt_content is None.""" diff --git a/tests_integ/test_bedrock_guardrails.py b/tests_integ/test_bedrock_guardrails.py index 37fa6028c..06195e57a 100644 --- a/tests_integ/test_bedrock_guardrails.py +++ b/tests_integ/test_bedrock_guardrails.py @@ -289,6 +289,60 @@ def list_users() -> str: assert tool_result["content"][0]["text"] == INPUT_REDACT_MESSAGE +def test_guardrail_last_turn_only(boto_session, bedrock_guardrail): + """Test that guardrail_last_turn_only only sends the last turn to guardrails.""" + bedrock_model = BedrockModel( + guardrail_id=bedrock_guardrail, + guardrail_version="DRAFT", + guardrail_last_turn_only=True, + boto_session=boto_session, + ) + + agent = Agent(model=bedrock_model, system_prompt="You are a helpful assistant.", callback_handler=None) + + # First conversation turn - should not trigger guardrail + response1 = agent("Hello, how are you?") + assert response1.stop_reason != "guardrail_intervened" + + # Second conversation turn with blocked word - should trigger guardrail + # Since guardrail_last_turn_only=True, only this message and the previous assistant response + # should be evaluated by the guardrail, not the entire conversation history + response2 = agent("CACTUS") + assert response2.stop_reason == "guardrail_intervened" + assert str(response2).strip() == BLOCKED_INPUT + + +def test_guardrail_last_turn_only_recovery_scenario(boto_session, bedrock_guardrail): + """Test guardrail recovery: blocked content followed by normal question. + + This tests the key benefit of guardrail_last_turn_only: + 1. First turn: blocked content triggers guardrail + 2. Second turn: normal question should work because only last turn is analyzed + """ + bedrock_model = BedrockModel( + guardrail_id=bedrock_guardrail, + guardrail_version="DRAFT", + guardrail_last_turn_only=True, + boto_session=boto_session, + ) + + agent = Agent(model=bedrock_model, system_prompt="You are a helpful assistant.", callback_handler=None) + + # First turn - should be blocked by guardrail + response1 = agent("CACTUS") + assert response1.stop_reason == "guardrail_intervened" + assert str(response1).strip() == BLOCKED_INPUT + + # Second turn - should work normally with last turn only + # This is the key test: normal questions should work after blocked content + response2 = agent("What is the weather like today?") + assert response2.stop_reason != "guardrail_intervened" + assert str(response2).strip() != BLOCKED_INPUT + + # Verify the conversation has both messages + assert len(agent.messages) == 4 # 2 user + 2 assistant messages + + def test_guardrail_input_intervention_properly_redacts_in_session(boto_session, bedrock_guardrail, temp_dir): bedrock_model = BedrockModel( guardrail_id=bedrock_guardrail, From 2ffbfd7f6617845ca091fb48a5eb679cc4f1830f Mon Sep 17 00:00:00 2001 From: Aleksei Iancheruk Date: Fri, 21 Nov 2025 15:42:02 +0100 Subject: [PATCH 2/4] fix(bedrock): include assistant response in guardrail_last_turn_only context --- src/strands/models/bedrock.py | 15 ++++---- tests/strands/models/test_bedrock.py | 51 +++++++++++++++++++++----- tests_integ/test_bedrock_guardrails.py | 33 +++++++++++++++++ 3 files changed, 81 insertions(+), 18 deletions(-) diff --git a/src/strands/models/bedrock.py b/src/strands/models/bedrock.py index 3225254f1..5de194d4d 100644 --- a/src/strands/models/bedrock.py +++ b/src/strands/models/bedrock.py @@ -311,14 +311,14 @@ def _format_request( def _get_last_turn_messages(self, messages: Messages) -> Messages: """Get the last turn messages for guardrail evaluation. - Returns the latest user message and the previous assistant message (if it exists). + Returns the latest user message and the assistant's response (if it exists). This reduces the conversation context sent to guardrails when guardrail_last_turn_only is True. Args: messages: Full conversation messages. Returns: - Messages containing only the last turn (user + previous assistant if exists). + Messages containing only the last turn (user + assistant response if exists). """ if not messages: return [] @@ -334,13 +334,12 @@ def _get_last_turn_messages(self, messages: Messages) -> Messages: # No user message found, return empty return [] - # Include the previous assistant message if it exists - result_messages: Messages = [] - if last_user_index > 0 and messages[last_user_index - 1]["role"] == "assistant": - result_messages.append(messages[last_user_index - 1]) + # Start with the last user message + result_messages: Messages = [messages[last_user_index]] - # Add the last user message - result_messages.append(messages[last_user_index]) + # Include the assistant's response if it exists (the message after the user message) + if last_user_index < len(messages) - 1 and messages[last_user_index + 1]["role"] == "assistant": + result_messages.append(messages[last_user_index + 1]) return result_messages diff --git a/tests/strands/models/test_bedrock.py b/tests/strands/models/test_bedrock.py index 279f712fe..23004db5d 100644 --- a/tests/strands/models/test_bedrock.py +++ b/tests/strands/models/test_bedrock.py @@ -2057,29 +2057,43 @@ def test_get_last_turn_messages(model): # Test empty messages assert model._get_last_turn_messages([]) == [] - # Test single user message + # Test single user message (no assistant response yet) messages = [{"role": "user", "content": [{"text": "Hello"}]}] result = model._get_last_turn_messages(messages) assert len(result) == 1 assert result[0]["role"] == "user" - # Test user-assistant pair + # Test user-assistant pair (user message with assistant response) messages = [ {"role": "user", "content": [{"text": "Hello"}]}, {"role": "assistant", "content": [{"text": "Hi"}]}, {"role": "user", "content": [{"text": "How are you?"}]}, + {"role": "assistant", "content": [{"text": "I'm doing well"}]}, ] result = model._get_last_turn_messages(messages) assert len(result) == 2 - assert result[0]["role"] == "assistant" - assert result[1]["role"] == "user" - assert result[1]["content"][0]["text"] == "How are you?" + assert result[0]["role"] == "user" + assert result[0]["content"][0]["text"] == "How are you?" + assert result[1]["role"] == "assistant" + assert result[1]["content"][0]["text"] == "I'm doing well" + + # Test last user message without assistant response + messages = [ + {"role": "user", "content": [{"text": "Hello"}]}, + {"role": "assistant", "content": [{"text": "Hi"}]}, + {"role": "user", "content": [{"text": "How are you?"}]}, + ] + result = model._get_last_turn_messages(messages) + assert len(result) == 1 + assert result[0]["role"] == "user" + assert result[0]["content"][0]["text"] == "How are you?" def test_format_request_with_guardrail_last_turn_only(model, model_id): """Test _format_request uses filtered messages when guardrail_last_turn_only=True.""" model.update_config(guardrail_id="test-guardrail", guardrail_version="DRAFT", guardrail_last_turn_only=True) + # Test with last user message only (no assistant response yet) messages = [ {"role": "user", "content": [{"text": "First message"}]}, {"role": "assistant", "content": [{"text": "First response"}]}, @@ -2088,12 +2102,29 @@ def test_format_request_with_guardrail_last_turn_only(model, model_id): request = model._format_request(messages) - # Should only include the last turn (assistant + user) + # Should only include the last user message (no assistant response after it yet) formatted_messages = request["messages"] - assert len(formatted_messages) == 2 - assert formatted_messages[0]["role"] == "assistant" - assert formatted_messages[1]["role"] == "user" - assert formatted_messages[1]["content"][0]["text"] == "Latest message" + assert len(formatted_messages) == 1 + assert formatted_messages[0]["role"] == "user" + assert formatted_messages[0]["content"][0]["text"] == "Latest message" + + # Test with last user message + assistant response + messages_with_response = [ + {"role": "user", "content": [{"text": "First message"}]}, + {"role": "assistant", "content": [{"text": "First response"}]}, + {"role": "user", "content": [{"text": "How are you?"}]}, + {"role": "assistant", "content": [{"text": "I'm good"}]}, + ] + + request2 = model._format_request(messages_with_response) + + # Should include last user + assistant response + formatted_messages2 = request2["messages"] + assert len(formatted_messages2) == 2 + assert formatted_messages2[0]["role"] == "user" + assert formatted_messages2[0]["content"][0]["text"] == "How are you?" + assert formatted_messages2[1]["role"] == "assistant" + assert formatted_messages2[1]["content"][0]["text"] == "I'm good" @pytest.mark.asyncio diff --git a/tests_integ/test_bedrock_guardrails.py b/tests_integ/test_bedrock_guardrails.py index 06195e57a..f65587342 100644 --- a/tests_integ/test_bedrock_guardrails.py +++ b/tests_integ/test_bedrock_guardrails.py @@ -343,6 +343,39 @@ def test_guardrail_last_turn_only_recovery_scenario(boto_session, bedrock_guardr assert len(agent.messages) == 4 # 2 user + 2 assistant messages +def test_guardrail_last_turn_only_output_intervention(boto_session, bedrock_guardrail): + """Test that guardrail_last_turn_only works with OUTPUT guardrails. + + This tests that when the assistant tries to output blocked content, + the OUTPUT guardrail intervenes, even with guardrail_last_turn_only=True. + Then verifies that subsequent normal responses work correctly. + """ + bedrock_model = BedrockModel( + guardrail_id=bedrock_guardrail, + guardrail_version="DRAFT", + guardrail_last_turn_only=True, + guardrail_stream_processing_mode="sync", + boto_session=boto_session, + ) + + agent = Agent( + model=bedrock_model, + system_prompt="When asked to say the word, say CACTUS. Otherwise respond normally.", + callback_handler=None, + load_tools_from_directory=False, + ) + + # First turn - assistant tries to output "CACTUS", should be blocked by OUTPUT guardrail + response1 = agent("Say the word.") + assert response1.stop_reason == "guardrail_intervened" + assert BLOCKED_OUTPUT in str(response1) + + # Second turn - normal question should work fine + response2 = agent("What is 2+2?") + assert response2.stop_reason != "guardrail_intervened" + assert BLOCKED_OUTPUT not in str(response2) + + def test_guardrail_input_intervention_properly_redacts_in_session(boto_session, bedrock_guardrail, temp_dir): bedrock_model = BedrockModel( guardrail_id=bedrock_guardrail, From 8cda99fda0596fa8076d276200abc6a853436eac Mon Sep 17 00:00:00 2001 From: Jack Yuan Date: Fri, 19 Dec 2025 11:44:25 -0800 Subject: [PATCH 3/4] fix: optimize code --- src/strands/models/bedrock.py | 25 +++++-------------------- 1 file changed, 5 insertions(+), 20 deletions(-) diff --git a/src/strands/models/bedrock.py b/src/strands/models/bedrock.py index 5de194d4d..f2003ffbc 100644 --- a/src/strands/models/bedrock.py +++ b/src/strands/models/bedrock.py @@ -320,28 +320,13 @@ def _get_last_turn_messages(self, messages: Messages) -> Messages: Returns: Messages containing only the last turn (user + assistant response if exists). """ - if not messages: - return [] - - # Find the last user message - last_user_index = -1 for i in range(len(messages) - 1, -1, -1): if messages[i]["role"] == "user": - last_user_index = i - break - - if last_user_index == -1: - # No user message found, return empty - return [] - - # Start with the last user message - result_messages: Messages = [messages[last_user_index]] - - # Include the assistant's response if it exists (the message after the user message) - if last_user_index < len(messages) - 1 and messages[last_user_index + 1]["role"] == "assistant": - result_messages.append(messages[last_user_index + 1]) - - return result_messages + # Include assistant response if it immediately follows + if i + 1 < len(messages) and messages[i + 1]["role"] == "assistant": + return [messages[i], messages[i + 1]] + return [messages[i]] + return [] def _format_bedrock_messages(self, messages: Messages) -> list[dict[str, Any]]: """Format messages for Bedrock API compatibility. From 398716575623d03b7641c93a4e7bf6f2c25a6878 Mon Sep 17 00:00:00 2001 From: Jack Yuan Date: Mon, 22 Dec 2025 14:41:59 -0800 Subject: [PATCH 4/4] feat: rewrtie the logic, include last user message in guardContent when feature flag is true --- src/strands/models/bedrock.py | 53 ++++----- tests/strands/models/test_bedrock.py | 154 ++++++++++++++------------- 2 files changed, 102 insertions(+), 105 deletions(-) diff --git a/src/strands/models/bedrock.py b/src/strands/models/bedrock.py index f2003ffbc..a458a48fd 100644 --- a/src/strands/models/bedrock.py +++ b/src/strands/models/bedrock.py @@ -209,14 +209,7 @@ def _format_request( Returns: A Bedrock converse stream request. """ - # Filter messages for guardrails if guardrail_last_turn_only is enabled messages_for_request = messages - if ( - self.config.get("guardrail_last_turn_only", False) - and self.config.get("guardrail_id") - and self.config.get("guardrail_version") - ): - messages_for_request = self._get_last_turn_messages(messages) if not tool_specs: has_tool_content = any( @@ -237,7 +230,10 @@ def _format_request( return { "modelId": self.config["model_id"], - "messages": self._format_bedrock_messages(messages_for_request), + "messages": self._format_bedrock_messages( + messages_for_request, + guardrail_last_turn_only=bool(self.config.get("guardrail_last_turn_only", False)), + ), "system": system_blocks, **( { @@ -308,36 +304,20 @@ def _format_request( ), } - def _get_last_turn_messages(self, messages: Messages) -> Messages: - """Get the last turn messages for guardrail evaluation. - - Returns the latest user message and the assistant's response (if it exists). - This reduces the conversation context sent to guardrails when guardrail_last_turn_only is True. - - Args: - messages: Full conversation messages. - - Returns: - Messages containing only the last turn (user + assistant response if exists). - """ - for i in range(len(messages) - 1, -1, -1): - if messages[i]["role"] == "user": - # Include assistant response if it immediately follows - if i + 1 < len(messages) and messages[i + 1]["role"] == "assistant": - return [messages[i], messages[i + 1]] - return [messages[i]] - return [] - - def _format_bedrock_messages(self, messages: Messages) -> list[dict[str, Any]]: + def _format_bedrock_messages( + self, messages: Messages, guardrail_last_turn_only: bool = False + ) -> list[dict[str, Any]]: """Format messages for Bedrock API compatibility. This function ensures messages conform to Bedrock's expected format by: - Filtering out SDK_UNKNOWN_MEMBER content blocks - Eagerly filtering content blocks to only include Bedrock-supported fields - Ensuring all message content blocks are properly formatted for the Bedrock API + - Optionally wrapping the last user message in guardrailConverseContent blocks Args: messages: List of messages to format + guardrail_last_turn_only: If True, wrap the last user message content in guardrailConverseContent blocks Returns: Messages formatted for Bedrock API compatibility @@ -354,7 +334,15 @@ def _format_bedrock_messages(self, messages: Messages) -> list[dict[str, Any]]: filtered_unknown_members = False dropped_deepseek_reasoning_content = False - for message in messages: + # Find the index of the last user message if wrapping is enabled + last_user_idx = -1 + if guardrail_last_turn_only: + for i in range(len(messages) - 1, -1, -1): + if messages[i]["role"] == "user": + last_user_idx = i + break + + for idx, message in enumerate(messages): cleaned_content: list[dict[str, Any]] = [] for content_block in message["content"]: @@ -371,6 +359,11 @@ def _format_bedrock_messages(self, messages: Messages) -> list[dict[str, Any]]: # Format content blocks for Bedrock API compatibility formatted_content = self._format_request_message_content(content_block) + + # Wrap text content in guardrailConverseContent if this is the last user message + if guardrail_last_turn_only and idx == last_user_idx and "text" in formatted_content: + formatted_content = {"guardContent": {"text": {"text": formatted_content["text"]}}} + cleaned_content.append(formatted_content) # Create new message with cleaned content (skip if empty) diff --git a/tests/strands/models/test_bedrock.py b/tests/strands/models/test_bedrock.py index 23004db5d..d648a5918 100644 --- a/tests/strands/models/test_bedrock.py +++ b/tests/strands/models/test_bedrock.py @@ -2052,81 +2052,6 @@ def test_format_request_filters_output_schema(model, messages, model_id): assert tool_spec["inputSchema"] == {"type": "object", "properties": {}} -def test_get_last_turn_messages(model): - """Test _get_last_turn_messages helper method.""" - # Test empty messages - assert model._get_last_turn_messages([]) == [] - - # Test single user message (no assistant response yet) - messages = [{"role": "user", "content": [{"text": "Hello"}]}] - result = model._get_last_turn_messages(messages) - assert len(result) == 1 - assert result[0]["role"] == "user" - - # Test user-assistant pair (user message with assistant response) - messages = [ - {"role": "user", "content": [{"text": "Hello"}]}, - {"role": "assistant", "content": [{"text": "Hi"}]}, - {"role": "user", "content": [{"text": "How are you?"}]}, - {"role": "assistant", "content": [{"text": "I'm doing well"}]}, - ] - result = model._get_last_turn_messages(messages) - assert len(result) == 2 - assert result[0]["role"] == "user" - assert result[0]["content"][0]["text"] == "How are you?" - assert result[1]["role"] == "assistant" - assert result[1]["content"][0]["text"] == "I'm doing well" - - # Test last user message without assistant response - messages = [ - {"role": "user", "content": [{"text": "Hello"}]}, - {"role": "assistant", "content": [{"text": "Hi"}]}, - {"role": "user", "content": [{"text": "How are you?"}]}, - ] - result = model._get_last_turn_messages(messages) - assert len(result) == 1 - assert result[0]["role"] == "user" - assert result[0]["content"][0]["text"] == "How are you?" - - -def test_format_request_with_guardrail_last_turn_only(model, model_id): - """Test _format_request uses filtered messages when guardrail_last_turn_only=True.""" - model.update_config(guardrail_id="test-guardrail", guardrail_version="DRAFT", guardrail_last_turn_only=True) - - # Test with last user message only (no assistant response yet) - messages = [ - {"role": "user", "content": [{"text": "First message"}]}, - {"role": "assistant", "content": [{"text": "First response"}]}, - {"role": "user", "content": [{"text": "Latest message"}]}, - ] - - request = model._format_request(messages) - - # Should only include the last user message (no assistant response after it yet) - formatted_messages = request["messages"] - assert len(formatted_messages) == 1 - assert formatted_messages[0]["role"] == "user" - assert formatted_messages[0]["content"][0]["text"] == "Latest message" - - # Test with last user message + assistant response - messages_with_response = [ - {"role": "user", "content": [{"text": "First message"}]}, - {"role": "assistant", "content": [{"text": "First response"}]}, - {"role": "user", "content": [{"text": "How are you?"}]}, - {"role": "assistant", "content": [{"text": "I'm good"}]}, - ] - - request2 = model._format_request(messages_with_response) - - # Should include last user + assistant response - formatted_messages2 = request2["messages"] - assert len(formatted_messages2) == 2 - assert formatted_messages2[0]["role"] == "user" - assert formatted_messages2[0]["content"][0]["text"] == "How are you?" - assert formatted_messages2[1]["role"] == "assistant" - assert formatted_messages2[1]["content"][0]["text"] == "I'm good" - - @pytest.mark.asyncio async def test_stream_backward_compatibility_system_prompt(bedrock_client, model, messages, alist): """Test that system_prompt is converted to system_prompt_content when system_prompt_content is None.""" @@ -2271,3 +2196,82 @@ async def test_citations_content_preserves_tagged_union_structure(bedrock_client "(documentChar, documentPage, documentChunk, searchResultLocation, or web) " "with the location fields nested inside." ) + + +@pytest.mark.asyncio +async def test_format_request_with_guardrail_last_turn_only(model): + """Test _format_request passes apply_last_turn flag correctly.""" + model.update_config(guardrail_id="test-guardrail", guardrail_version="DRAFT", guardrail_last_turn_only=True) + + messages = [ + {"role": "user", "content": [{"text": "First message"}]}, + {"role": "assistant", "content": [{"text": "First response"}]}, + {"role": "user", "content": [{"text": "Latest message"}]}, + ] + + request = model._format_request(messages) + + # All messages should be in the request + formatted_messages = request["messages"] + assert len(formatted_messages) == 3 + + # Last user message should be wrapped + assert "guardContent" in formatted_messages[2]["content"][0] + assert formatted_messages[2]["content"][0]["guardContent"]["text"]["text"] == "Latest message" + + # First user message should NOT be wrapped + assert "text" in formatted_messages[0]["content"][0] + assert formatted_messages[0]["content"][0]["text"] == "First message" + + +def test_format_bedrock_messages_multimodal_content(model): + """Test that only text blocks are wrapped, not images.""" + messages = [ + { + "role": "user", + "content": [ + {"text": "Look at this image"}, + {"image": {"format": "png", "source": {"bytes": b"fake_image_data"}}}, + ], + } + ] + + result = model._format_bedrock_messages(messages, guardrail_last_turn_only=True) + + # Should have 2 content blocks + assert len(result[0]["content"]) == 2 + + # Text should be wrapped + assert "guardContent" in result[0]["content"][0] + assert result[0]["content"][0]["guardContent"]["text"]["text"] == "Look at this image" + + # Image should NOT be wrapped + assert "image" in result[0]["content"][1] + + +def test_format_bedrock_messages_wraps_last_user_text(model): + """Test that only the last user message text is wrapped in guardContent.""" + messages = [ + {"role": "user", "content": [{"text": "First message"}]}, + {"role": "assistant", "content": [{"text": "First response"}]}, + {"role": "user", "content": [{"text": "Latest message"}]}, + ] + + result = model._format_bedrock_messages(messages, guardrail_last_turn_only=True) + + # All messages should be present + assert len(result) == 3 + + # First user message should NOT be wrapped + assert result[0]["role"] == "user" + assert "text" in result[0]["content"][0] + assert result[0]["content"][0]["text"] == "First message" + + # Assistant message should be unchanged + assert result[1]["role"] == "assistant" + assert result[1]["content"][0]["text"] == "First response" + + # Last user message should be wrapped in guardContent + assert result[2]["role"] == "user" + assert "guardContent" in result[2]["content"][0] + assert result[2]["content"][0]["guardContent"]["text"]["text"] == "Latest message"