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_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-3-5-sonnet-20240620", "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-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 # 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 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]} " "(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-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}" ) @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-4o", object="chat.completion", ) data = { "model": "gpt-4o", "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")