scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

fix(bedrock_guardrails.py): respect bedrock runtime endpoint when using guardrails

Browse source 
Closes LIT-983
This commit is contained in:
Krrish Dholakia 2025-09-17 14:41:40 -07:00
parent 06f3991ed0
commit 1ada663959
2 changed files with 90 additions and 46 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/proxy/_new_secret_config.yaml
View file

@ -18,3 +18,15 @@ model_list:
guardrails:
- guardrail_name: "intel-bedrock-guard-cfg"
litellm_params:
guardrail: bedrock
mode: [pre_call, post_call]
guardrailIdentifier: "1234"
guardrailVersion: "1"
aws_access_key_id: "os.environ/AWS_ACCESS_KEY_ID"
aws_secret_access_key: "os.environ/AWS_SECRET_ACCESS_KEY"
aws_bedrock_runtime_endpoint: "os.environ/AWS_BEDROCK_RUNTIME_ENDPOINT"
default_on: true

124
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -15,7 +15,7 @@ sys.path.insert(
import json
import sys
from typing import Any, AsyncGenerator, List, Literal, Optional, Tuple, Union
from litellm.secret_managers.main import get_secret_str
import httpx
from fastapi import HTTPException
@ -32,6 +32,7 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
@ -118,18 +119,17 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""
If True, will not raise an exception when the guardrail is blocked.
"""
# Set supported event hooks to include MCP hooks
if 'supported_event_hooks' not in kwargs:
kwargs['supported_event_hooks'] = [
if "supported_event_hooks" not in kwargs:
kwargs["supported_event_hooks"] = [
GuardrailEventHooks.pre_call,
GuardrailEventHooks.post_call,
GuardrailEventHooks.during_call,
GuardrailEventHooks.pre_mcp_call,
GuardrailEventHooks.during_mcp_call,
]
super().__init__(**kwargs)
BaseAWSLLM.__init__(self)
@ -138,9 +138,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.guardrailIdentifier,
self.guardrailVersion,
)
def _create_bedrock_input_content_request(self, messages: Optional[List[AllMessageValues]]) -> BedrockRequest:
def _create_bedrock_input_content_request(
self, messages: Optional[List[AllMessageValues]]
) -> BedrockRequest:
"""
Create a bedrock request for the input content - the LLM request.
"""
@ -149,8 +150,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
if messages is None:
return bedrock_request
for message in messages:
message_text_content: Optional[List[str]] = (
self.get_content_for_message(message=message)
message_text_content: Optional[List[str]] = self.get_content_for_message(
message=message
)
if message_text_content is None:
continue
@ -163,7 +164,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
bedrock_request["content"] = bedrock_request_content
return bedrock_request
def _create_bedrock_output_content_request(self, response: Union[Any, ModelResponse]) -> BedrockRequest:
def _create_bedrock_output_content_request(
self, response: Union[Any, ModelResponse]
) -> BedrockRequest:
"""
Create a bedrock request for the output content - the LLM response.
"""
@ -199,9 +202,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""
bedrock_request: BedrockRequest = BedrockRequest(source=source)
if source == "INPUT":
bedrock_request = self._create_bedrock_input_content_request(messages=messages)
bedrock_request = self._create_bedrock_input_content_request(
messages=messages
)
elif source == "OUTPUT":
bedrock_request = self._create_bedrock_output_content_request(response=response)
bedrock_request = self._create_bedrock_output_content_request(
response=response
)
return bedrock_request
#### CALL HOOKS - proxy only ####
@ -255,9 +262,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
headers = {"Content-Type": "application/json"}
if extra_headers is not None:
headers = {"Content-Type": "application/json", **extra_headers}
api_base = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
aws_bedrock_runtime_endpoint = self.optional_params.get(
"aws_bedrock_runtime_endpoint", None
)
_, proxy_endpoint_url = self.get_runtime_endpoint(
api_base=None,
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
aws_region_name=aws_region_name,
)
proxy_endpoint_url = f"{proxy_endpoint_url}/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
# api_base = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
encoded_data = json.dumps(data).encode("utf-8")
# first check api-key, if none, fall back to sigV4
if api_key is not None:
aws_bearer_token: Optional[str] = api_key
@ -268,21 +285,31 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
try:
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
raise ImportError(
"Missing boto3 to call bedrock. Run 'pip install boto3'."
)
headers["Authorization"] = f"Bearer {aws_bearer_token}"
request = AWSRequest(
method="POST", url=api_base, data=encoded_data, headers=headers
method="POST",
url=proxy_endpoint_url,
data=encoded_data,
headers=headers,
)
else:
try:
from botocore.auth import SigV4Auth
from botocore.awsrequest import AWSRequest
except ImportError:
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
raise ImportError(
"Missing boto3 to call bedrock. Run 'pip install boto3'."
)
sigv4 = SigV4Auth(credentials, "bedrock", aws_region_name)
request = AWSRequest(
method="POST", url=api_base, data=encoded_data, headers=headers
method="POST",
url=proxy_endpoint_url,
data=encoded_data,
headers=headers,
)
sigv4.add_auth(request)
if (
@ -294,20 +321,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
return prepped_request
async def make_bedrock_api_request(
self,
self,
source: Literal["INPUT", "OUTPUT"],
messages: Optional[List[AllMessageValues]] = None,
response: Optional[Union[Any, litellm.ModelResponse]] = None,
request_data: Optional[dict] = None
request_data: Optional[dict] = None,
) -> BedrockGuardrailResponse:
from datetime import datetime
start_time = datetime.now()
credentials, aws_region_name = self._load_credentials()
bedrock_request_data: dict = dict(
self.convert_to_bedrock_format(
source=source,
messages=messages,
response=response
source=source, messages=messages, response=response
)
)
bedrock_guardrail_response: BedrockGuardrailResponse = (
@ -316,11 +342,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
api_key: Optional[str] = None
if request_data:
bedrock_request_data.update(
self.get_guardrail_dynamic_request_body_params(request_data=request_data)
self.get_guardrail_dynamic_request_body_params(
request_data=request_data
)
)
if request_data.get("api_key") is not None:
api_key = request_data["api_key"]
prepared_request = self._prepare_request(
credentials=credentials,
data=bedrock_request_data,
@ -346,7 +374,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=response.json(),
request_data=request_data or {},
guardrail_status=self._get_bedrock_guardrail_response_status(response=response),
guardrail_status=self._get_bedrock_guardrail_response_status(
response=response
),
start_time=start_time.timestamp(),
end_time=datetime.now().timestamp(),
duration=(datetime.now() - start_time).total_seconds(),
@ -372,8 +402,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
return bedrock_guardrail_response
def _get_bedrock_guardrail_response_status(self, response: httpx.Response) -> Literal["success", "failure"]:
def _get_bedrock_guardrail_response_status(
self, response: httpx.Response
) -> Literal["success", "failure"]:
"""
Get the status of the bedrock guardrail response.
"""
@ -381,7 +413,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
return "success"
return "failure"
def _get_http_exception_for_blocked_guardrail(self, response: BedrockGuardrailResponse) -> HTTPException:
def _get_http_exception_for_blocked_guardrail(
self, response: BedrockGuardrailResponse
) -> HTTPException:
"""
Get the HTTP exception for a blocked guardrail.
"""
@ -393,17 +427,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
for output in outputs:
if output.get("text"):
bedrock_guardrail_output_text += output.get("text") or ""
return HTTPException(
status_code=400,
detail={
"error": "Violated guardrail policy",
"error": "Violated guardrail policy",
"bedrock_guardrail_response": bedrock_guardrail_output_text,
}
},
)
def _should_raise_guardrail_blocked_exception(
self, response: BedrockGuardrailResponse
) -> bool:
@ -416,7 +448,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
# if user opted into masking, return False. since we'll use the masked output from the guardrail
if self.mask_request_content or self.mask_response_content:
return False
if self.disable_exception_on_block is True:
return False
@ -631,9 +663,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
########## 1. Make parallel Bedrock API requests ##########
#########################################################
output_content_bedrock = await self.make_bedrock_api_request(
source="OUTPUT",
response=response,
request_data=data
source="OUTPUT", response=response, request_data=data
) # Only response
#########################################################
@ -729,16 +759,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
###################################################################
# Create tasks for parallel execution
input_task = self.make_bedrock_api_request(
source="INPUT", messages=request_data.get("messages"), request_data=request_data
source="INPUT",
messages=request_data.get("messages"),
request_data=request_data,
) # Only input messages
output_task = self.make_bedrock_api_request(
source="OUTPUT", response=assembled_model_response
) # Only response
# Execute both requests in parallel
_, output_guardrail_response = await asyncio.gather(
input_task, output_task
)
_, output_guardrail_response = await asyncio.gather(input_task, output_task)
#########################################################################
########## 2. Apply masking to response with output guardrail response ##########
@ -891,7 +921,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
) -> None:
"""
Apply masked content from bedrock guardrail to the response object.
Args:
response: The response object to modify
bedrock_guardrail_response: Response from Bedrock guardrail containing masked content
@ -902,7 +932,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
if not masked_texts:
verbose_proxy_logger.debug("No masked outputs found, skipping response masking")
verbose_proxy_logger.debug(
"No masked outputs found, skipping response masking"
)
return
verbose_proxy_logger.debug(
@ -922,13 +954,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
) -> None:
"""
Apply masked texts to a ModelResponse object.
Args:
response: The ModelResponse object to modify in-place
masked_texts: List of masked text strings from guardrail
"""
masking_index = 0
for choice in response.choices:
if isinstance(choice, Choices):
# For chat completions