Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
14 changes: 0 additions & 14 deletions litellm/exceptions.py
Original file line number Diff line number Diff line change
Expand Up @@ -1180,20 +1180,6 @@ def __init__(
super().__init__(message)


class GuardrailInterventionNormalStringError(
Exception
): # custom exception to raise when a guardrail intervenes, but we want to return a normal string to the user
def __init__(self, message: str):
self.message = message
super().__init__(self.message)

def __str__(self):
return self.message

def __repr__(self):
return self.__str__()


class SensitiveDataRouteException(Exception):
"""
Exception raised when a guardrail detects sensitive data and wants to reroute the request.
Expand Down
145 changes: 82 additions & 63 deletions litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
from litellm._logging import verbose_proxy_logger
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
from litellm.caching import DualCache
from litellm.exceptions import GuardrailInterventionNormalStringError
from litellm.exceptions import ModifyResponseException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.custom_httpx.http_handler import (
Expand Down Expand Up @@ -754,7 +754,9 @@ async def make_bedrock_api_request(
)
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response):
raise self._get_http_exception_for_blocked_guardrail(bedrock_guardrail_response)
raise self._get_http_exception_for_blocked_guardrail(
bedrock_guardrail_response, request_data=request_data
)
else:
status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
verbose_proxy_logger.error(
Expand Down Expand Up @@ -1027,8 +1029,8 @@ def _extract_blocked_assessments(self, response: BedrockGuardrailResponse) -> Li
return blocked

def _get_http_exception_for_blocked_guardrail(
self, response: BedrockGuardrailResponse
) -> Union[HTTPException, GuardrailInterventionNormalStringError]:
self, response: BedrockGuardrailResponse, request_data: Optional[dict] = None
) -> Union[HTTPException, ModifyResponseException]:
"""
Get the HTTP exception for a blocked guardrail.
"""
Expand All @@ -1040,7 +1042,13 @@ def _get_http_exception_for_blocked_guardrail(
bedrock_guardrail_output_text += output.get("text") or ""

if self.disable_exception_on_block is True:
return GuardrailInterventionNormalStringError(message=bedrock_guardrail_output_text)
_request_data = request_data or {}
return ModifyResponseException(
message=bedrock_guardrail_output_text,
model=_request_data.get("model", "bedrock-guardrail"),
request_data=_request_data,
guardrail_name=self.guardrail_name,
)

detail: Dict[str, Any] = {
"error": "Violated guardrail policy",
Expand Down Expand Up @@ -1134,18 +1142,6 @@ def _should_raise_guardrail_blocked_exception(self, response: BedrockGuardrailRe
# This means all actions were ANONYMIZED or NONE, so don't raise exception
return False

def create_guardrail_blocked_response(self, response: str) -> ModelResponse:
from litellm.types.utils import Choices, Message, ModelResponse

return ModelResponse(
choices=[
Choices(
message=Message(content=response),
)
],
model="bedrock-guardrail",
)

async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
Expand Down Expand Up @@ -1183,16 +1179,15 @@ async def async_pre_call_hook(
#########################################################
########## 1. Make the Bedrock API request ##########
#########################################################
bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = None
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=data,
logging_event_type=GuardrailEventHooks.pre_call,
)
except GuardrailInterventionNormalStringError as e:
bedrock_guardrail_response = e.message
# A block with disable_exception_on_block=True raises ModifyResponseException
# from make_bedrock_api_request; that propagates to the endpoint handler,
# which returns a 200 whose message is the guardrail's blockedInputMessaging.
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=data,
logging_event_type=GuardrailEventHooks.pre_call,
)
#########################################################

#########################################################
Expand All @@ -1207,8 +1202,6 @@ async def async_pre_call_hook(
updated_target_messages=updated_subset,
target_indices=filter_result.target_indices,
)
if isinstance(bedrock_guardrail_response, str):
data["mock_response"] = self.create_guardrail_blocked_response(response=bedrock_guardrail_response)

#########################################################
########## 3. Add the guardrail to the applied guardrails header ##########
Expand Down Expand Up @@ -1248,16 +1241,19 @@ async def async_moderation_hook(
#########################################################
########## 1. Make the Bedrock API request ##########
#########################################################
bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = None
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=data,
logging_event_type=GuardrailEventHooks.during_call,
)
except GuardrailInterventionNormalStringError as e:
bedrock_guardrail_response = e.message
# A block with disable_exception_on_block=True raises ModifyResponseException
# from make_bedrock_api_request. Because during_call runs in an asyncio.gather
# alongside the LLM call (common_request_processing.py), swallowing the
# exception here to set data["mock_response"] was ineffective: route_request
# unpacked kwargs before this hook ran, and the LLM task's response was taken
# unconditionally. Letting the exception propagate cancels the LLM task and
# the endpoint handler returns the block response.
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT",
messages=filtered_messages,
request_data=data,
logging_event_type=GuardrailEventHooks.during_call,
)
#########################################################

#########################################################
Expand All @@ -1272,8 +1268,6 @@ async def async_moderation_hook(
updated_target_messages=updated_subset,
target_indices=filter_result.target_indices,
)
if isinstance(bedrock_guardrail_response, str):
data["mock_response"] = self.create_guardrail_blocked_response(response=bedrock_guardrail_response)

#########################################################
########## 3. Add the guardrail to the applied guardrails header ##########
Expand Down Expand Up @@ -1323,7 +1317,11 @@ async def async_post_call_success_hook(
# users should configure if they want input validation. Running an
# extra INPUT scan here produced a duplicate post-call entry in the
# trace and made no semantic sense for a "post-call" event.
output_content_bedrock: Optional[Union[BedrockGuardrailResponse, str]] = None
# A block with disable_exception_on_block=True raises ModifyResponseException
# from make_bedrock_api_request; that propagates to the endpoint handler,
# which returns a 200 whose message is the guardrail's blockedInputMessaging.
# Attach the LLM response to original_response so the synthetic block reply
# reports the real token usage the upstream call consumed instead of zero.
try:
output_content_bedrock = await self.make_bedrock_api_request(
source="OUTPUT",
Expand All @@ -1332,15 +1330,15 @@ async def async_post_call_success_hook(
request_data=data,
logging_event_type=GuardrailEventHooks.post_call,
)
except GuardrailInterventionNormalStringError as e:
output_content_bedrock = e.message
except ModifyResponseException as e:
if e.original_response is None:
e.original_response = response
raise

#########################################################
########## 2. Apply masking to response with output guardrail response ##########
#########################################################
if isinstance(output_content_bedrock, str):
response = self.create_guardrail_blocked_response(response=output_content_bedrock)
elif output_content_bedrock is not None:
if output_content_bedrock is not None:
self._apply_masking_to_response(
response=response,
bedrock_guardrail_response=output_content_bedrock,
Expand All @@ -1357,7 +1355,7 @@ async def async_post_call_success_hook(
def _update_messages_with_updated_bedrock_guardrail_response(
self,
messages: List[AllMessageValues],
bedrock_guardrail_response: Union[BedrockGuardrailResponse, str],
bedrock_guardrail_response: BedrockGuardrailResponse,
) -> List[AllMessageValues]:
"""
Use the output from the bedrock guardrail to mask sensitive content in messages.
Expand All @@ -1369,8 +1367,6 @@ def _update_messages_with_updated_bedrock_guardrail_response(
Returns:
List of messages with content masked according to guardrail response
"""
if isinstance(bedrock_guardrail_response, str):
return messages
# Get masked texts from guardrail response
masked_texts = self._extract_masked_texts_from_response(bedrock_guardrail_response)

Expand Down Expand Up @@ -1422,7 +1418,14 @@ async def async_post_call_streaming_iterator_hook(
# pre_call / during_call. Bedrock will raise if the response
# violates the guardrail policy.
###################################################################
output_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = None
# A block with disable_exception_on_block=True raises ModifyResponseException
# from make_bedrock_api_request. Non-streaming paths let it propagate so
# the endpoint handler turns it into a 200. Streaming can't do that: the
# SSE response headers are already flushed, so a raise would be serialized
# as an error frame by async_streaming_data_generator. Instead, replace
# the assembled response with the synthetic block content in-place and
# yield it as a normal stream, matching the shape a non-streaming block
# produces.
try:
output_guardrail_response = await self.make_bedrock_api_request(
source="OUTPUT",
Expand All @@ -1431,15 +1434,31 @@ async def async_post_call_streaming_iterator_hook(
request_data=request_data,
logging_event_type=GuardrailEventHooks.post_call,
)
except GuardrailInterventionNormalStringError as e:
output_guardrail_response = e.message
except ModifyResponseException as e:
# Preserve upstream usage from the LLM call we already
# consumed. Non-streaming blocks carry it via
# ModifyResponseException.original_response +
# _blocked_response_usage; streaming has to do the copy
# itself since the exception can't escape this generator.
_original_usage = getattr(assembled_model_response, "usage", None)
assembled_model_response = ModelResponse(
choices=[
Choices(
index=0,
message=Message(role="assistant", content=e.message),
finish_reason="content_filter",
)
],
model=e.model,
)
if _original_usage is not None:
assembled_model_response.usage = _original_usage
output_guardrail_response = None
Comment thread
cursor[bot] marked this conversation as resolved.

#########################################################################
########## 2. Apply masking to response with output guardrail response ##########
#########################################################################
if isinstance(output_guardrail_response, str):
assembled_model_response = self.create_guardrail_blocked_response(response=output_guardrail_response)
elif output_guardrail_response is not None:
if output_guardrail_response is not None:
self._apply_masking_to_response(
response=assembled_model_response,
bedrock_guardrail_response=output_guardrail_response,
Expand Down Expand Up @@ -1732,13 +1751,13 @@ async def apply_guardrail(
inputs["texts"] = masked_texts
return inputs

except (HTTPException, GuardrailInterventionNormalStringError):
# Let guardrail blocking exceptions propagate as-is so the proxy
# can return the correct HTTP status (400) or handle the
# GuardrailInterventionNormalStringError for disable_exception_on_block mode.
# Without this, the generic except below wraps them into a plain
# Exception, losing the HTTP semantics and preventing the proxy
# from properly blocking the call.
except (HTTPException, ModifyResponseException):
# Let guardrail blocking exceptions propagate as-is so the proxy can
# return the correct HTTP status (400 for HTTPException, 200 with the
# block message for ModifyResponseException in disable_exception_on_block
# mode). Without this, the generic except below wraps them into a plain
# Exception, losing the semantics and preventing the proxy from
# properly blocking the call.
raise
except Exception as e:
verbose_proxy_logger.error("Bedrock Guardrail: Failed to apply guardrail: %s", str(e))
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -329,12 +329,13 @@ def test_bedrock_guardrail_filters_latest_user_message_when_enabled():
@pytest.mark.asyncio
async def test_bedrock_apply_guardrail_blocked_with_disable_exception_on_block():
"""
Regression test for issue #20045: when disable_exception_on_block=True,
make_bedrock_api_request raises GuardrailInterventionNormalStringError.
apply_guardrail must let it propagate as-is so the proxy can handle it
properly instead of wrapping it in a generic Exception.
Regression test for LIT-4186: when disable_exception_on_block=True, a
Bedrock block raises ModifyResponseException. apply_guardrail must let it
propagate as-is so the endpoint handler (proxy_server.py) can turn it into
a 200 response with the block message as content, instead of the exception
surfacing as a bare 500.
"""
from litellm.exceptions import GuardrailInterventionNormalStringError
from litellm.exceptions import ModifyResponseException

guardrail = BedrockGuardrail(
guardrail_name="test-bedrock-guard",
Expand All @@ -346,18 +347,21 @@ async def test_bedrock_apply_guardrail_blocked_with_disable_exception_on_block()
with patch.object(
guardrail, "make_bedrock_api_request", new_callable=AsyncMock
) as mock_api:
mock_api.side_effect = GuardrailInterventionNormalStringError(
message="Sorry, your question in its current format is unable to be answered."
mock_api.side_effect = ModifyResponseException(
message="Sorry, your question in its current format is unable to be answered.",
model="bedrock-guardrail",
request_data={},
guardrail_name="test-bedrock-guard",
)

with pytest.raises(GuardrailInterventionNormalStringError) as exc_info:
with pytest.raises(ModifyResponseException) as exc_info:
await guardrail.apply_guardrail(
inputs={"texts": ["harmful prompt content"]},
request_data={},
input_type="request",
)

assert "unable to be answered" in str(exc_info.value.message)
assert "unable to be answered" in exc_info.value.message


@pytest.mark.asyncio
Expand Down
Loading
Loading