Merge pull request #18115 from BerriAI/litellm_bedrock_image_gaurdrail_support1

Add support for bedrock image gaurdrails

litellm/llms/bedrock/image/image_handler.py

@@ -170,6 +170,21 @@ class BedrockImageGeneration(BaseAWSLLM):
         )
         return model_response
 
+    def _extract_headers_from_optional_params(self, optional_params: dict) -> dict:
+        """
+        Extract guardrail parameters from optional_params and convert them to headers.
+        """
+        headers = {}
+        guardrail_identifier = optional_params.pop("guardrailIdentifier", None)
+        guardrail_version = optional_params.pop("guardrailVersion", None)
+        
+        if guardrail_identifier is not None:
+            headers["x-amz-bedrock-guardrail-identifier"] = guardrail_identifier
+        if guardrail_version is not None:
+            headers["x-amz-bedrock-guardrail-version"] = guardrail_version
+            
+        return headers
+
     def _prepare_request(
         self,
         model: str,
@@ -228,6 +243,10 @@ class BedrockImageGeneration(BaseAWSLLM):
         if extra_headers is not None:
             headers = {"Content-Type": "application/json", **extra_headers}
 
+        # Extract guardrail parameters and add them as headers
+        guardrail_headers = self._extract_headers_from_optional_params(optional_params)
+        headers.update(guardrail_headers)
+
         prepped = self.get_request_headers(
             credentials=boto3_credentials_info.credentials,
             aws_region_name=boto3_credentials_info.aws_region_name,

tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py

@@ -541,3 +541,28 @@ def test_amazon_titan_image_gen():
     print(f"response cost: {response._hidden_params['response_cost']}")
 
     assert response._hidden_params["response_cost"] > 0
+
+
+def test_extract_headers_from_optional_params_with_guardrails():
+    """Test that guardrail parameters are correctly extracted from optional_params and converted to headers"""
+    handler = BedrockImageGeneration()
+    
+    # Test with both guardrail parameters
+    optional_params = {
+        "guardrailIdentifier": "4cf5knqaeq15",
+        "guardrailVersion": "1",
+        "someOtherParam": "value",
+    }
+    
+    headers = handler._extract_headers_from_optional_params(optional_params)
+    
+    # Verify headers are correctly set
+    assert headers["x-amz-bedrock-guardrail-identifier"] == "4cf5knqaeq15"
+    assert headers["x-amz-bedrock-guardrail-version"] == "1"
+    
+    # Verify guardrail params are removed from optional_params
+    assert "guardrailIdentifier" not in optional_params
+    assert "guardrailVersion" not in optional_params
+    
+    # Verify other params remain in optional_params
+    assert optional_params["someOtherParam"] == "value"