diff --git a/docs/my-website/docs/proxy/guardrails/noma_security.md b/docs/my-website/docs/proxy/guardrails/noma_security.md index 3a50841d65e..4aebb29eb57 100644 --- a/docs/my-website/docs/proxy/guardrails/noma_security.md +++ b/docs/my-website/docs/proxy/guardrails/noma_security.md @@ -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: diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py index dc3e4d9768e..4a96219b0d6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py @@ -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, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 1a1ed2acb1f..2601a73927d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -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, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index f31f304bda9..c35b0ff55a0 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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=()) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py index f1e91db7d59..895ad40faa2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py @@ -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", + )