scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-15 23:31:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 165
litellm/tests/guardrails_tests/test_bedrock_guardrails.py
yucheng-berri 6eed38bcfb
fix(guardrails/bedrock): honor disable_exception_on_block by raising ModifyResponseException (#32289) 
* fix(guardrails/bedrock): honor disable_exception_on_block by raising ModifyResponseException

The Bedrock-specific GuardrailInterventionNormalStringError predates the
unified guardrails refactor and no proxy code path handles it, so a block
with the flag set surfaced as an uncaught Exception -> HTTP 500 in pre_call
mode and was silently discarded in during_call mode (model call proceeded
in the parallel asyncio.gather; the block hook's data["mock_response"]
mutation happened after route_request had already unpacked kwargs).

Convert the block to ModifyResponseException at the raise site inside
make_bedrock_api_request. That exception is the industry-standard proxy
contract already caught in proxy_server, anthropic_endpoints, response_api
_endpoints, and pass_through_endpoints; it turns into a 200 response with
finish_reason=content_filter and the block message as content, which is
exactly what the flag was documented to yield. Post-call blocks attach
the LLM response to original_response so the synthetic reply reports the
upstream call's real token usage instead of zero.

Deletes the now-orphaned GuardrailInterventionNormalStringError class and
the dead create_guardrail_blocked_response / mock_response plumbing in the
Bedrock hooks; updates the existing tests that had locked in the buggy
contract.

Resolves LIT-4186

* chore(guardrails/bedrock): drop dead str branch in _update_messages_with_updated_bedrock_guardrail_response

Follow-up to the disable_exception_on_block fix. That method used to
receive either a BedrockGuardrailResponse or a plain string (the block
message, when the flag was set). Now that a block always raises
ModifyResponseException before this method runs, the string branch is
unreachable; tighten the type to BedrockGuardrailResponse and delete
the guard.

* fix(guardrails/bedrock): streaming post_call block yields synthetic stream instead of surfacing as SSE 500

Regression from the LIT-4186 refactor: pre-refactor, the streaming
post_call iterator caught GuardrailInterventionNormalStringError locally
and replaced the assembled response with a synthetic content-filter
message, then re-emitted it as chunks via MockResponseIterator. After
the refactor the exception was re-raised as ModifyResponseException,
which async_streaming_data_generator serializes as a proxy 500 error
frame because the SSE response headers are already flushed by the time
the block fires.

Non-streaming paths still let ModifyResponseException propagate to the
endpoint handler (which converts it into a 200). Streaming can't do
that, so keep the local synthesis: on the exception, rebind the
assembled response to a ModelResponse whose single choice carries the
block message as content and finish_reason=content_filter, and let the
downstream MockResponseIterator emit it as chunks. Same shape a
non-streaming block produces.

Adds a mapped-file regression test that mutation-kills the raise
behavior and locks in the synthetic-stream contract.

* fix(guardrails/bedrock): preserve upstream usage on streaming post_call block

Non-streaming post_call blocks report the upstream LLM call's real
token usage via ModifyResponseException.original_response, which the
endpoint handler unwraps through _blocked_response_usage. Streaming
post_call synthesizes its own ModelResponse locally (the exception
can't escape the SSE generator), and previously left .usage unset,
so the client saw accurate billing on non-streaming blocks and zero
on streaming blocks -- silent revenue leak.

Copy the assembled response's .usage onto the synthetic block
response before yielding. Pre-refactor code had the same gap
(create_guardrail_blocked_response never set usage); this is a net
improvement, not a regression fix.
2026-07-09 13:18:51 -07:00

1755 lines
60 KiB
Python
Raw Permalink Blame History

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
)
Reference in a new issue View git blame Copy permalink