diff --git a/litellm/exceptions.py b/litellm/exceptions.py index adf7b3ef05a..aca3fb551cc 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -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. diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 6e46f971dd8..a45719d2eb8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -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)) diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index 87c5e3bf2a9..f257b47404e 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -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 diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index a23e89e576c..823ee05839f 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index f43d8e85aca..bf237d8017b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -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