import sys import os import io, asyncio import pytest sys.path.insert(0, os.path.abspath("../..")) import litellm from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail from litellm.proxy._types import UserAPIKeyAuth from litellm.caching import DualCache from unittest.mock import MagicMock, AsyncMock, patch @pytest.mark.asyncio async def test_bedrock_guardrails_pii_masking(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", ) request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, {"role": "assistant", "content": "Hello, how can I help you today?"}, {"role": "user", "content": "I need to cancel my order"}, {"role": "user", "content": "ok, my credit card number is 1234-5678-9012-3456"}, ], } response = await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, call_type="completion" ) print("response after moderation hook", response) if response: # Only assert if response is not None assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}" assert response["messages"][1]["content"] == "Hello, how can I help you today?" assert response["messages"][2]["content"] == "I need to cancel my order" assert response["messages"][3]["content"] == "ok, my credit card number is {CREDIT_DEBIT_CARD_NUMBER}" @pytest.mark.asyncio async def test_bedrock_guardrails_pii_masking_content_list(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", ) request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": [ {"type": "text", "text": "Hello, my phone number is +1 412 555 1212"}, {"type": "text", "text": "what time is it?"}, ]}, {"role": "assistant", "content": "Hello, how can I help you today?"}, { "role": "user", "content": "who is the president of the united states?" } ], } response = await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, call_type="completion" ) print(response) if response: # Only assert if response is not None # Verify that the list content is properly masked assert isinstance(response["messages"][0]["content"], list) assert response["messages"][0]["content"][0]["text"] == "Hello, my phone number is {PHONE}" assert response["messages"][0]["content"][1]["text"] == "what time is it?" assert response["messages"][1]["content"] == "Hello, how can I help you today?" assert response["messages"][2]["content"] == "who is the president of the united states?" @pytest.mark.asyncio async def test_bedrock_guardrails_with_streaming(): from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks # Create proper mock objects mock_user_api_key_cache = MagicMock(spec=DualCache) mock_user_api_key_dict = UserAPIKeyAuth() with pytest.raises(Exception): # Assert that this raises an exception proxy_logging_obj = ProxyLogging( user_api_key_cache=mock_user_api_key_cache, premium_user=True, ) guardrail = BedrockGuardrail( guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", ) litellm.callbacks.append(guardrail) request_data = { "model": "gpt-4o", "messages": [ { "role": "user", "content": "Hi I like coffee" } ], "stream": True, "metadata": {"guardrails": ["bedrock-post-guard"]} } response = await litellm.acompletion( **request_data, ) response = proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, response=response, request_data=request_data, ) async for chunk in response: print(chunk) @pytest.mark.asyncio async def test_bedrock_guardrails_with_streaming_no_violation(): from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks # Create proper mock objects mock_user_api_key_cache = MagicMock(spec=DualCache) mock_user_api_key_dict = UserAPIKeyAuth() proxy_logging_obj = ProxyLogging( user_api_key_cache=mock_user_api_key_cache, premium_user=True, ) guardrail = BedrockGuardrail( guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", ) litellm.callbacks.append(guardrail) request_data = { "model": "gpt-4o", "messages": [ { "role": "user", "content": "hi" } ], "stream": True, "metadata": {"guardrails": ["bedrock-post-guard"]} } response = await litellm.acompletion( **request_data, ) response = proxy_logging_obj.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, response=response, request_data=request_data, ) async for chunk in response: print(chunk) @pytest.mark.asyncio async def test_bedrock_guardrails_streaming_request_body_mock(): """Test that the exact request body sent to Bedrock matches expected format when using streaming""" import json from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from litellm.caching import DualCache from litellm.types.guardrails import GuardrailEventHooks # Create mock objects mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) # Create the guardrail guardrail = BedrockGuardrail( guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", ) # Mock the assembled response from streaming mock_response = litellm.ModelResponse( id="test-id", choices=[ litellm.Choices( index=0, message=litellm.Message( role="assistant", content="The capital of Spain is Madrid." ), finish_reason="stop" ) ], created=1234567890, model="gpt-4o", object="chat.completion" ) # Mock Bedrock API response mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "NONE", "outputs": [] } # Patch the async_handler.post method to capture the request body with patch.object(guardrail, 'async_handler') as mock_async_handler: mock_async_handler.post = AsyncMock(return_value=mock_bedrock_response) # Test data - simulating request data and assembled response request_data = { "model": "gpt-4o", "messages": [ { "role": "user", "content": "what's the capital of spain?" } ], "stream": True, "metadata": {"guardrails": ["bedrock-post-guard"]} } # Call the method that should make the Bedrock API request await guardrail.make_bedrock_api_request( source="OUTPUT", response=mock_response, request_data=request_data ) # Verify the API call was made mock_async_handler.post.assert_called_once() # Get the request data that was passed call_args = mock_async_handler.post.call_args # The data should be in the 'data' parameter of the prepared request # We need to parse the JSON from the prepared request body prepared_request_body = call_args.kwargs.get('data') # Parse the JSON body if isinstance(prepared_request_body, bytes): actual_body = json.loads(prepared_request_body.decode('utf-8')) else: actual_body = json.loads(prepared_request_body) # Expected body based on the convert_to_bedrock_format method behavior expected_body = { 'source': 'OUTPUT', 'content': [ {'text': {'text': 'The capital of Spain is Madrid.'}} ] } print("Actual Bedrock request body:", json.dumps(actual_body, indent=2)) print("Expected Bedrock request body:", json.dumps(expected_body, indent=2)) # Assert the request body matches exactly assert actual_body == expected_body, f"Request body mismatch. Expected: {expected_body}, Got: {actual_body}" @pytest.mark.asyncio async def test_bedrock_guardrail_aws_param_persistence(): """Test that AWS auth params set on init are used for every request and not popped out.""" from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks guardrail = BedrockGuardrail( guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", aws_access_key_id="test-access-key", aws_secret_access_key="test-secret-key", aws_region_name="us-east-1", supported_event_hooks=[GuardrailEventHooks.post_call], guardrail_name="bedrock-post-guard", ) with patch.object(guardrail, "get_credentials", wraps=guardrail.get_credentials) as mock_get_creds: for i in range(3): request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": f"request {i}"} ], "stream": False, "metadata": {"guardrails": ["bedrock-post-guard"]} } with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: # Configure the mock response properly mock_response = AsyncMock() mock_response.status_code = 200 mock_response.json = MagicMock(return_value={"action": "NONE", "outputs": []}) mock_post.return_value = mock_response await guardrail.make_bedrock_api_request(source="INPUT", messages=request_data.get("messages"), request_data=request_data) assert mock_get_creds.call_count == 3 for call in mock_get_creds.call_args_list: kwargs = call.kwargs print("used the following kwargs to get credentials=", kwargs) assert kwargs["aws_access_key_id"] == "test-access-key" assert kwargs["aws_secret_access_key"] == "test-secret-key" assert kwargs["aws_region_name"] == "us-east-1" @pytest.mark.asyncio async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" from unittest.mock import MagicMock from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrailResponse guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) # Test 1: ANONYMIZED action should NOT raise exception anonymized_response: BedrockGuardrailResponse = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "Hello, my phone number is {PHONE}" }], "assessments": [{ "sensitiveInformationPolicy": { "piiEntities": [{ "type": "PHONE", "match": "+1 412 555 1212", "action": "ANONYMIZED" }] } }] } should_raise = guardrail._should_raise_guardrail_blocked_exception(anonymized_response) assert should_raise is False, "ANONYMIZED actions should not raise exceptions" # Test 2: BLOCKED action should raise exception blocked_response: BedrockGuardrailResponse = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "I can't provide that information." }], "assessments": [{ "topicPolicy": { "topics": [{ "name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED" }] } }] } should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response) assert should_raise is True, "BLOCKED actions should raise exceptions" # Test 3: Mixed actions - should raise if ANY action is BLOCKED mixed_response: BedrockGuardrailResponse = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "I can't provide that information." }], "assessments": [{ "sensitiveInformationPolicy": { "piiEntities": [{ "type": "PHONE", "match": "+1 412 555 1212", "action": "ANONYMIZED" }] }, "topicPolicy": { "topics": [{ "name": "Blocked Topic", "type": "DENY", "action": "BLOCKED" }] } }] } should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) assert should_raise is True, "Mixed actions with any BLOCKED should raise exceptions" # Test 4: NONE action should not raise exception none_response: BedrockGuardrailResponse = { "action": "NONE", "outputs": [], "assessments": [] } should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response) assert should_raise is False, "NONE actions should not raise exceptions" # Test 5: Test other policy types with BLOCKED actions content_blocked_response: BedrockGuardrailResponse = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "I can't provide that information." }], "assessments": [{ "contentPolicy": { "filters": [{ "type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED" }] } }] } should_raise = guardrail._should_raise_guardrail_blocked_exception(content_blocked_response) assert should_raise is True, "Content policy BLOCKED actions should raise exceptions" @pytest.mark.asyncio async def test_bedrock_guardrail_masking_with_anonymized_response(): """Test that masking works correctly when guardrail returns ANONYMIZED actions""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from litellm.caching import DualCache # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", mask_request_content=True, ) # Mock the Bedrock API response with ANONYMIZED action mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "Hello, my phone number is {PHONE}" }], "assessments": [{ "sensitiveInformationPolicy": { "piiEntities": [{ "type": "PHONE", "match": "+1 412 555 1212", "action": "ANONYMIZED" }] } }] } request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, ], } # 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 NOT raise an exception since action is ANONYMIZED try: response = await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, call_type="completion" ) # Should succeed and return data with masked content assert response is not None assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}" except Exception as e: pytest.fail(f"Should not raise exception for ANONYMIZED actions, but got: {e}") @pytest.mark.asyncio async def test_bedrock_guardrail_uses_masked_output_without_masking_flags(): """Test that masked output from guardrails is used even when masking flags are not enabled""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() # Create guardrail WITHOUT masking flags enabled guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", # Note: No mask_request_content=True or mask_response_content=True ) # Mock the Bedrock API response with ANONYMIZED action and masked output mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "Hello, my phone number is {PHONE} and email is {EMAIL}" }], "assessments": [{ "sensitiveInformationPolicy": { "piiEntities": [ { "type": "PHONE", "match": "+1 412 555 1212", "action": "ANONYMIZED" }, { "type": "EMAIL", "match": "user@example.com", "action": "ANONYMIZED" } ] } }] } request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "Hello, my phone number is +1 412 555 1212 and email is user@example.com"}, ], } # 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 use the masked output even without masking flags response = await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, call_type="completion" ) # Should use the masked content from guardrail output assert response is not None assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE} and email is {EMAIL}" print("✅ Masked output was applied even without masking flags enabled") @pytest.mark.asyncio async def test_bedrock_guardrail_response_pii_masking_non_streaming(): """Test that PII masking is applied to response content in non-streaming scenarios""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() # Create guardrail with response masking enabled guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", ) # Mock the Bedrock API response with ANONYMIZED PII mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" }], "assessments": [{ "sensitiveInformationPolicy": { "piiEntities": [ { "type": "CREDIT_DEBIT_CARD_NUMBER", "match": "1234-5678-9012-3456", "action": "ANONYMIZED" }, { "type": "PHONE", "match": "+1 412 555 1212", "action": "ANONYMIZED" } ] } }] } # Create a mock response that contains PII mock_response = litellm.ModelResponse( id="test-id", choices=[ litellm.Choices( index=0, message=litellm.Message( role="assistant", content="My credit card number is 1234-5678-9012-3456 and my phone is +1 412 555 1212" ), finish_reason="stop" ) ], created=1234567890, model="gpt-4o", object="chat.completion" ) request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "What's your credit card and phone number?"}, ], } # 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 # Call the post-call success hook await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, response=mock_response ) # Verify that the response content was masked assert mock_response.choices[0].message.content == "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" print("✓ Non-streaming response PII masking test passed") @pytest.mark.asyncio async def test_bedrock_guardrail_response_pii_masking_streaming(): """Test that PII masking is applied to response content in streaming scenarios""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ModelResponseStream # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() # Create guardrail with response masking enabled guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", ) # Mock the Bedrock API response with ANONYMIZED PII mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "Sure! My email is {EMAIL} and SSN is {US_SSN}" }], "assessments": [{ "sensitiveInformationPolicy": { "piiEntities": [ { "type": "EMAIL", "match": "john@example.com", "action": "ANONYMIZED" }, { "type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED" } ] } }] } # Create mock streaming chunks async def mock_streaming_response(): chunks = [ ModelResponseStream( id="test-id", choices=[ litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="Sure! My email is "), finish_reason=None ) ], created=1234567890, model="gpt-4o", object="chat.completion.chunk" ), ModelResponseStream( id="test-id", choices=[ litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="john@example.com and SSN is "), finish_reason=None ) ], created=1234567890, model="gpt-4o", object="chat.completion.chunk" ), ModelResponseStream( id="test-id", choices=[ litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="123-45-6789"), finish_reason="stop" ) ], created=1234567890, model="gpt-4o", object="chat.completion.chunk" ) ] for chunk in chunks: yield chunk request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "What's your email and SSN?"}, ], "stream": True, } # 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 # Call the streaming hook masked_stream = guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, response=mock_streaming_response(), request_data=request_data ) # Collect all chunks from the masked stream masked_chunks = [] async for chunk in masked_stream: masked_chunks.append(chunk) # Verify that we got chunks back assert len(masked_chunks) > 0 # Reconstruct the full response from chunks to verify masking full_content = "" for chunk in masked_chunks: if hasattr(chunk, 'choices') and chunk.choices: if hasattr(chunk.choices[0], 'delta') and chunk.choices[0].delta: if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content: full_content += chunk.choices[0].delta.content # Verify that the reconstructed content contains the masked PII assert "Sure! My email is {EMAIL} and SSN is {US_SSN}" == full_content print("✓ Streaming response PII masking test passed") @pytest.mark.asyncio async def test_convert_to_bedrock_format_input_source(): """Test convert_to_bedrock_format with INPUT source and mock messages""" from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockRequest from unittest.mock import patch # Create the guardrail instance guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) # Mock messages mock_messages = [ {"role": "user", "content": "Hello, how are you?"}, {"role": "assistant", "content": "I'm doing well, thank you!"}, {"role": "user", "content": [ {"type": "text", "text": "What's the weather like?"}, {"type": "text", "text": "Is it sunny today?"} ]} ] # Call the method result = guardrail.convert_to_bedrock_format( source="INPUT", messages=mock_messages ) # Verify the result structure assert isinstance(result, dict) assert result.get("source") == "INPUT" assert "content" in result assert isinstance(result.get("content"), list) # Verify content items expected_content_items = [ {"text": {"text": "Hello, how are you?"}}, {"text": {"text": "I'm doing well, thank you!"}}, {"text": {"text": "What's the weather like?"}}, {"text": {"text": "Is it sunny today?"}} ] assert result.get("content") == expected_content_items print("✅ INPUT source test passed - result:", result) @pytest.mark.asyncio async def test_convert_to_bedrock_format_output_source(): """Test convert_to_bedrock_format with OUTPUT source and mock ModelResponse""" from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockRequest import litellm from unittest.mock import patch # Create the guardrail instance guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) # Mock ModelResponse mock_response = litellm.ModelResponse( id="test-response-id", choices=[ litellm.Choices( index=0, message=litellm.Message( role="assistant", content="This is a test response from the model." ), finish_reason="stop" ), litellm.Choices( index=1, message=litellm.Message( role="assistant", content="This is a second choice response." ), finish_reason="stop" ) ], created=1234567890, model="gpt-4o", object="chat.completion" ) # Call the method result = guardrail.convert_to_bedrock_format( source="OUTPUT", response=mock_response ) # Verify the result structure assert isinstance(result, dict) assert result.get("source") == "OUTPUT" assert "content" in result assert isinstance(result.get("content"), list) # Verify content items - should contain both choice contents expected_content_items = [ {"text": {"text": "This is a test response from the model."}}, {"text": {"text": "This is a second choice response."}} ] assert result.get("content") == expected_content_items print("✅ OUTPUT source test passed - result:", result) @pytest.mark.asyncio async def test_convert_to_bedrock_format_post_call_streaming_hook(): """Test async_post_call_streaming_iterator_hook makes OUTPUT bedrock request and applies masking""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ModelResponseStream import litellm # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() # Create guardrail instance guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) # Mock streaming chunks that contain PII async def mock_streaming_response(): chunks = [ ModelResponseStream( id="test-id", choices=[ litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="My email is "), finish_reason=None ) ], created=1234567890, model="gpt-4o", object="chat.completion.chunk" ), ModelResponseStream( id="test-id", choices=[ litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="john@example.com"), finish_reason="stop" ) ], created=1234567890, model="gpt-4o", object="chat.completion.chunk" ) ] for chunk in chunks: yield chunk # Mock Bedrock API response with PII masking mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "My email is {EMAIL}" }], "assessments": [{ "sensitiveInformationPolicy": { "piiEntities": [{ "type": "EMAIL", "match": "john@example.com", "action": "ANONYMIZED" }] } }] } request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "What's your email?"} ], "stream": True } # Track which bedrock API calls were made bedrock_calls = [] # Mock the make_bedrock_api_request method to track calls async def mock_make_bedrock_api_request(source, messages=None, response=None, request_data=None): bedrock_calls.append({ "source": source, "messages": messages, "response": response, "request_data": request_data }) # Return the mock bedrock response from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrailResponse return BedrockGuardrailResponse(**mock_bedrock_response.json()) # Patch the bedrock API request method with patch.object(guardrail, 'make_bedrock_api_request', side_effect=mock_make_bedrock_api_request): # Call the streaming hook result_generator = guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, response=mock_streaming_response(), request_data=request_data ) # Collect all chunks from the result result_chunks = [] async for chunk in result_generator: result_chunks.append(chunk) # Verify bedrock API calls were made assert len(bedrock_calls) == 2, f"Expected 2 bedrock calls (INPUT and OUTPUT), got {len(bedrock_calls)}" # Find the OUTPUT call output_calls = [call for call in bedrock_calls if call["source"] == "OUTPUT"] assert len(output_calls) == 1, f"Expected 1 OUTPUT call, got {len(output_calls)}" output_call = output_calls[0] assert output_call["source"] == "OUTPUT" assert output_call["response"] is not None assert output_call["messages"] is None # OUTPUT calls don't need messages # Verify that the response content was masked # The streaming chunks should now contain the masked content full_content = "" for chunk in result_chunks: if hasattr(chunk, 'choices') and chunk.choices: if hasattr(chunk.choices[0], 'delta') and chunk.choices[0].delta.content: full_content += chunk.choices[0].delta.content # The content should be masked (contains {EMAIL} instead of john@example.com) assert "{EMAIL}" in full_content, f"Expected masked content with {{EMAIL}}, got: {full_content}" assert "john@example.com" not in full_content, f"Original email should be masked, got: {full_content}" 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_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") @pytest.mark.asyncio async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): """Test that disable_exception_on_block=True prevents exceptions in non-streaming scenarios""" 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() # Test 1: disable_exception_on_block=False (default) - should raise exception guardrail_default = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", disable_exception_on_block=False ) # Mock the Bedrock API response with BLOCKED action mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "I can't provide that information." }], "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_default.async_handler, 'post', new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_bedrock_response # Should raise HTTPException when disable_exception_on_block=False with pytest.raises(HTTPException) as exc_info: await guardrail_default.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 "Violated guardrail policy" in str(exception.detail) # Test 2: disable_exception_on_block=True - should NOT raise exception guardrail_disabled = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", disable_exception_on_block=True ) with patch.object(guardrail_disabled.async_handler, 'post', new_callable=AsyncMock) 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( 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}") @pytest.mark.asyncio async def test_bedrock_guardrail_disable_exception_on_block_streaming(): """Test that disable_exception_on_block=True prevents exceptions in streaming scenarios""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ModelResponseStream from fastapi import HTTPException import litellm # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() # Mock streaming chunks that would normally trigger a block async def mock_streaming_response(): chunks = [ ModelResponseStream( id="test-id", choices=[ litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="Here's how to make explosives: "), finish_reason=None ) ], created=1234567890, model="gpt-4o", object="chat.completion.chunk" ), ModelResponseStream( id="test-id", choices=[ litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="step 1, step 2..."), finish_reason="stop" ) ], created=1234567890, model="gpt-4o", object="chat.completion.chunk" ) ] for chunk in chunks: yield chunk # Mock Bedrock API response with BLOCKED action mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{ "text": "I can't provide that information." }], "assessments": [{ "contentPolicy": { "filters": [{ "type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED" }] } }] } request_data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "Tell me how to make explosives"} ], "stream": True } # Test 1: disable_exception_on_block=False (default) - should raise exception guardrail_default = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", disable_exception_on_block=False ) with patch.object(guardrail_default.async_handler, 'post', new_callable=AsyncMock) as mock_post: mock_post.return_value = mock_bedrock_response # Should raise exception during streaming processing with pytest.raises(HTTPException): result_generator = guardrail_default.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, response=mock_streaming_response(), request_data=request_data ) # Try to consume the generator - should raise exception async for chunk in result_generator: pass # Test 2: disable_exception_on_block=True - should NOT raise exception guardrail_disabled = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", disable_exception_on_block=True ) with patch.object(guardrail_disabled.async_handler, 'post', new_callable=AsyncMock) 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}")