mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(guardrails/bedrock): honor disable_exception_on_block by raising ModifyResponseException (#32289)
* fix(guardrails/bedrock): honor disable_exception_on_block by raising ModifyResponseException The Bedrock-specific GuardrailInterventionNormalStringError predates the unified guardrails refactor and no proxy code path handles it, so a block with the flag set surfaced as an uncaught Exception -> HTTP 500 in pre_call mode and was silently discarded in during_call mode (model call proceeded in the parallel asyncio.gather; the block hook's data["mock_response"] mutation happened after route_request had already unpacked kwargs). Convert the block to ModifyResponseException at the raise site inside make_bedrock_api_request. That exception is the industry-standard proxy contract already caught in proxy_server, anthropic_endpoints, response_api _endpoints, and pass_through_endpoints; it turns into a 200 response with finish_reason=content_filter and the block message as content, which is exactly what the flag was documented to yield. Post-call blocks attach the LLM response to original_response so the synthetic reply reports the upstream call's real token usage instead of zero. Deletes the now-orphaned GuardrailInterventionNormalStringError class and the dead create_guardrail_blocked_response / mock_response plumbing in the Bedrock hooks; updates the existing tests that had locked in the buggy contract. Resolves LIT-4186 * chore(guardrails/bedrock): drop dead str branch in _update_messages_with_updated_bedrock_guardrail_response Follow-up to the disable_exception_on_block fix. That method used to receive either a BedrockGuardrailResponse or a plain string (the block message, when the flag was set). Now that a block always raises ModifyResponseException before this method runs, the string branch is unreachable; tighten the type to BedrockGuardrailResponse and delete the guard. * fix(guardrails/bedrock): streaming post_call block yields synthetic stream instead of surfacing as SSE 500 Regression from the LIT-4186 refactor: pre-refactor, the streaming post_call iterator caught GuardrailInterventionNormalStringError locally and replaced the assembled response with a synthetic content-filter message, then re-emitted it as chunks via MockResponseIterator. After the refactor the exception was re-raised as ModifyResponseException, which async_streaming_data_generator serializes as a proxy 500 error frame because the SSE response headers are already flushed by the time the block fires. Non-streaming paths still let ModifyResponseException propagate to the endpoint handler (which converts it into a 200). Streaming can't do that, so keep the local synthesis: on the exception, rebind the assembled response to a ModelResponse whose single choice carries the block message as content and finish_reason=content_filter, and let the downstream MockResponseIterator emit it as chunks. Same shape a non-streaming block produces. Adds a mapped-file regression test that mutation-kills the raise behavior and locks in the synthetic-stream contract. * fix(guardrails/bedrock): preserve upstream usage on streaming post_call block Non-streaming post_call blocks report the upstream LLM call's real token usage via ModifyResponseException.original_response, which the endpoint handler unwraps through _blocked_response_usage. Streaming post_call synthesizes its own ModelResponse locally (the exception can't escape the SSE generator), and previously left .usage unset, so the client saw accurate billing on non-streaming blocks and zero on streaming blocks -- silent revenue leak. Copy the assembled response's .usage onto the synthetic block response before yielding. Pre-refactor code had the same gap (create_guardrail_blocked_response never set usage); this is a net improvement, not a regression fix.
This commit is contained in:
yucheng-berri
GitHub
parent
1d9a86eac4
commit
6eed38bcfb
5 changed files with 529 additions and 123 deletions
|
|
@ -1180,20 +1180,6 @@ class ModifyResponseException(Exception):
|
|||
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.
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ import litellm
|
|||
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 (
|
||||
|
|
@ -754,7 +754,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
)
|
||||
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(
|
||||
|
|
@ -1027,8 +1029,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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.
|
||||
"""
|
||||
|
|
@ -1040,7 +1042,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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",
|
||||
|
|
@ -1134,18 +1142,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# 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,
|
||||
|
|
@ -1183,16 +1179,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################
|
||||
########## 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,
|
||||
)
|
||||
#########################################################
|
||||
|
||||
#########################################################
|
||||
|
|
@ -1207,8 +1202,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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 ##########
|
||||
|
|
@ -1248,16 +1241,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################
|
||||
########## 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,
|
||||
)
|
||||
#########################################################
|
||||
|
||||
#########################################################
|
||||
|
|
@ -1272,8 +1268,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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 ##########
|
||||
|
|
@ -1323,7 +1317,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# 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",
|
||||
|
|
@ -1332,15 +1330,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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,
|
||||
|
|
@ -1357,7 +1355,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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.
|
||||
|
|
@ -1369,8 +1367,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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)
|
||||
|
||||
|
|
@ -1422,7 +1418,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# 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",
|
||||
|
|
@ -1431,15 +1434,31 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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
|
||||
|
||||
#########################################################################
|
||||
########## 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,
|
||||
|
|
@ -1732,13 +1751,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
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))
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1390,7 +1390,14 @@ async def test_bedrock_guardrail_disable_exception_on_block_non_streaming():
|
|||
assert exception.status_code == 400
|
||||
assert "Violated guardrail policy" in str(exception.detail)
|
||||
|
||||
# Test 2: disable_exception_on_block=True - should NOT raise exception
|
||||
# Test 2: disable_exception_on_block=True - raises ModifyResponseException.
|
||||
# LIT-4186: pre-fix, the native hook swallowed the block and set
|
||||
# data["mock_response"], which was dead code (route_request already
|
||||
# unpacked kwargs) so during_call let the model call proceed anyway.
|
||||
# The correct contract is to raise ModifyResponseException so the endpoint
|
||||
# handler returns a 200 with the block message as content.
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
|
||||
guardrail_disabled = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
|
|
@ -1402,20 +1409,13 @@ async def test_bedrock_guardrail_disable_exception_on_block_non_streaming():
|
|||
) as mock_post:
|
||||
mock_post.return_value = mock_bedrock_response
|
||||
|
||||
# Should NOT raise exception when disable_exception_on_block=True
|
||||
try:
|
||||
response = await guardrail_disabled.async_moderation_hook(
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
await guardrail_disabled.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
call_type="completion",
|
||||
)
|
||||
# Should succeed and return data (even though content was blocked)
|
||||
assert response is not None
|
||||
print("✅ No exception raised when disable_exception_on_block=True")
|
||||
except Exception as e:
|
||||
pytest.fail(
|
||||
f"Should not raise exception when disable_exception_on_block=True, but got: {e}"
|
||||
)
|
||||
assert exc_info.value.message == "I can't provide that information."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1514,7 +1514,10 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming():
|
|||
async for chunk in result_generator:
|
||||
pass
|
||||
|
||||
# Test 2: disable_exception_on_block=True - should NOT raise exception
|
||||
# Test 2: disable_exception_on_block=True. Streaming can't raise up to the
|
||||
# endpoint handler (SSE headers already flushed), so the block is delivered
|
||||
# as a synthetic stream with finish_reason=content_filter and the block
|
||||
# message as content -- same shape a non-streaming block produces.
|
||||
guardrail_disabled = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
|
|
@ -1526,31 +1529,20 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming():
|
|||
) as mock_post:
|
||||
mock_post.return_value = mock_bedrock_response
|
||||
|
||||
# Should NOT raise exception when disable_exception_on_block=True
|
||||
try:
|
||||
result_generator = (
|
||||
guardrail_disabled.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
response=mock_streaming_response(),
|
||||
request_data=request_data,
|
||||
)
|
||||
)
|
||||
|
||||
# Consume the generator - should succeed without exceptions
|
||||
result_chunks = []
|
||||
async for chunk in result_generator:
|
||||
result_chunks.append(chunk)
|
||||
|
||||
# Should have received chunks back even though content was blocked
|
||||
assert len(result_chunks) > 0
|
||||
print(
|
||||
"✅ Streaming completed without exception when disable_exception_on_block=True"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail(
|
||||
f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}"
|
||||
)
|
||||
result_generator = guardrail_disabled.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
response=mock_streaming_response(),
|
||||
request_data=request_data,
|
||||
)
|
||||
chunks = [c async for c in result_generator]
|
||||
assert chunks, "streaming block should yield synthetic chunks, not empty"
|
||||
assembled_content = "".join(
|
||||
(c.choices[0].delta.content or "")
|
||||
for c in chunks
|
||||
if getattr(c, "choices", None) and getattr(c.choices[0], "delta", None)
|
||||
)
|
||||
assert assembled_content == "I can't provide that information."
|
||||
assert chunks[-1].choices[0].finish_reason == "content_filter"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -2767,3 +2767,408 @@ async def test_grounding_output_blocked_raises_400():
|
|||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
|
||||
###############################################################################
|
||||
# LIT-4186: disable_exception_on_block regression tests
|
||||
#
|
||||
# Before the fix, a Bedrock block with disable_exception_on_block=True raised
|
||||
# GuardrailInterventionNormalStringError, which no proxy code handled: the
|
||||
# unified pre_call path re-raised it, so the client saw HTTP 500 with the block
|
||||
# message; the native during_call hook swallowed it and set data["mock_response"],
|
||||
# which was dead code because route_request already unpacked kwargs.
|
||||
#
|
||||
# The fix converts blocks to ModifyResponseException at the raise site inside
|
||||
# make_bedrock_api_request. That exception is already the industry-standard
|
||||
# proxy contract (caught in proxy_server.py, anthropic_endpoints, etc.) and
|
||||
# turns into a 200 response whose content is the block message.
|
||||
###############################################################################
|
||||
|
||||
|
||||
def _blocked_bedrock_httpx_response() -> MagicMock:
|
||||
response = MagicMock()
|
||||
response.status_code = 200
|
||||
response.json.return_value = {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"outputs": [{"text": "Sorry, the model cannot answer this question."}],
|
||||
"assessments": [
|
||||
{
|
||||
"topicPolicy": {
|
||||
"topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}]
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
return response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_make_bedrock_api_request_block_raises_modify_response_when_flag_set():
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
disable_exception_on_block=True,
|
||||
)
|
||||
|
||||
request_data = {"model": "bedrock-nova-micro"}
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "k"
|
||||
mock_credentials.secret_key = "s"
|
||||
mock_credentials.token = None
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
guardrail.async_handler, "post", new_callable=AsyncMock
|
||||
) as mock_post,
|
||||
patch.object(
|
||||
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
|
||||
),
|
||||
patch.object(guardrail, "_prepare_request", return_value=MagicMock()),
|
||||
):
|
||||
mock_post.return_value = _blocked_bedrock_httpx_response()
|
||||
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
await guardrail.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=[{"role": "user", "content": "My name is John Doe"}],
|
||||
request_data=request_data,
|
||||
)
|
||||
|
||||
assert exc_info.value.message == "Sorry, the model cannot answer this question."
|
||||
assert exc_info.value.model == "bedrock-nova-micro"
|
||||
assert exc_info.value.guardrail_name == "test-bedrock-guard"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_make_bedrock_api_request_block_raises_http_400_when_flag_unset():
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
disable_exception_on_block=False,
|
||||
)
|
||||
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "k"
|
||||
mock_credentials.secret_key = "s"
|
||||
mock_credentials.token = None
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
guardrail.async_handler, "post", new_callable=AsyncMock
|
||||
) as mock_post,
|
||||
patch.object(
|
||||
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
|
||||
),
|
||||
patch.object(guardrail, "_prepare_request", return_value=MagicMock()),
|
||||
):
|
||||
mock_post.return_value = _blocked_bedrock_httpx_response()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
request_data={"model": "bedrock-nova-micro"},
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_pre_call_hook_propagates_modify_response_on_block():
|
||||
"""pre_call: block with disable_exception_on_block=True must raise
|
||||
ModifyResponseException so the endpoint handler returns 200 with the block
|
||||
message. Before LIT-4186 the exception was swallowed and only data
|
||||
["mock_response"] was mutated, which the unified pre_call path never read
|
||||
(surfaced as HTTP 500)."""
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
disable_exception_on_block=True,
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"model": "bedrock-nova-micro",
|
||||
"messages": [{"role": "user", "content": "My name is John Doe"}],
|
||||
}
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "k"
|
||||
mock_credentials.secret_key = "s"
|
||||
mock_credentials.token = None
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
guardrail.async_handler, "post", new_callable=AsyncMock
|
||||
) as mock_post,
|
||||
patch.object(
|
||||
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
|
||||
),
|
||||
patch.object(guardrail, "_prepare_request", return_value=MagicMock()),
|
||||
):
|
||||
mock_post.return_value = _blocked_bedrock_httpx_response()
|
||||
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
cache=DualCache(),
|
||||
data=request_data,
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
assert exc_info.value.message == "Sorry, the model cannot answer this question."
|
||||
# No `mock_response` mutation: the old broken contract must be gone
|
||||
# (route_request unpacks kwargs before this hook runs, so `mock_response`
|
||||
# would never reach the LLM call anyway).
|
||||
assert "mock_response" not in request_data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_moderation_hook_propagates_modify_response_on_block():
|
||||
"""during_call: block must raise ModifyResponseException from the moderation
|
||||
task so the surrounding asyncio.gather cancels the LLM call, instead of
|
||||
the old behavior of swallowing the block and letting the model call proceed
|
||||
(LIT-4186 symptom 2: silent bypass, model billed)."""
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
disable_exception_on_block=True,
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"model": "bedrock-nova-micro",
|
||||
"messages": [{"role": "user", "content": "My name is John Doe"}],
|
||||
}
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "k"
|
||||
mock_credentials.secret_key = "s"
|
||||
mock_credentials.token = None
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
guardrail.async_handler, "post", new_callable=AsyncMock
|
||||
) as mock_post,
|
||||
patch.object(
|
||||
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
|
||||
),
|
||||
patch.object(guardrail, "_prepare_request", return_value=MagicMock()),
|
||||
):
|
||||
mock_post.return_value = _blocked_bedrock_httpx_response()
|
||||
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
await guardrail.async_moderation_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
call_type="acompletion",
|
||||
)
|
||||
|
||||
assert exc_info.value.message == "Sorry, the model cannot answer this question."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_success_hook_attaches_original_response_on_block():
|
||||
"""post_call: block must raise ModifyResponseException and attach the LLM
|
||||
response to `original_response` so the synthetic block reply reports the
|
||||
upstream call's real token usage instead of zero."""
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
disable_exception_on_block=True,
|
||||
)
|
||||
|
||||
request_data = {
|
||||
"model": "bedrock-nova-micro",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
llm_response = _model_response("Hello John Doe! The capital of France is Paris.")
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "k"
|
||||
mock_credentials.secret_key = "s"
|
||||
mock_credentials.token = None
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
guardrail.async_handler, "post", new_callable=AsyncMock
|
||||
) as mock_post,
|
||||
patch.object(
|
||||
guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")
|
||||
),
|
||||
patch.object(guardrail, "_prepare_request", return_value=MagicMock()),
|
||||
):
|
||||
mock_post.return_value = _blocked_bedrock_httpx_response()
|
||||
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
await guardrail.async_post_call_success_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=llm_response,
|
||||
)
|
||||
|
||||
assert exc_info.value.original_response is llm_response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_propagates_modify_response_on_block():
|
||||
"""apply_guardrail (unified path used by pre_call / /apply_guardrail
|
||||
endpoint) must let ModifyResponseException propagate as-is so the endpoint
|
||||
handler catches it and returns a 200."""
|
||||
from litellm.exceptions import ModifyResponseException
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
disable_exception_on_block=True,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "make_bedrock_api_request", new_callable=AsyncMock
|
||||
) as mock_api:
|
||||
mock_api.side_effect = ModifyResponseException(
|
||||
message="Sorry, the model cannot answer this question.",
|
||||
model="bedrock-nova-micro",
|
||||
request_data={},
|
||||
guardrail_name="test-bedrock-guard",
|
||||
)
|
||||
|
||||
with pytest.raises(ModifyResponseException) as exc_info:
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["My name is John Doe"]},
|
||||
request_data={"model": "bedrock-nova-micro"},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
assert exc_info.value.message == "Sorry, the model cannot answer this question."
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_post_call_block_yields_synthetic_stream_not_raise():
|
||||
"""LIT-4186 regression: with disable_exception_on_block=True, streaming
|
||||
post_call blocks must be delivered as a synthetic stream (finish_reason=
|
||||
content_filter, block message as content), NOT raised. Pre-fix the local
|
||||
handler already produced this shape; the LIT-4186 refactor briefly turned
|
||||
it into an SSE 500 by letting ModifyResponseException escape the streaming
|
||||
generator. This test locks in the correct streaming contract.
|
||||
"""
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
disable_exception_on_block=True,
|
||||
)
|
||||
|
||||
async def _stream():
|
||||
yield ModelResponseStream(
|
||||
choices=[
|
||||
StreamingChoices(
|
||||
index=0,
|
||||
delta=Delta(role="assistant", content="Coffee is a popular"),
|
||||
)
|
||||
]
|
||||
)
|
||||
yield ModelResponseStream(
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content=" beverage."), finish_reason="stop")]
|
||||
)
|
||||
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "k"
|
||||
mock_credentials.secret_key = "s"
|
||||
mock_credentials.token = None
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post,
|
||||
patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")),
|
||||
patch.object(guardrail, "_prepare_request", return_value=MagicMock()),
|
||||
):
|
||||
mock_post.return_value = _blocked_bedrock_httpx_response()
|
||||
|
||||
chunks = [
|
||||
c
|
||||
async for c in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=_stream(),
|
||||
request_data={"model": "bedrock-nova-micro"},
|
||||
)
|
||||
]
|
||||
|
||||
assert chunks, "streaming block should yield synthetic chunks, not error out"
|
||||
assembled_content = "".join(
|
||||
(c.choices[0].delta.content or "")
|
||||
for c in chunks
|
||||
if getattr(c, "choices", None) and getattr(c.choices[0], "delta", None)
|
||||
)
|
||||
assert assembled_content == "Sorry, the model cannot answer this question."
|
||||
assert chunks[-1].choices[0].finish_reason == "content_filter"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_post_call_block_preserves_upstream_usage():
|
||||
"""LIT-4186: streaming block must report the usage the upstream LLM call
|
||||
actually consumed. Non-streaming blocks carry it via original_response +
|
||||
_blocked_response_usage in the endpoint handler; streaming has to copy it
|
||||
onto the synthetic ModelResponse directly since the exception can't escape
|
||||
the SSE generator. Without this, clients see accurate billing on
|
||||
non-streaming blocks and zero on streaming blocks -- silent revenue leak."""
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices, Usage
|
||||
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrail_name="test-bedrock-guard",
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
disable_exception_on_block=True,
|
||||
)
|
||||
|
||||
async def _stream_with_usage():
|
||||
# Terminal chunk carrying usage, as OpenAI-style streams do with
|
||||
# stream_options={"include_usage": True}. stream_chunk_builder
|
||||
# aggregates this into the assembled ModelResponse's .usage.
|
||||
yield ModelResponseStream(
|
||||
choices=[StreamingChoices(index=0, delta=Delta(role="assistant", content="Coffee is delicious"))]
|
||||
)
|
||||
yield ModelResponseStream(
|
||||
choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")],
|
||||
usage=Usage(prompt_tokens=42, completion_tokens=17, total_tokens=59),
|
||||
)
|
||||
|
||||
mock_credentials = MagicMock()
|
||||
mock_credentials.access_key = "k"
|
||||
mock_credentials.secret_key = "s"
|
||||
mock_credentials.token = None
|
||||
|
||||
with (
|
||||
patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post,
|
||||
patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")),
|
||||
patch.object(guardrail, "_prepare_request", return_value=MagicMock()),
|
||||
):
|
||||
mock_post.return_value = _blocked_bedrock_httpx_response()
|
||||
|
||||
chunks = [
|
||||
c
|
||||
async for c in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
response=_stream_with_usage(),
|
||||
request_data={"model": "bedrock-nova-micro"},
|
||||
)
|
||||
]
|
||||
|
||||
# Find the chunk carrying usage (MockResponseIterator emits it on the
|
||||
# terminating chunk when the source ModelResponse has .usage set)
|
||||
usage_chunks = [c for c in chunks if getattr(c, "usage", None) is not None]
|
||||
assert usage_chunks, "streaming block should carry the upstream call's usage on at least one chunk"
|
||||
reported_usage = usage_chunks[-1].usage
|
||||
assert reported_usage.prompt_tokens == 42
|
||||
assert reported_usage.completion_tokens == 17
|
||||
assert reported_usage.total_tokens == 59
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue