mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
1758 lines
60 KiB
Python
1758 lines
60 KiB
Python
import io, asyncio
|
|
import pytest
|
|
|
|
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 fastapi import HTTPException
|
|
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()
|
|
|
|
async def _stream_through_guardrail():
|
|
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)
|
|
|
|
with pytest.raises(HTTPException):
|
|
await _stream_through_guardrail()
|
|
|
|
|
|
@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
|
|
async def _drain():
|
|
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,
|
|
)
|
|
)
|
|
|
|
async for chunk in result_generator:
|
|
pass
|
|
|
|
with pytest.raises(HTTPException):
|
|
await _drain()
|
|
|
|
# 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
|
|
)
|