diff --git a/litellm/llms/bedrock/image/image_handler.py b/litellm/llms/bedrock/image/image_handler.py index 89e37bbdd8d..2e76596eefe 100644 --- a/litellm/llms/bedrock/image/image_handler.py +++ b/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, diff --git a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py b/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py index 5526f22cd5e..73331547772 100644 --- a/tests/image_gen_tests/test_bedrock_image_gen_unit_tests.py +++ b/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"