scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

Add conversion of gaurdrail params into bedrock headers

Browse source
This commit is contained in:
Sameer Kankute 2025-12-17 14:04:00 +05:30
parent 1c5f16fb6e
commit a1536fdd87
1 changed files with 19 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm/llms/bedrock/image/image_handler.py
View file

@ -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,