From 1371abf8809f81a4a186630786b4f8b3a4e36f8b Mon Sep 17 00:00:00 2001 From: Sameerlite Date: Wed, 17 Sep 2025 21:08:19 +0530 Subject: [PATCH] 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?" +