From 1371abf8809f81a4a186630786b4f8b3a4e36f8b Mon Sep 17 00:00:00 2001 From: Sameerlite Date: Wed, 17 Sep 2025 21:08:19 +0530 Subject: [PATCH 1/2] add last message as default in gaurdrail --- .../bedrock/chat/converse_transformation.py | 71 +++++- .../chat/test_converse_transformation.py | 205 ++++++++++++++++++ 2 files changed, 275 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 88103a86cf..6f6e07dd51 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -102,6 +102,69 @@ class AmazonConverseConfig(BaseConfig): "performanceConfig": PerformanceConfigBlock, } + @staticmethod + def _convert_last_user_message_to_guarded_text( + messages: List[AllMessageValues], optional_params: dict + ) -> List[AllMessageValues]: + """ + Convert the last user message to guarded_text type if guardrailConfig is present + and no guarded_text is already present in the last user message. + """ + # Check if guardrailConfig is present + if "guardrailConfig" not in optional_params: + return messages + + # Find the last user message + last_user_message = None + last_user_message_index = -1 + for i in range(len(messages) - 1, -1, -1): + if messages[i].get("role") == "user": + last_user_message = messages[i] + last_user_message_index = i + break + + if last_user_message is None: + return messages + + # Check if the last user message already has guarded_text + content = last_user_message.get("content", []) + if isinstance(content, list): + has_guarded_text = any( + isinstance(item, dict) and item.get("type") == "guarded_text" + for item in content + ) + if has_guarded_text: + return messages + + # Convert text elements to guarded_text + new_content = [] + for item in content: + if isinstance(item, dict) and item.get("type") == "text": + new_item = { + "type": "guarded_text", + "text": item["text"] + } + new_content.append(new_item) + else: + new_content.append(item) + + # Create a copy of messages and update the last user message + messages_copy = copy.deepcopy(messages) + messages_copy[last_user_message_index]["content"] = new_content + return messages_copy + elif isinstance(content, str): + # If content is a string, convert it to guarded_text + messages_copy = copy.deepcopy(messages) + messages_copy[last_user_message_index]["content"] = [ + { + "type": "guarded_text", + "text": content + } + ] + return messages_copy + + return messages + @classmethod def get_config(cls): return { @@ -769,6 +832,9 @@ class AmazonConverseConfig(BaseConfig): headers: Optional[dict] = None, ) -> RequestObject: messages, system_content_blocks = self._transform_system_message(messages) + + # Convert last user message to guarded_text if guardrailConfig is present + messages = self._convert_last_user_message_to_guarded_text(messages, optional_params) ## TRANSFORMATION ## _data: CommonRequestObject = self._transform_request_helper( @@ -821,6 +887,9 @@ class AmazonConverseConfig(BaseConfig): ) -> RequestObject: messages, system_content_blocks = self._transform_system_message(messages) + # Convert last user message to guarded_text if guardrailConfig is present + messages = self._convert_last_user_message_to_guarded_text(messages, optional_params) + _data: CommonRequestObject = self._transform_request_helper( model=model, system_content_blocks=system_content_blocks, @@ -1277,4 +1346,4 @@ class AmazonConverseConfig(BaseConfig): ################################################################### if "ai21" in model: return True - return False + return False \ No newline at end of file diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index df003850f7..78b07fb088 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -1868,3 +1868,208 @@ def test_guarded_text_guardrail_config_preserved(): assert result["inferenceConfig"]["guardrailConfig"]["guardrailIdentifier"] == "gr-abc123" +def test_auto_convert_last_user_message_to_guarded_text(): + """Test that last user message is automatically converted to guarded_text when guardrailConfig is present.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the main topic of this legal document?" + } + ] + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the helper method directly + converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + + # Verify the conversion + assert len(converted_messages) == 1 + assert converted_messages[0]["role"] == "user" + assert len(converted_messages[0]["content"]) == 1 + assert converted_messages[0]["content"][0]["type"] == "guarded_text" + assert converted_messages[0]["content"][0]["text"] == "What is the main topic of this legal document?" + + +def test_auto_convert_last_user_message_string_content(): + """Test that last user message with string content is automatically converted to guarded_text when guardrailConfig is present.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": "What is the main topic of this legal document?" + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the helper method directly + converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + + # Verify the conversion + assert len(converted_messages) == 1 + assert converted_messages[0]["role"] == "user" + assert len(converted_messages[0]["content"]) == 1 + assert converted_messages[0]["content"][0]["type"] == "guarded_text" + assert converted_messages[0]["content"][0]["text"] == "What is the main topic of this legal document?" + + +def test_no_conversion_when_no_guardrail_config(): + """Test that no conversion happens when guardrailConfig is not present.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the main topic of this legal document?" + } + ] + } + ] + + optional_params = {} + + # Test the helper method directly + converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + + # Verify no conversion happened + assert converted_messages == messages + + +def test_no_conversion_when_guarded_text_already_present(): + """Test that no conversion happens when guarded_text is already present in the last user message.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "guarded_text", + "text": "This is already guarded" + } + ] + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the helper method directly + converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + + # Verify no conversion happened + assert converted_messages == messages + + +def test_auto_convert_with_mixed_content(): + """Test that only text elements are converted to guarded_text, other content types are preserved.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the main topic of this legal document?" + }, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.jpg"} + } + ] + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the helper method directly + converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + + # Verify the conversion + assert len(converted_messages) == 1 + assert converted_messages[0]["role"] == "user" + assert len(converted_messages[0]["content"]) == 2 + + # First element should be converted to guarded_text + assert converted_messages[0]["content"][0]["type"] == "guarded_text" + assert converted_messages[0]["content"][0]["text"] == "What is the main topic of this legal document?" + + # Second element should remain unchanged + assert converted_messages[0]["content"][1]["type"] == "image_url" + assert converted_messages[0]["content"][1]["image_url"]["url"] == "https://example.com/image.jpg" + + +def test_auto_convert_in_full_transformation(): + """Test that the automatic conversion works in the full transformation pipeline.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What is the main topic of this legal document?" + } + ] + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the full transformation + result = config._transform_request( + model="anthropic.claude-3-sonnet-20240229-v1:0", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={} + ) + + # Verify the transformation worked + assert "messages" in result + assert len(result["messages"]) == 1 + + # The message should have guardrailConverseContent + message = result["messages"][0] + assert "content" in message + assert len(message["content"]) == 1 + assert "guardrailConverseContent" in message["content"][0] + assert message["content"][0]["guardrailConverseContent"]["text"] == "What is the main topic of this legal document?" + From edf95966c9419d3ae5d252e0dd6333399938ac7f Mon Sep 17 00:00:00 2001 From: Sameerlite Date: Thu, 18 Sep 2025 10:35:25 +0530 Subject: [PATCH 2/2] Handle consecutive user messages --- .../bedrock/chat/converse_transformation.py | 94 ++++--- .../chat/test_converse_transformation.py | 234 +++++++++++++++++- 2 files changed, 274 insertions(+), 54 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 6f6e07dd51..d99cff6b6b 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -103,67 +103,59 @@ class AmazonConverseConfig(BaseConfig): } @staticmethod - def _convert_last_user_message_to_guarded_text( + def _convert_consecutive_user_messages_to_guarded_text( messages: List[AllMessageValues], optional_params: dict ) -> List[AllMessageValues]: """ - Convert the last user message to guarded_text type if guardrailConfig is present - and no guarded_text is already present in the last user message. + Convert consecutive user messages at the end to guarded_text type if guardrailConfig is present + and no guarded_text is already present in those messages. """ # Check if guardrailConfig is present if "guardrailConfig" not in optional_params: return messages - # Find the last user message - last_user_message = None - last_user_message_index = -1 + # Find all consecutive user messages at the end + consecutive_user_message_indices = [] for i in range(len(messages) - 1, -1, -1): if messages[i].get("role") == "user": - last_user_message = messages[i] - last_user_message_index = i + consecutive_user_message_indices.append(i) + else: break - if last_user_message is None: + if not consecutive_user_message_indices: return messages - # Check if the last user message already has guarded_text - content = last_user_message.get("content", []) - if isinstance(content, list): - has_guarded_text = any( - isinstance(item, dict) and item.get("type") == "guarded_text" - for item in content - ) - if has_guarded_text: - return messages + # Process each consecutive user message + messages_copy = copy.deepcopy(messages) + for user_message_index in consecutive_user_message_indices: + user_message = messages_copy[user_message_index] + content = user_message.get("content", []) - # Convert text elements to guarded_text - new_content = [] - for item in content: - if isinstance(item, dict) and item.get("type") == "text": - new_item = { - "type": "guarded_text", - "text": item["text"] - } - new_content.append(new_item) - else: - new_content.append(item) - - # Create a copy of messages and update the last user message - messages_copy = copy.deepcopy(messages) - messages_copy[last_user_message_index]["content"] = new_content - return messages_copy - elif isinstance(content, str): - # If content is a string, convert it to guarded_text - messages_copy = copy.deepcopy(messages) - messages_copy[last_user_message_index]["content"] = [ - { - "type": "guarded_text", - "text": content - } - ] - return messages_copy + if isinstance(content, list): + has_guarded_text = any( + isinstance(item, dict) and item.get("type") == "guarded_text" + for item in content + ) + if has_guarded_text: + continue # Skip this message if it already has guarded_text - return messages + # Convert text elements to guarded_text + new_content = [] + for item in content: + if isinstance(item, dict) and item.get("type") == "text": + new_item = {"type": "guarded_text", "text": item["text"]} # type: ignore + new_content.append(new_item) + else: + new_content.append(item) + + messages_copy[user_message_index]["content"] = new_content # type: ignore + elif isinstance(content, str): + # If content is a string, convert it to guarded_text + messages_copy[user_message_index]["content"] = [ # type: ignore + {"type": "guarded_text", "text": content} # type: ignore + ] + + return messages_copy @classmethod def get_config(cls): @@ -832,9 +824,11 @@ class AmazonConverseConfig(BaseConfig): headers: Optional[dict] = None, ) -> RequestObject: messages, system_content_blocks = self._transform_system_message(messages) - + # Convert last user message to guarded_text if guardrailConfig is present - messages = self._convert_last_user_message_to_guarded_text(messages, optional_params) + messages = self._convert_consecutive_user_messages_to_guarded_text( + messages, optional_params + ) ## TRANSFORMATION ## _data: CommonRequestObject = self._transform_request_helper( @@ -888,7 +882,9 @@ class AmazonConverseConfig(BaseConfig): messages, system_content_blocks = self._transform_system_message(messages) # Convert last user message to guarded_text if guardrailConfig is present - messages = self._convert_last_user_message_to_guarded_text(messages, optional_params) + messages = self._convert_consecutive_user_messages_to_guarded_text( + messages, optional_params + ) _data: CommonRequestObject = self._transform_request_helper( model=model, @@ -1346,4 +1342,4 @@ class AmazonConverseConfig(BaseConfig): ################################################################### if "ai21" in model: return True - return False \ No newline at end of file + return False diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 78b07fb088..257f5be3ee 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -1892,7 +1892,7 @@ def test_auto_convert_last_user_message_to_guarded_text(): } # Test the helper method directly - converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) # Verify the conversion assert len(converted_messages) == 1 @@ -1921,7 +1921,7 @@ def test_auto_convert_last_user_message_string_content(): } # Test the helper method directly - converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) # Verify the conversion assert len(converted_messages) == 1 @@ -1950,7 +1950,7 @@ def test_no_conversion_when_no_guardrail_config(): optional_params = {} # Test the helper method directly - converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) # Verify no conversion happened assert converted_messages == messages @@ -1980,7 +1980,7 @@ def test_no_conversion_when_guarded_text_already_present(): } # Test the helper method directly - converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) # Verify no conversion happened assert converted_messages == messages @@ -2014,7 +2014,7 @@ def test_auto_convert_with_mixed_content(): } # Test the helper method directly - converted_messages = config._convert_last_user_message_to_guarded_text(messages, optional_params) + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) # Verify the conversion assert len(converted_messages) == 1 @@ -2073,3 +2073,227 @@ def test_auto_convert_in_full_transformation(): assert "guardrailConverseContent" in message["content"][0] assert message["content"][0]["guardrailConverseContent"]["text"] == "What is the main topic of this legal document?" + +def test_convert_consecutive_user_messages_to_guarded_text(): + """Test that consecutive user messages at the end are converted to guarded_text.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "First user message" + } + ] + }, + { + "role": "assistant", + "content": "Assistant response" + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Second user message" + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Third user message" + } + ] + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the helper method directly + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) + + # Verify the conversion - only the last two user messages should be converted + assert len(converted_messages) == 4 + + # First user message should remain unchanged + assert converted_messages[0]["role"] == "user" + assert converted_messages[0]["content"][0]["type"] == "text" + assert converted_messages[0]["content"][0]["text"] == "First user message" + + # Assistant message should remain unchanged + assert converted_messages[1]["role"] == "assistant" + assert converted_messages[1]["content"] == "Assistant response" + + # Second user message should be converted to guarded_text + assert converted_messages[2]["role"] == "user" + assert converted_messages[2]["content"][0]["type"] == "guarded_text" + assert converted_messages[2]["content"][0]["text"] == "Second user message" + + # Third user message should be converted to guarded_text + assert converted_messages[3]["role"] == "user" + assert converted_messages[3]["content"][0]["type"] == "guarded_text" + assert converted_messages[3]["content"][0]["text"] == "Third user message" + + +def test_convert_all_user_messages_when_all_consecutive(): + """Test that all user messages are converted when they are all consecutive at the end.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "First user message" + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Second user message" + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Third user message" + } + ] + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the helper method directly + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) + + # Verify all three user messages are converted + assert len(converted_messages) == 3 + + for i in range(3): + assert converted_messages[i]["role"] == "user" + assert converted_messages[i]["content"][0]["type"] == "guarded_text" + + assert converted_messages[0]["content"][0]["text"] == "First user message" + assert converted_messages[1]["content"][0]["text"] == "Second user message" + assert converted_messages[2]["content"][0]["text"] == "Third user message" + + +def test_convert_consecutive_user_messages_with_string_content(): + """Test that consecutive user messages with string content are converted to guarded_text.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "assistant", + "content": "Assistant response" + }, + { + "role": "user", + "content": "First user message" + }, + { + "role": "user", + "content": "Second user message" + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the helper method directly + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) + + # Verify the conversion + assert len(converted_messages) == 3 + + # Assistant message should remain unchanged + assert converted_messages[0]["role"] == "assistant" + assert converted_messages[0]["content"] == "Assistant response" + + # Both user messages should be converted to guarded_text + assert converted_messages[1]["role"] == "user" + assert len(converted_messages[1]["content"]) == 1 + assert converted_messages[1]["content"][0]["type"] == "guarded_text" + assert converted_messages[1]["content"][0]["text"] == "First user message" + + assert converted_messages[2]["role"] == "user" + assert len(converted_messages[2]["content"]) == 1 + assert converted_messages[2]["content"][0]["type"] == "guarded_text" + assert converted_messages[2]["content"][0]["text"] == "Second user message" + + +def test_skip_consecutive_user_messages_with_existing_guarded_text(): + """Test that consecutive user messages with existing guarded_text are skipped.""" + config = AmazonConverseConfig() + + messages = [ + { + "role": "user", + "content": [ + { + "type": "guarded_text", + "text": "Already guarded" + } + ] + }, + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Should be converted" + } + ] + } + ] + + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "gr-abc123", + "guardrailVersion": "1" + } + } + + # Test the helper method directly + converted_messages = config._convert_consecutive_user_messages_to_guarded_text(messages, optional_params) + + # Verify the conversion + assert len(converted_messages) == 2 + + # First message should remain unchanged (already has guarded_text) + assert converted_messages[0]["role"] == "user" + assert converted_messages[0]["content"][0]["type"] == "guarded_text" + assert converted_messages[0]["content"][0]["text"] == "Already guarded" + + # Second message should be converted + assert converted_messages[1]["role"] == "user" + assert converted_messages[1]["content"][0]["type"] == "guarded_text" + assert converted_messages[1]["content"][0]["text"] == "Should be converted" +