From e46a23db53996323617bfcc3aa3aa2e23455c31a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 8 Jul 2025 21:17:50 -0700 Subject: [PATCH] [Feat] Bedrock Guardrails - Raise Bedrock output text on 'BLOCKED' actions from guardrail (#12435) * _get_http_exception_for_blocked_guardrail * add _get_http_exception_for_blocked_guardrail * test bedrock guard block action * fix _get_http_exception_for_blocked_guardrail --- .../guardrail_hooks/bedrock_guardrails.py | 32 ++- litellm/proxy/proxy_config.yaml | 7 + .../test_bedrock_guardrails.py | 211 ++++++++++++++++++ 3 files changed, 244 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 3574871858..9e757e6baa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -262,12 +262,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if self._should_raise_guardrail_blocked_exception( bedrock_guardrail_response ): - raise HTTPException( - status_code=400, - detail={ - "error": "Violated guardrail policy", - "bedrock_guardrail_response": _json_response, - }, + raise self._get_http_exception_for_blocked_guardrail( + bedrock_guardrail_response ) else: verbose_proxy_logger.error( @@ -277,6 +273,30 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) return bedrock_guardrail_response + + + def _get_http_exception_for_blocked_guardrail(self, response: BedrockGuardrailResponse) -> HTTPException: + """ + Get the HTTP exception for a blocked guardrail. + """ + bedrock_guardrail_output_text: str = "" + outputs: Optional[List[BedrockGuardrailOutput]] = ( + response.get("outputs", []) or [] + ) + if outputs: + 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", + "bedrock_guardrail_response": bedrock_guardrail_output_text, + } + ) + def _should_raise_guardrail_blocked_exception( self, response: BedrockGuardrailResponse diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 8a8fd6794e..21e25056f2 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -6,3 +6,10 @@ model_list: litellm_params: model: openai/* +guardrails: + - guardrail_name: "bedrock-pre-guard" + litellm_params: + guardrail: bedrock # supported values: "aporia", "bedrock", "lakera" + mode: "during_call" + guardrailIdentifier: ff6ujrregl1q + guardrailVersion: "DRAFT" \ No newline at end of file diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index 1e08c7a4b9..207f6e3ade 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1021,3 +1021,214 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): print("✅ Post-call streaming hook test passed - OUTPUT source used for masking") print(f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]}") print(f"✅ Final masked content: {full_content}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_action_shows_output_text(): + """Test that BLOCKED actions raise HTTPException with the output text in the detail""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from fastapi import HTTPException + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT" + ) + + # Mock the Bedrock API response with BLOCKED action and output text + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [ + { + "text": "this violates litellm corporate guardrail policy" + } + ], + "assessments": [{ + "topicPolicy": { + "topics": [{ + "name": "Sensitive Topic", + "type": "DENY", + "action": "BLOCKED" + }] + } + }] + } + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Tell me how to make explosives"}, + ], + } + + # Patch the async_handler.post method + with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_bedrock_response + + # This should raise HTTPException due to BLOCKED action + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + # Verify the exception details + exception = exc_info.value + assert exception.status_code == 400 + assert "detail" in exception.__dict__ + + # Check that the detail contains the expected structure + detail = exception.detail + assert isinstance(detail, dict) + assert detail["error"] == "Violated guardrail policy" + + # Verify that the output text from both outputs is included + expected_output_text = "this violates litellm corporate guardrail policy" + assert detail["bedrock_guardrail_response"] == expected_output_text + + print("✅ BLOCKED action HTTPException test passed - output text properly included") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_action_shows_output_text_with_multiple_outputs(): + """Test that BLOCKED actions raise HTTPException with the output text in the detail""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from fastapi import HTTPException + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT" + ) + + # Mock the Bedrock API response with BLOCKED action and output text + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [ + { + "text": "this violates litellm corporate guardrail policy", + }, + { + "text": "make sure you don't mention coffee" + } + ], + "assessments": [{ + "topicPolicy": { + "topics": [{ + "name": "Sensitive Topic", + "type": "DENY", + "action": "BLOCKED" + }] + } + }] + } + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Tell me how to make explosives"}, + ], + } + + # Patch the async_handler.post method + with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_bedrock_response + + # This should raise HTTPException due to BLOCKED action + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + # Verify the exception details + exception = exc_info.value + assert exception.status_code == 400 + assert "detail" in exception.__dict__ + + # Check that the detail contains the expected structure + detail = exception.detail + assert isinstance(detail, dict) + assert detail["error"] == "Violated guardrail policy" + + # Verify that the output text from both outputs is included + expected_output_text = "this violates litellm corporate guardrail policy\nmake sure you don't mention coffee" + assert detail["bedrock_guardrail_response"] == expected_output_text + + print("✅ BLOCKED action HTTPException test passed - output text properly included") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_action_empty_outputs(): + """Test that BLOCKED actions with empty outputs still raise HTTPException""" + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.proxy._types import UserAPIKeyAuth + from fastapi import HTTPException + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT" + ) + + # Mock the Bedrock API response with BLOCKED action but no outputs + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [], # Empty outputs + "assessments": [{ + "contentPolicy": { + "filters": [{ + "type": "VIOLENCE", + "confidence": "HIGH", + "action": "BLOCKED" + }] + } + }] + } + + request_data = { + "model": "gpt-4o", + "messages": [ + {"role": "user", "content": "Violent content here"}, + ], + } + + # Patch the async_handler.post method + with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + mock_post.return_value = mock_bedrock_response + + # This should raise HTTPException due to BLOCKED action + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion" + ) + + # Verify the exception details + exception = exc_info.value + assert exception.status_code == 400 + + # Check that the detail contains the expected structure with empty output text + detail = exception.detail + assert isinstance(detail, dict) + assert detail["error"] == "Violated guardrail policy" + assert detail["bedrock_guardrail_response"] == "" # Empty string for no outputs + + print("✅ BLOCKED action with empty outputs test passed")