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, _redact_pii_matches, ) 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-5.5", "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-5.5", "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_block_messages_api(): """ Test that guardrails block messages API requests containing 'coffee' and raise the expected exception. """ from fastapi import HTTPException # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", ) request_data = { "model": "claude-sonnet-4-5-20250929", "messages": [ { "role": "user", "content": [ { "type": "text", "text": "Hello, my phone number is +1 412 555 1212", }, {"type": "text", "text": "what time is it?"}, ], }, {"role": "user", "content": "tell me about coffee"}, ], } with pytest.raises(HTTPException) as exc_info: await guardrail.async_pre_call_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, call_type="anthropic_messages", cache=MagicMock(spec=DualCache), ) exception = exc_info.value assert exception.status_code == 400 detail = exception.detail assert isinstance(detail, dict) assert detail["error"] == "Violated guardrail policy" assert ( detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question. coffee guardrail applied " ) @pytest.mark.asyncio async def test_bedrock_guardrails_block_responses_api(): """ Test that guardrails block responses API requests containing 'coffee' and raise the expected exception. """ from fastapi import HTTPException # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() guardrail = BedrockGuardrail( guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", ) request_data = { "model": "gpt-4.1", "input": "Tell me a three sentence bedtime story about a unicorn drinking coffee", "stream": False, } with pytest.raises(HTTPException) as exc_info: await guardrail.async_pre_call_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, call_type="responses", cache=MagicMock(spec=DualCache), ) exception = exc_info.value assert exception.status_code == 400 detail = exception.detail assert isinstance(detail, dict) assert detail["error"] == "Violated guardrail policy" assert ( detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question. coffee guardrail applied " ) @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-5.5", "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-5.5", "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-5.5", 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-5.5", "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-5.5", "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-5.5", "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-5.5", "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-5.5", object="chat.completion", ) request_data = { "model": "gpt-5.5", "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-5.5", 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-5.5", 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-5.5", object="chat.completion.chunk", ), ] for chunk in chunks: yield chunk request_data = { "model": "gpt-5.5", "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-5.5", 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-5.5", 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-5.5", 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-5.5", "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, logging_event_type=None, **kwargs, ): bedrock_calls.append( { "source": source, "messages": messages, "response": response, "request_data": request_data, "logging_event_type": logging_event_type, } ) # 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 # Note: When event_hook is None (default), the guardrail is considered enabled for all hooks. # In post_call, INPUT validation is skipped if pre_call/during_call is already enabled # to avoid redundant validation. Since event_hook=None means all hooks are enabled, # only OUTPUT validation should be performed in post_call. assert ( len(bedrock_calls) == 1 ), f"Expected 1 bedrock call (OUTPUT only), got {len(bedrock_calls)}" # Verify the OUTPUT call output_call = bedrock_calls[0] assert output_call["source"] == "OUTPUT" assert output_call["response"] is not None # OUTPUT forwards the request messages so contextual grounding can pull # grounding_source/query blocks from them even on streamed responses. A # plain-text (non-grounding) request still yields the single-block payload. assert output_call["messages"] == request_data["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]} " "(INPUT validation skipped due to event_hook=None implying pre_call/during_call enabled)" ) 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-5.5", "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-5.5", "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-5.5", "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 - 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", 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 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", ) assert exc_info.value.message == "I can't provide that information." @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-5.5", 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-5.5", 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-5.5", "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. 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", 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 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 async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): """Test that async_post_call_success_hook skips when there's no output text""" 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" ) # Create a ModelResponse with tool calls (no text content) # This simulates a response where the LLM is making a tool call mock_response = litellm.ModelResponse( id="test-id", choices=[ litellm.Choices( index=0, message=litellm.Message( role="assistant", content=None, # No text content tool_calls=[ litellm.utils.ChatCompletionMessageToolCall( id="tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", function=litellm.utils.Function( name="top_song", arguments='{"sign": "WZPZ"}' ), type="function", ) ], ), finish_reason="tool_calls", ) ], created=1234567890, model="gpt-5.5", object="chat.completion", ) data = { "model": "gpt-5.5", "messages": [ {"role": "user", "content": "Hello"}, ], } mock_user_api_key_dict = UserAPIKeyAuth() result = await guardrail.async_post_call_success_hook( data=data, response=mock_response, user_api_key_dict=mock_user_api_key_dict, ) # If no error is raised and result is None, then the test passes assert result is None print("✅ No output text in response test passed") @pytest.mark.asyncio async def test__redact_pii_matches_null_list_fields(): """Test that explicit null values from Bedrock API are handled correctly. The Bedrock API can return explicit JSON null for list fields like piiEntities, regexes, customWords, managedWordLists. This would cause TypeError: 'NoneType' object is not iterable if not handled. """ # Test 1: null piiEntities and regexes response_with_null_pii = { "action": "GUARDRAIL_INTERVENED", "assessments": [ { "sensitiveInformationPolicy": { "piiEntities": None, "regexes": None, } } ], } redacted = _redact_pii_matches(response_with_null_pii) assert redacted is not None assert ( redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] is None ) assert redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] is None # Test 2: null customWords and managedWordLists response_with_null_words = { "action": "GUARDRAIL_INTERVENED", "assessments": [ { "wordPolicy": { "customWords": None, "managedWordLists": None, } } ], } redacted = _redact_pii_matches(response_with_null_words) assert redacted is not None assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None # Test 3: null assessments at top level response_with_null_assessments = { "action": "GUARDRAIL_INTERVENED", "assessments": None, } redacted = _redact_pii_matches(response_with_null_assessments) assert redacted is not None @pytest.mark.asyncio async def test__redact_pii_matches_malformed_response(): """Test _redact_pii_matches with malformed response (should not crash)""" # Test with completely malformed response malformed_response = { "action": "GUARDRAIL_INTERVENED", "assessments": "not_a_list", } redacted_response = _redact_pii_matches(malformed_response) assert redacted_response == malformed_response # Test with missing keys missing_keys_response = { "action": "GUARDRAIL_INTERVENED", } redacted_response = _redact_pii_matches(missing_keys_response) assert redacted_response == missing_keys_response @pytest.mark.asyncio async def test_should_raise_guardrail_blocked_exception_null_fields(): """Test that _should_raise_guardrail_blocked_exception handles null list fields. Validates the or [] null-safety pattern works for all policy fields in _should_raise_guardrail_blocked_exception. """ guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) # Test with null assessments response_null_assessments = { "action": "GUARDRAIL_INTERVENED", "assessments": None, } assert ( guardrail._should_raise_guardrail_blocked_exception(response_null_assessments) is False ) # Test with null topics in topicPolicy response_null_topics = { "action": "GUARDRAIL_INTERVENED", "assessments": [{"topicPolicy": {"topics": None}}], } assert ( guardrail._should_raise_guardrail_blocked_exception(response_null_topics) is False ) # Test with null filters in contentPolicy response_null_filters = { "action": "GUARDRAIL_INTERVENED", "assessments": [{"contentPolicy": {"filters": None}}], } assert ( guardrail._should_raise_guardrail_blocked_exception(response_null_filters) is False ) # Test with null customWords and managedWordLists in wordPolicy response_null_words = { "action": "GUARDRAIL_INTERVENED", "assessments": [ {"wordPolicy": {"customWords": None, "managedWordLists": None}} ], } assert ( guardrail._should_raise_guardrail_blocked_exception(response_null_words) is False ) # Test with null piiEntities and regexes in sensitiveInformationPolicy response_null_pii = { "action": "GUARDRAIL_INTERVENED", "assessments": [ {"sensitiveInformationPolicy": {"piiEntities": None, "regexes": None}} ], } assert ( guardrail._should_raise_guardrail_blocked_exception(response_null_pii) is False ) # Test with null filters in contextualGroundingPolicy response_null_grounding = { "action": "GUARDRAIL_INTERVENED", "assessments": [{"contextualGroundingPolicy": {"filters": None}}], } assert ( guardrail._should_raise_guardrail_blocked_exception(response_null_grounding) is False )