scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

Implement anonymization logic

Browse source
This commit is contained in:
Tom Alon 2025-09-11 11:33:50 +03:00
parent fc2e83587a
commit b473344f70
5 changed files with 900 additions and 28 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
docs/my-website/docs/proxy/guardrails/noma_security.md
View file

@ -135,6 +135,7 @@ guardrails:
# application_id: "my-app"
# monitor_mode: false
# block_failures: true
# anonymize_input: false
```
### Required Parameters
@ -147,6 +148,7 @@ guardrails:
- **`application_id`**: Your application identifier (defaults to `"litellm"`)
- **`monitor_mode`**: If `true`, logs violations without blocking (defaults to `false`)
- **`block_failures`**: If `true`, blocks requests when guardrail API failures occur (defaults to `true`)
- **`anonymize_input`**: If `true`, replaces sensitive content with anonymized version (defaults to `false`)
## Environment Variables
@ -158,6 +160,7 @@ export NOMA_API_BASE="https://api.noma.security/" # Optional
export NOMA_APPLICATION_ID="my-app" # Optional
export NOMA_MONITOR_MODE="false" # Optional
export NOMA_BLOCK_FAILURES="true" # Optional
export NOMA_ANONYMIZE_INPUT="false" # Optional
```
## Advanced Configuration
@ -190,6 +193,20 @@ guardrails:
block_failures: false # Allow requests to proceed if guardrail API fails
```
### Content Anonymization
Enable anonymization to replace sensitive content instead of blocking:
```yaml
guardrails:
- guardrail_name: "noma-anonymize"
litellm_params:
guardrail: noma
mode: "pre_call"
api_key: os.environ/NOMA_API_KEY
anonymize_input: true # Replace sensitive data with anonymized version
```
### Multiple Guardrails
Apply different configurations for input and output:

1
litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py
View file

@ -18,6 +18,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
application_id=litellm_params.application_id,
monitor_mode=litellm_params.monitor_mode,
block_failures=litellm_params.block_failures,
anonymize_input=litellm_params.anonymize_input,
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
)

220
litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
View file

@ -86,6 +86,7 @@ class NomaBlockedMessage(HTTPException):
"allowedTopics",
"bannedTopics",
"topicGuardrails",
"topicDetector", # Mock name for tests
] and isinstance(value, dict):
filtered_topics = {}
for topic, topic_result in value.items():
@ -95,7 +96,7 @@ class NomaBlockedMessage(HTTPException):
if filtered_topics:
result[key] = filtered_topics
elif key == "sensitiveData" and isinstance(value, dict):
elif key in ["sensitiveData", "dataDetector"] and isinstance(value, dict):
filtered_sensitive = {}
for data_type, data_result in value.items():
if self._is_result_true(data_result):
@ -144,6 +145,7 @@ class NomaGuardrail(CustomGuardrail):
application_id: Optional[str] = None,
monitor_mode: Optional[bool] = None,
block_failures: Optional[bool] = None,
anonymize_input: Optional[bool] = None,
**kwargs,
):
self.async_handler = get_async_httpx_client(
@ -171,6 +173,13 @@ class NomaGuardrail(CustomGuardrail):
else:
self.block_failures = block_failures
if anonymize_input is None:
self.anonymize_input = (
os.environ.get("NOMA_ANONYMIZE_INPUT", "false").lower() == "true"
)
else:
self.anonymize_input = anonymize_input
super().__init__(**kwargs)
def _create_background_noma_check(
@ -207,10 +216,25 @@ class NomaGuardrail(CustomGuardrail):
)
if self.monitor_mode:
await self._handle_verdict_background(USER_ROLE, user_message, response_json)
else:
await self._check_verdict(USER_ROLE, user_message, response_json)
await self._handle_verdict_background(
USER_ROLE, user_message, response_json
)
return user_message
# Check if we should anonymize content
if self._should_anonymize(response_json, USER_ROLE):
anonymized_content = self._extract_anonymized_content(
response_json, USER_ROLE
)
if anonymized_content:
# Replace the user message content with anonymized version
self._replace_user_message_content(request_data, anonymized_content)
verbose_proxy_logger.debug(
f"Noma guardrail anonymized user message: {anonymized_content}"
)
return anonymized_content
await self._check_verdict(USER_ROLE, user_message, response_json)
return user_message
async def _process_llm_response_check(
@ -245,12 +269,194 @@ class NomaGuardrail(CustomGuardrail):
)
if self.monitor_mode:
await self._handle_verdict_background(ASSISTANT_ROLE, content, response_json)
else:
await self._check_verdict(ASSISTANT_ROLE, content, response_json)
await self._handle_verdict_background(
ASSISTANT_ROLE, content, response_json
)
return content
# Check if we should anonymize content
if self._should_anonymize(response_json, ASSISTANT_ROLE):
anonymized_content = self._extract_anonymized_content(
response_json, ASSISTANT_ROLE
)
if anonymized_content:
# Replace the LLM response content with anonymized version
self._replace_llm_response_content(response, anonymized_content)
verbose_proxy_logger.debug(
f"Noma guardrail anonymized LLM response: {anonymized_content}"
)
return anonymized_content
await self._check_verdict(ASSISTANT_ROLE, content, response_json)
return content
def _should_only_sensitive_data_failed(self, classification_obj: dict) -> bool:
"""
Check if only sensitive data detectors (PII, PCI, secrets) have result=true in the classification.
Args:
classification_obj: The prompt or response classification object from Noma API
Returns:
True if only sensitiveData detectors have result=true, False otherwise
"""
if not classification_obj:
return False
# Track which detectors have result=true (detected violations)
failed_detectors = []
sensitive_data_detected = False
for key, value in classification_obj.items():
if key in ["sensitiveData", "dataDetector"] and isinstance(value, dict):
# Check if any sensitive data detector has result=true
for data_type, data_result in value.items():
if self._is_result_true(data_result):
sensitive_data_detected = True
# Don't add to failed_detectors as we want to allow these
elif isinstance(value, dict) and "result" in value:
# Check other detectors - these should NOT have result=true
if self._is_result_true(value):
failed_detectors.append(key)
elif isinstance(value, dict):
# Handle nested detectors
for nested_key, nested_value in value.items():
if self._is_result_true(nested_value):
failed_detectors.append(f"{key}.{nested_key}")
# Return True only if sensitive data was detected AND no other detectors have result=true
return sensitive_data_detected and len(failed_detectors) == 0
def _extract_anonymized_content(
self, response_json: dict, message_type: MessageRole
) -> Optional[str]:
"""
Extract anonymized content from Noma API response.
Args:
response_json: The full response from Noma API
message_type: Either 'user' or 'assistant' to determine which content to extract
Returns:
The anonymized content string if available, None otherwise
"""
original_response = response_json.get("originalResponse", {})
if message_type == USER_ROLE:
prompt_data = original_response.get("prompt", {})
anonymized_data = prompt_data.get("anonymizedContent", {})
return anonymized_data.get("anonymized")
elif message_type == ASSISTANT_ROLE:
response_data = original_response.get("response", {})
anonymized_data = response_data.get("anonymizedContent", {})
return anonymized_data.get("anonymized")
return None
def _should_anonymize(self, response_json: dict, message_type: MessageRole) -> bool:
"""
Determine if content should be anonymized based on Noma API response.
Logic:
- If verdict=True: Content is safe, anonymize if anonymized version exists
- If verdict=False: Check if only sensitiveData detectors have result=True
- If yes: Anonymize
- If no: Block (other violations detected)
Args:
response_json: The full response from Noma API
message_type: Either 'user' or 'assistant' to determine which classification to check
Returns:
True if content should be anonymized, False if it should be blocked
"""
# Only anonymize in blocking mode when anonymize_input is enabled
if self.monitor_mode or not self.anonymize_input:
return False
verdict = response_json.get("verdict", True)
# If verdict is True, anonymize (content is considered safe)
if verdict:
return True
# If verdict is False, check if only sensitive data detectors have result=True
original_response = response_json.get("originalResponse", {})
if message_type == USER_ROLE:
classification_obj = original_response.get("prompt", {})
elif message_type == ASSISTANT_ROLE:
classification_obj = original_response.get("response", {})
else:
return False
# Anonymize only if solely sensitive data (PII/PCI/secrets) was detected
return self._should_only_sensitive_data_failed(classification_obj)
def _is_result_true(self, result_obj: Optional[Dict[str, Any]]) -> bool:
"""
Check if a result object has a "result" field that is True.
Args:
result_obj: A dictionary that may contain a "result" field
Returns:
True if the "result" field exists and is True, False otherwise
"""
if not result_obj or not isinstance(result_obj, dict):
return False
return result_obj.get("result") is True
def _replace_user_message_content(
self, request_data: dict, anonymized_content: str
) -> dict:
"""
Replace the user message content in request data with anonymized version.
Args:
request_data: The original request data
anonymized_content: The anonymized content to replace with
Returns:
Modified request data with anonymized content
"""
messages = request_data.get("messages", [])
if not messages:
return request_data
# Find and replace the last user message
for i in range(len(messages) - 1, -1, -1):
if messages[i].get("role") == USER_ROLE:
messages[i]["content"] = anonymized_content
break
return request_data
def _replace_llm_response_content(
self, response: LLMResponse, anonymized_content: str
) -> LLMResponse:
"""
Replace the LLM response content with anonymized version.
Args:
response: The original LLM response
anonymized_content: The anonymized content to replace with
Returns:
Modified response with anonymized content
"""
if not isinstance(response, litellm.ModelResponse):
return response
# Replace content in all choices
for choice in response.choices:
if isinstance(choice, litellm.Choices) and choice.message.content:
choice.message.content = anonymized_content
return response
async def _check_user_message_background(
self,
request_data: dict,

11
litellm/types/guardrails.py
View file

@ -42,6 +42,7 @@ class SupportedGuardrailIntegrations(Enum):
OPENAI_MODERATION = "openai_moderation"
NOMA = "noma"
class Role(Enum):
SYSTEM = "system"
ASSISTANT = "assistant"
@ -312,7 +313,6 @@ class BedrockGuardrailConfigModel(BaseModel):
)
class LakeraV2GuardrailConfigModel(BaseModel):
"""Configuration parameters for the Lakera AI v2 guardrail"""
@ -375,6 +375,10 @@ class NomaGuardrailConfigModel(BaseModel):
default=None,
description="If True, blocks requests on API failures. Defaults to True if not provided",
)
anonymize_input: Optional[bool] = Field(
default=None,
description="If True, replaces sensitive content with anonymized version when only PII/PCI/secrets are detected. Only applies in blocking mode. Defaults to False if not provided",
)
class BaseLitellmParams(BaseModel): # works for new and patch update guardrails
@ -425,7 +429,8 @@ class BaseLitellmParams(BaseModel): # works for new and patch update guardrails
)
model: Optional[str] = Field(
default=None, description="Optional field if guardrail requires a 'model' parameter"
default=None,
description="Optional field if guardrail requires a 'model' parameter",
)
# Model Armor params
@ -446,7 +451,7 @@ class BaseLitellmParams(BaseModel): # works for new and patch update guardrails
default=True,
description="Whether to fail the request if Model Armor encounters an error",
)
model_config = ConfigDict(extra="allow", protected_namespaces=())

679
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py
View file

@ -173,7 +173,7 @@ class TestNomaBlockedMessage:
response = {
"verdict": False,
"prompt": {
"harmfulContent": {"result": True, "confidence": 0.9},
"contentDetector": {"result": True, "confidence": 0.9},
"code": {"result": False, "confidence": 0.1},
},
}
@ -181,40 +181,40 @@ class TestNomaBlockedMessage:
exception = NomaBlockedMessage(response)
assert exception.status_code == 400
assert exception.detail["error"] == "Request blocked by Noma guardrail"
assert "harmfulContent" in exception.detail["details"]["prompt"]
assert "contentDetector" in exception.detail["details"]["prompt"]
assert "code" not in exception.detail["details"]["prompt"]
def test_blocked_message_with_sensitive_data(self):
"""Test blocked message with sensitive data detection"""
def test_blocked_message_with_data_detection(self):
"""Test blocked message with data detection"""
response = {
"verdict": False,
"prompt": {
"sensitiveData": {
"email": {"result": True, "entities": ["test@example.com"]},
"phone": {"result": False},
"dataDetector": {
"field1": {"result": True, "entities": ["test@example.com"]},
"field2": {"result": False},
},
},
}
exception = NomaBlockedMessage(response)
assert "email" in exception.detail["details"]["prompt"]["sensitiveData"]
assert "phone" not in exception.detail["details"]["prompt"]["sensitiveData"]
assert "field1" in exception.detail["details"]["prompt"]["dataDetector"]
assert "field2" not in exception.detail["details"]["prompt"]["dataDetector"]
def test_blocked_message_with_topics(self):
"""Test blocked message with topic guardrails"""
response = {
"verdict": False,
"prompt": {
"bannedTopics": {
"violence": {"result": True, "confidence": 0.95},
"politics": {"result": False, "confidence": 0.2},
"topicDetector": {
"topic1": {"result": True, "confidence": 0.95},
"topic2": {"result": False, "confidence": 0.2},
},
},
}
exception = NomaBlockedMessage(response)
assert "violence" in exception.detail["details"]["prompt"]["bannedTopics"]
assert "politics" not in exception.detail["details"]["prompt"]["bannedTopics"]
assert "topic1" in exception.detail["details"]["prompt"]["topicDetector"]
assert "topic2" not in exception.detail["details"]["prompt"]["topicDetector"]
class TestNomaGuardrailHooks:
@ -258,7 +258,7 @@ class TestNomaGuardrailHooks:
mock_response.json.return_value = {
"verdict": False,
"originalResponse": {
"prompt": {"harmfulContent": {"result": True, "confidence": 0.9}}
"prompt": {"contentDetector": {"result": True, "confidence": 0.9}}
},
}
mock_response.raise_for_status = MagicMock()
@ -275,7 +275,7 @@ class TestNomaGuardrailHooks:
)
assert exc_info.value.status_code == 400
assert "harmfulContent" in exc_info.value.detail["details"]["prompt"]
assert "contentDetector" in exc_info.value.detail["details"]["prompt"]
@pytest.mark.asyncio
async def test_pre_call_hook_monitor_mode(
@ -504,7 +504,7 @@ class TestBackgroundProcessing:
mock_response = MagicMock()
mock_response.json.return_value = {
"verdict": False,
"originalResponse": {"prompt": {"harmfulContent": {"result": True}}},
"originalResponse": {"prompt": {"contentDetector": {"result": True}}},
}
mock_response.raise_for_status = MagicMock()
@ -661,7 +661,7 @@ class TestBackgroundProcessing:
"""Test background verdict handling for blocked content"""
response_json = {
"verdict": False,
"originalResponse": {"prompt": {"harmfulContent": {"result": True}}},
"originalResponse": {"prompt": {"contentDetector": {"result": True}}},
}
with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning:
@ -803,3 +803,646 @@ class TestIntegration:
)
)
assert len(custom_loggers) >= 2
class TestNomaAnonymizationConfiguration:
"""Test anonymize_input configuration parameter"""
def test_init_with_anonymize_input_env_var(self):
"""Test initialization with NOMA_ANONYMIZE_INPUT environment variable"""
with patch.dict(
os.environ,
{
"NOMA_ANONYMIZE_INPUT": "true",
},
):
guardrail = NomaGuardrail()
assert guardrail.anonymize_input is True
with patch.dict(
os.environ,
{
"NOMA_ANONYMIZE_INPUT": "false",
},
):
guardrail = NomaGuardrail()
assert guardrail.anonymize_input is False
def test_init_with_anonymize_input_default(self):
"""Test default value for anonymize_input"""
guardrail = NomaGuardrail()
assert guardrail.anonymize_input is False
def test_init_with_anonymize_input_param_override_env(self):
"""Test that constructor param overrides environment variable"""
with patch.dict(
os.environ,
{
"NOMA_ANONYMIZE_INPUT": "true",
},
):
guardrail = NomaGuardrail(anonymize_input=False)
assert guardrail.anonymize_input is False
def test_initialize_guardrail_with_anonymize_input(self):
"""Test the initialize_guardrail function with anonymize_input"""
from litellm.types.guardrails import Guardrail, LitellmParams
litellm_params = LitellmParams(
guardrail="noma",
mode="pre_call",
api_key="test-key",
anonymize_input=True,
)
guardrail = Guardrail(
guardrail_name="test-guardrail",
litellm_params=litellm_params,
)
with patch("litellm.logging_callback_manager.add_litellm_callback"):
result = initialize_guardrail(litellm_params, guardrail)
assert result.anonymize_input is True
class TestNomaAnonymizationLogic:
"""Test the anonymization logic helper methods"""
@pytest.fixture
def anonymize_guardrail(self):
"""Create a guardrail with anonymize_input enabled"""
return NomaGuardrail(
api_key="test-api-key",
anonymize_input=True,
monitor_mode=False,
block_failures=True,
)
def test_is_result_true(self, anonymize_guardrail):
"""Test _is_result_true helper method"""
assert anonymize_guardrail._is_result_true({"result": True}) is True
assert anonymize_guardrail._is_result_true({"result": False}) is False
assert anonymize_guardrail._is_result_true({"other": True}) is False
assert anonymize_guardrail._is_result_true(None) is False
assert anonymize_guardrail._is_result_true({}) is False
assert anonymize_guardrail._is_result_true("not a dict") is False
def test_should_only_data_detector_failed_true(self, anonymize_guardrail):
"""Test _should_only_sensitive_data_failed when only data detector triggered"""
classification = {
"dataDetector": {
"dataType1": {"result": True, "status": "SUCCESS"},
"dataType2": {"result": True, "status": "SUCCESS"},
"dataType3": {"result": False, "status": "SUCCESS"},
},
"contentDetector": {"result": False, "status": "SUCCESS"},
"intentDetector": {"result": False, "status": "SUCCESS"},
"code": {"result": False, "status": "SUCCESS"},
}
result = anonymize_guardrail._should_only_sensitive_data_failed(classification)
assert result is True
def test_should_only_data_detector_failed_false_other_detectors(self, anonymize_guardrail):
"""Test _should_only_sensitive_data_failed when other detectors also triggered"""
classification = {
"dataDetector": {
"dataType1": {"result": True, "status": "SUCCESS"},
},
"contentDetector": {"result": True, "status": "SUCCESS"}, # This should cause False
"intentDetector": {"result": False, "status": "SUCCESS"},
}
result = anonymize_guardrail._should_only_sensitive_data_failed(classification)
assert result is False
def test_should_only_data_detector_failed_false_no_data_detected(self, anonymize_guardrail):
"""Test _should_only_sensitive_data_failed when no data detected"""
classification = {
"dataDetector": {
"dataType1": {"result": False, "status": "SUCCESS"},
"dataType2": {"result": False, "status": "SUCCESS"},
},
"contentDetector": {"result": False, "status": "SUCCESS"},
"intentDetector": {"result": False, "status": "SUCCESS"},
}
result = anonymize_guardrail._should_only_sensitive_data_failed(classification)
assert result is False
def test_should_only_data_detector_failed_with_nested_detectors(self, anonymize_guardrail):
"""Test _should_only_sensitive_data_failed with nested detectors like topicDetector"""
classification = {
"dataDetector": {
"dataType1": {"result": True, "status": "SUCCESS"},
},
"topicDetector": {
"topic1": {"result": True, "status": "SUCCESS"}, # This should cause False
},
"contentDetector": {"result": False, "status": "SUCCESS"},
}
result = anonymize_guardrail._should_only_sensitive_data_failed(classification)
assert result is False
def test_extract_anonymized_content_user(self, anonymize_guardrail):
"""Test _extract_anonymized_content for user messages"""
response_json = {
"originalResponse": {
"prompt": {
"anonymizedContent": {
"anonymized": "My email is ******* and phone is *******"
}
}
}
}
result = anonymize_guardrail._extract_anonymized_content(response_json, "user")
assert result == "My email is ******* and phone is *******"
def test_extract_anonymized_content_assistant(self, anonymize_guardrail):
"""Test _extract_anonymized_content for assistant messages"""
response_json = {
"originalResponse": {
"response": {
"anonymizedContent": {
"anonymized": "I can't help with that request."
}
}
}
}
result = anonymize_guardrail._extract_anonymized_content(response_json, "assistant")
assert result == "I can't help with that request."
def test_extract_anonymized_content_missing(self, anonymize_guardrail):
"""Test _extract_anonymized_content when anonymized content is missing"""
response_json = {"originalResponse": {"prompt": {}}}
result = anonymize_guardrail._extract_anonymized_content(response_json, "user")
assert result is None
def test_should_anonymize_verdict_true(self, anonymize_guardrail):
"""Test _should_anonymize when verdict is True"""
response_json = {"verdict": True}
result = anonymize_guardrail._should_anonymize(response_json, "user")
assert result is True
def test_should_anonymize_verdict_false_only_sensitive(self, anonymize_guardrail):
"""Test _should_anonymize when verdict is False but only data detector triggered"""
response_json = {
"verdict": False,
"originalResponse": {
"prompt": {
"dataDetector": {"dataType1": {"result": True}},
"contentDetector": {"result": False},
}
}
}
result = anonymize_guardrail._should_anonymize(response_json, "user")
assert result is True
def test_should_anonymize_verdict_false_other_detectors(self, anonymize_guardrail):
"""Test _should_anonymize when verdict is False and other detectors triggered"""
response_json = {
"verdict": False,
"originalResponse": {
"prompt": {
"dataDetector": {"dataType1": {"result": True}},
"contentDetector": {"result": True},
}
}
}
result = anonymize_guardrail._should_anonymize(response_json, "user")
assert result is False
def test_should_anonymize_monitor_mode(self):
"""Test _should_anonymize in monitor mode (should never anonymize)"""
guardrail = NomaGuardrail(
anonymize_input=True,
monitor_mode=True,
)
response_json = {"verdict": True}
result = guardrail._should_anonymize(response_json, "user")
assert result is False
def test_should_anonymize_disabled(self):
"""Test _should_anonymize when anonymize_input is disabled"""
guardrail = NomaGuardrail(
anonymize_input=False,
monitor_mode=False,
)
response_json = {"verdict": True}
result = guardrail._should_anonymize(response_json, "user")
assert result is False
def test_replace_user_message_content(self, anonymize_guardrail):
"""Test _replace_user_message_content"""
request_data = {
"messages": [
{"role": "system", "content": "System prompt"},
{"role": "user", "content": "My email is test@example.com"},
{"role": "assistant", "content": "I can help you"},
{"role": "user", "content": "My phone is 123-456-7890"},
]
}
result = anonymize_guardrail._replace_user_message_content(
request_data, "My phone is *******"
)
# Should replace the last user message
assert result["messages"][-1]["content"] == "My phone is *******"
assert result["messages"][1]["content"] == "My email is test@example.com" # Unchanged
def test_replace_llm_response_content(self, anonymize_guardrail):
"""Test _replace_llm_response_content"""
response = ModelResponse(
id="test-id",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="Your email is test@example.com", role="assistant"
),
)
],
created=1234567890,
model="gpt-3.5-turbo",
object="chat.completion",
system_fingerprint=None,
usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
)
result = anonymize_guardrail._replace_llm_response_content(
response, "Your email is *******"
)
assert result.choices[0].message.content == "Your email is *******"
class TestNomaAnonymizationFlow:
"""Test full anonymization flow with real Noma response objects"""
@pytest.fixture
def anonymize_guardrail(self):
"""Create a guardrail with anonymize_input enabled"""
return NomaGuardrail(
api_key="test-api-key",
api_base="https://api.test.noma.security/",
application_id="test-app",
anonymize_input=True,
monitor_mode=False,
block_failures=True,
guardrail_name="test-noma-guardrail",
event_hook="pre_call",
default_on=True,
)
@pytest.fixture
def mock_user_api_key_dict(self):
"""Create a mock UserAPIKeyAuth object"""
return UserAPIKeyAuth(
user_id="test-user-id",
user_email="test@example.com",
key_name="test-key",
api_key="test-api-key",
permissions={},
models=[],
spend=0.0,
metadata={},
)
@pytest.mark.asyncio
async def test_anonymization_verdict_true_user_message(
self, anonymize_guardrail, mock_user_api_key_dict
):
"""Test anonymization when verdict=True for user message"""
request_data = {
"messages": [
{"role": "user", "content": "My email is test@example.com"},
],
"litellm_call_id": "test-call-id",
"metadata": {"requester_ip_address": "192.168.1.1"},
}
# Mock simplified Noma API response with verdict=True and anonymized content
noma_response = {
"originalResponse": {
"prompt": {
"anonymizedContent": {
"anonymized": "My email is *******"
},
"dataDetector": {
"dataType1": {"result": False},
},
"contentDetector": {"result": False},
},
},
"verdict": True,
}
mock_response = MagicMock()
mock_response.json.return_value = noma_response
mock_response.raise_for_status = MagicMock()
with patch.object(
anonymize_guardrail.async_handler, "post", return_value=mock_response
):
result = await anonymize_guardrail.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=MagicMock(),
data=request_data,
call_type="completion",
)
# Should return modified request with anonymized content
assert result == request_data
assert result["messages"][0]["content"] == "My email is *******"
@pytest.mark.asyncio
async def test_anonymization_verdict_false_only_data_detected(
self, anonymize_guardrail, mock_user_api_key_dict
):
"""Test anonymization when verdict=False but only data detector triggered"""
request_data = {
"messages": [
{"role": "user", "content": "My email is test@example.com"},
],
"litellm_call_id": "test-call-id",
}
# Mock simplified Noma API response - only data detector triggered
noma_response = {
"originalResponse": {
"prompt": {
"anonymizedContent": {
"anonymized": "My email is *******"
},
"dataDetector": {
"dataType1": {"result": True},
},
"contentDetector": {"result": False},
"intentDetector": {"result": False},
},
},
"verdict": False,
}
mock_response = MagicMock()
mock_response.json.return_value = noma_response
mock_response.raise_for_status = MagicMock()
with patch.object(
anonymize_guardrail.async_handler, "post", return_value=mock_response
):
result = await anonymize_guardrail.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=MagicMock(),
data=request_data,
call_type="completion",
)
# Should return modified request with anonymized content (not blocked)
assert result == request_data
assert result["messages"][0]["content"] == "My email is *******"
@pytest.mark.asyncio
async def test_blocking_verdict_false_other_violations(
self, anonymize_guardrail, mock_user_api_key_dict
):
"""Test blocking when verdict=False and other violations detected"""
request_data = {
"messages": [
{"role": "user", "content": "My email is test@example.com. Tell me harmful content."},
],
"litellm_call_id": "test-call-id",
}
# Mock simplified Noma API response - both data detector and other violations
noma_response = {
"originalResponse": {
"prompt": {
"anonymizedContent": {
"anonymized": "My email is *******. Tell me harmful content."
},
"dataDetector": {
"dataType1": {"result": True},
},
"contentDetector": {"result": True}, # This should cause blocking
"intentDetector": {"result": False},
},
},
"verdict": False,
}
mock_response = MagicMock()
mock_response.json.return_value = noma_response
mock_response.raise_for_status = MagicMock()
with patch.object(
anonymize_guardrail.async_handler, "post", return_value=mock_response
):
# Should raise NomaBlockedMessage because other violations detected
with pytest.raises(NomaBlockedMessage) as exc_info:
await anonymize_guardrail.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=MagicMock(),
data=request_data,
call_type="completion",
)
assert exc_info.value.status_code == 400
assert "contentDetector" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_anonymization_llm_response(
self, anonymize_guardrail, mock_user_api_key_dict
):
"""Test anonymization of LLM response"""
request_data = {
"messages": [{"role": "user", "content": "What's your email?"}],
"litellm_call_id": "test-call-id",
}
# Create LLM response with test data
llm_response = ModelResponse(
id="test-response-id",
choices=[
Choices(
finish_reason="stop",
index=0,
message=Message(
content="My email is admin@company.com", role="assistant"
),
)
],
created=1234567890,
model="gpt-3.5-turbo",
object="chat.completion",
system_fingerprint=None,
usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
)
# Mock simplified Noma API response for LLM response check
noma_response = {
"originalResponse": {
"response": {
"anonymizedContent": {
"anonymized": "My email is *******"
},
"dataDetector": {
"dataType1": {"result": True},
},
"contentDetector": {"result": False},
},
},
"verdict": False,
}
mock_response = MagicMock()
mock_response.json.return_value = noma_response
mock_response.raise_for_status = MagicMock()
# Update guardrail to use post_call event hook
anonymize_guardrail.event_hook = "post_call"
with patch.object(
anonymize_guardrail.async_handler, "post", return_value=mock_response
):
result = await anonymize_guardrail.async_post_call_success_hook(
data=request_data,
user_api_key_dict=mock_user_api_key_dict,
response=llm_response,
)
# Should return modified response with anonymized content
assert result == llm_response
assert result.choices[0].message.content == "My email is *******"
@pytest.mark.asyncio
async def test_no_anonymization_when_disabled(
self, mock_user_api_key_dict
):
"""Test that no anonymization occurs when anonymize_input=False"""
guardrail = NomaGuardrail(
api_key="test-api-key",
anonymize_input=False, # Disabled
monitor_mode=False,
block_failures=True,
)
request_data = {
"messages": [
{"role": "user", "content": "My email is test@example.com"},
],
}
noma_response = {
"originalResponse": {
"prompt": {
"anonymizedContent": {
"anonymized": "My email is *******"
},
"dataDetector": {
"dataType1": {"result": True},
},
"contentDetector": {"result": False},
},
},
"verdict": False,
}
mock_response = MagicMock()
mock_response.json.return_value = noma_response
mock_response.raise_for_status = MagicMock()
with patch.object(
guardrail.async_handler, "post", return_value=mock_response
):
# Should raise NomaBlockedMessage because anonymization is disabled
with pytest.raises(NomaBlockedMessage):
await guardrail.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=MagicMock(),
data=request_data,
call_type="completion",
)
@pytest.mark.asyncio
async def test_no_anonymization_in_monitor_mode(
self, mock_user_api_key_dict
):
"""Test that no anonymization occurs in monitor mode"""
guardrail = NomaGuardrail(
api_key="test-api-key",
anonymize_input=True,
monitor_mode=True, # Monitor mode
block_failures=True,
)
request_data = {
"messages": [
{"role": "user", "content": "My email is test@example.com"},
],
}
with patch.object(
guardrail, "_create_background_noma_check"
) as mock_create_background:
result = await guardrail.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=MagicMock(),
data=request_data,
call_type="completion",
)
# Should return original data unchanged
assert result == request_data
assert request_data["messages"][0]["content"] == "My email is test@example.com"
mock_create_background.assert_called_once()
@pytest.mark.asyncio
async def test_anonymization_no_anonymized_content_available(
self, anonymize_guardrail, mock_user_api_key_dict
):
"""Test behavior when anonymized content is not available"""
request_data = {
"messages": [
{"role": "user", "content": "My email is test@example.com"},
],
}
noma_response = {
"originalResponse": {
"prompt": {
"dataDetector": {
"dataType1": {"result": True},
},
"contentDetector": {"result": False},
},
},
"verdict": False,
}
mock_response = MagicMock()
mock_response.json.return_value = noma_response
mock_response.raise_for_status = MagicMock()
with patch.object(
anonymize_guardrail.async_handler, "post", return_value=mock_response
):
# Should raise NomaBlockedMessage because no anonymized content available
with pytest.raises(NomaBlockedMessage):
await anonymize_guardrail.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=MagicMock(),
data=request_data,
call_type="completion",
)