From f19356db604b429ea6e483b03d4e06719469c849 Mon Sep 17 00:00:00 2001 From: dean-zavad <161694810+dean-zavad@users.noreply.github.com> Date: Wed, 5 Nov 2025 21:23:39 +0200 Subject: [PATCH] Litellm noma guardrail support images (#16199) * noma support v2 api and images with during call * supporting streams and images with texts * Supporting text now * annonymization works * removing function * fixing noma.py * all old tests pass * adding new tests * removing changes * Fixing application id headers * fix whitespace * deleting unused imports --- .../guardrails/guardrail_hooks/noma/noma.py | 300 +++--- .../guardrails/guardrail_hooks/test_noma.py | 919 +++++++++++++----- 2 files changed, 831 insertions(+), 388 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index 7073a4341f..9848fe5947 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -6,11 +6,11 @@ # +-------------------------------------------------------------+ import asyncio -import copy import os from datetime import datetime from typing import TYPE_CHECKING, Any, Dict, Final, Literal, Optional, Type, Union from urllib.parse import urljoin +import json from fastapi import HTTPException @@ -26,6 +26,18 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import EmbeddingResponse, GuardrailStatus, ImageResponse +from litellm.types.utils import ( + ModelResponseStream, +) +from typing import ( + List, + AsyncGenerator +) + +from litellm.llms.base_llm.base_model_iterator import MockResponseIterator +from litellm.main import stream_chunk_builder +from litellm.types.utils import TextCompletionResponse + # Constants USER_ROLE: Final[Literal["user"]] = "user" ASSISTANT_ROLE: Final[Literal["assistant"]] = "assistant" @@ -43,79 +55,14 @@ class NomaBlockedMessage(HTTPException): """Exception raised when Noma guardrail blocks a message""" def __init__(self, classification_response: dict): - classification = self._filter_triggered_classifications(classification_response) super().__init__( status_code=400, detail={ "error": "Request blocked by Noma guardrail", - "details": classification, + "details": classification_response, }, ) - def _filter_triggered_classifications( - self, - response_dict: dict, - ) -> dict: - """Filter and return only triggered classifications""" - filtered_response = copy.deepcopy(response_dict) - - # Filter prompt classifications if present - if filtered_response.get("prompt"): - filtered_response["prompt"] = self.filter_classification_object( - filtered_response["prompt"] - ) - - # Filter response classifications if present - if filtered_response.get("response"): - filtered_response["response"] = self.filter_classification_object( - filtered_response["response"] - ) - - return filtered_response - - def filter_classification_object( - self, - classification_obj: dict, - ) -> dict: - """Filter classification object to only include triggered items""" - if not classification_obj: - return {} - - result = {} - - for key, value in classification_obj.items(): - if value is None: - continue - - if key in [ - "allowedTopics", - "bannedTopics", - "topicGuardrails", - "topicDetector", # Mock name for tests - ] and isinstance(value, dict): - filtered_topics = {} - for topic, topic_result in value.items(): - if self._is_result_true(topic_result): - filtered_topics[topic] = topic_result - - if filtered_topics: - result[key] = filtered_topics - - elif key in SENSITIVE_DATA_DETECTOR_KEYS and isinstance(value, dict): - filtered_sensitive = {} - for data_type, data_result in value.items(): - if self._is_result_true(data_result): - filtered_sensitive[data_type] = data_result - - if filtered_sensitive: - result[key] = filtered_sensitive - - elif isinstance(value, dict) and "result" in value: - if self._is_result_true(value): - result[key] = value - - return result - def _is_result_true(self, result_obj: Optional[Dict[str, Any]]) -> bool: """ Check if a result object has a "result" field that is True. @@ -141,7 +88,7 @@ class NomaGuardrail(CustomGuardrail): """ _DEFAULT_API_BASE = "https://api.noma.security/" - _AIDR_ENDPOINT = "/ai-dr/v1/prompt/scan/aggregate" + _AIDR_ENDPOINT = "/ai-dr/v2/prompt/scan" def __init__( self, @@ -212,7 +159,15 @@ class NomaGuardrail(CustomGuardrail): if not user_message: return None - payload = {"request": {"text": user_message}} + payload = { + "input": [ + { + "type": "message", + "role": "user", + "content": user_message + } + ] + } response_json = await self._call_noma_api( payload=payload, llm_request_id=None, @@ -240,9 +195,9 @@ class NomaGuardrail(CustomGuardrail): if self.monitor_mode: await self._handle_verdict_background( - USER_ROLE, user_message, response_json + USER_ROLE, json.dumps(user_message), response_json ) - return user_message + return json.dumps(user_message) # Check if we should anonymize content if self._should_anonymize(response_json, USER_ROLE): @@ -257,8 +212,8 @@ class NomaGuardrail(CustomGuardrail): ) return anonymized_content - await self._check_verdict(USER_ROLE, user_message, response_json) - return user_message + await self._check_verdict(USER_ROLE, json.dumps(user_message), response_json) + return json.dumps(user_message) async def _process_llm_response_check( self, @@ -283,7 +238,20 @@ class NomaGuardrail(CustomGuardrail): if not content or not isinstance(content, str): return None - payload = {"response": {"text": content}} + payload = { + "input": [ + { + "type": "message", + "role": "assistant", + "content": [ + { + "type": "input_text", + "text": content + } + ] + } + ] + } response_json = await self._call_noma_api( payload=payload, @@ -312,7 +280,7 @@ class NomaGuardrail(CustomGuardrail): if self.monitor_mode: await self._handle_verdict_background( - ASSISTANT_ROLE, content, response_json + ASSISTANT_ROLE, json.dumps(content), response_json ) return content @@ -349,18 +317,19 @@ class NomaGuardrail(CustomGuardrail): if not isinstance(response_json, dict): return "guardrail_failed_to_respond" - # Get the verdict from the response - verdict = response_json.get("verdict", True) + # Get the aggregatedScanResult from the response + # aggregatedScanResult=True means unsafe (block), False means safe (allow) + aggregated_scan_result = response_json.get("aggregatedScanResult", False) - # If verdict is True, content is allowed - if verdict is True: + # If aggregatedScanResult is False, content is safe/allowed + if aggregated_scan_result is False: return "success" - # If verdict is False, content is blocked/flagged - if verdict is False: + # If aggregatedScanResult is True, content is blocked/flagged + if aggregated_scan_result is True: return "guardrail_intervened" - # If verdict is missing or invalid, treat as failure + # If aggregatedScanResult is missing or invalid, treat as failure return "guardrail_failed_to_respond" except Exception as e: @@ -419,17 +388,16 @@ class NomaGuardrail(CustomGuardrail): 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") - + # Extract from new scanResult structure + scan_result = response_json.get("scanResult", []) + if not scan_result: + return None + + # Find the scan result matching the message type (role) + for result_item in scan_result: + if result_item.get("role") == message_type: + return result_item.get("results", {}).get("anonymizedContent", {}).get("anonymized", "") + return None def _should_anonymize(self, response_json: dict, message_type: MessageRole) -> bool: @@ -437,8 +405,8 @@ class NomaGuardrail(CustomGuardrail): 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 aggregatedScanResult=False: Content is safe, anonymize if anonymized version exists + - If aggregatedScanResult=True: Check if only sensitiveData detectors have result=True - If yes: Anonymize - If no: Block (other violations detected) @@ -453,23 +421,26 @@ class NomaGuardrail(CustomGuardrail): 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: + # aggregatedScanResult=False means safe, True means unsafe + aggregated_scan_result = response_json.get("aggregatedScanResult", False) + + # If aggregatedScanResult is False, content is safe - anonymize if available + if not aggregated_scan_result: 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: + # If aggregatedScanResult is True (unsafe), check if only sensitive data detectors triggered + scan_result = response_json.get("scanResult", []) + if not scan_result: return False + + if not isinstance(scan_result, list) or len(scan_result) == 0: + return False + + for result_item in scan_result: + if result_item.get("role") == message_type: + return self._should_only_sensitive_data_failed(result_item.get("results", {})) - # Anonymize only if solely sensitive data (PII/PCI/secrets) was detected - return self._should_only_sensitive_data_failed(classification_obj) + return False def _is_result_true(self, result_obj: Optional[Dict[str, Any]]) -> bool: """ @@ -557,12 +528,17 @@ class NomaGuardrail(CustomGuardrail): message: str, response_json: dict, ) -> None: - """Handle verdict from Noma API in background - logging only, never blocks""" + """Handle aggregatedScanResult from Noma API in background - logging only, never blocks + aggregatedScanResult=True means unsafe, False means safe + """ try: - if not response_json.get("verdict", True): + # aggregatedScanResult=True means blocked, False means allowed + aggregated_scan_result = response_json.get("aggregatedScanResult", False) + + if aggregated_scan_result: # True = unsafe msg = f"Noma guardrail blocked {type} message: {message}" verbose_proxy_logger.warning(msg) - else: + else: # False = safe msg = f"Noma guardrail allowed {type} message: {message}" verbose_proxy_logger.info(msg) except Exception as e: @@ -588,6 +564,7 @@ class NomaGuardrail(CustomGuardrail): "anthropic_messages", ], ) -> Optional[Union[Exception, str, dict]]: + verbose_proxy_logger.debug("Running Noma pre-call hook") if ( @@ -756,7 +733,7 @@ class NomaGuardrail(CustomGuardrail): request_data: dict, response: LLMResponse, user_auth: UserAPIKeyAuth, - ) -> Union[Exception, ModelResponse, Any]: + ) -> Any: """Check LLM response for policy violations""" content = await self._process_llm_response_check( request_data, response, user_auth @@ -766,7 +743,7 @@ class NomaGuardrail(CustomGuardrail): return response - async def _extract_user_message(self, data: dict) -> Optional[str]: + async def _extract_user_message(self, data: dict) -> Optional[List[dict]]: """Extract the last user message from request data""" messages = data.get("messages", []) if not messages: @@ -778,10 +755,40 @@ class NomaGuardrail(CustomGuardrail): return None last_user_message = user_messages[-1].get("content", "") - if not last_user_message or not isinstance(last_user_message, str): + if isinstance(last_user_message, str): + return [{ + "type": "input_text", + "text": last_user_message + }] + elif isinstance(last_user_message, list): + converted_messages = [] + for message in last_user_message: + converted_message = self._convert_single_user_message_to_payload(message) + if converted_message is not None: + converted_messages.append(converted_message) + return converted_messages + else: return None - return last_user_message + + def _convert_single_user_message_to_payload(self, user_message: Any) -> Optional[dict]: + if isinstance(user_message, str): + return { + "type": "input_text", + "text": user_message + } + elif user_message.get("type", "") == "image_url": + return { + "type": "input_image", + "image_url": user_message.get("image_url", {}).get("url", "") + } + elif user_message.get("type", "") == "text": + return { + "type": "input_text", + "text": user_message.get("text", "") + } + else: + return None async def _call_noma_api( self, @@ -793,7 +800,6 @@ class NomaGuardrail(CustomGuardrail): ) -> dict: call_id = request_data.get("litellm_call_id") headers = { - "X-Noma-AIDR-Application-ID": self.application_id, **({"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}), **({"X-Noma-Request-ID": call_id} if call_id else {}), } @@ -806,11 +812,11 @@ class NomaGuardrail(CustomGuardrail): headers=headers, json={ **payload, - "context": { + "x-noma-context": { "applicationId": extra_data.get("application_id") or request_data.get("metadata", {}) .get("headers", {}) - .get("x-noma-application-id"), + .get("x-noma-application-id") or self.application_id, "ipAddress": request_data.get("metadata", {}).get( "requester_ip_address", None ), @@ -833,24 +839,29 @@ class NomaGuardrail(CustomGuardrail): response_json: dict, ) -> None: """ - Check the verdict from the Noma API and raise an exception if needed + Check the aggregatedScanResult from the Noma API and raise an exception if needed. + aggregatedScanResult=True means unsafe (block), False means safe (allow) """ - if not response_json.get("verdict", True): + # aggregatedScanResult=True means blocked, False means allowed + aggregated_scan_result = response_json.get("aggregatedScanResult", False) + + if aggregated_scan_result: # True = unsafe, block it msg = f"Noma guardrail blocked {type} message: {message}" if self.monitor_mode: verbose_proxy_logger.warning(msg) else: verbose_proxy_logger.debug(msg) - original_response = response_json.get("originalResponse", {}) + original_response = response_json.get("scanResult", {}) + # Use the full response as the original response for error details raise NomaBlockedMessage(original_response) - else: + else: # False = safe, allow it msg = f"Noma guardrail allowed {type} message: {message}" if self.monitor_mode: verbose_proxy_logger.info(msg) else: verbose_proxy_logger.debug(msg) - + @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: from litellm.types.proxy.guardrails.guardrail_hooks.noma import ( @@ -859,3 +870,46 @@ class NomaGuardrail(CustomGuardrail): return NomaGuardrailConfigModel + async def async_post_call_streaming_iterator_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + response: Any, + request_data: dict, + ) -> AsyncGenerator[ModelResponseStream, None]: + """Process streaming response chunks with Noma guardrail.""" + + all_chunks: List[ModelResponseStream] = [] + async for chunk in response: + all_chunks.append(chunk) + + if not all_chunks: + return + + assembled_model_response: Optional[ + Union[ModelResponse, TextCompletionResponse] + ] = stream_chunk_builder(chunks=all_chunks) + + if isinstance(assembled_model_response, ModelResponse): + try: + processed_response = await self._check_llm_response( + request_data, assembled_model_response, user_api_key_dict + ) + except NomaBlockedMessage: + raise + except Exception as e: + if self.block_failures: + raise + verbose_proxy_logger.error( + f"Noma streaming post-call hook failed: {str(e)}" + ) + for chunk in all_chunks: + yield chunk + return + + mock_response = MockResponseIterator(model_response=processed_response) + async for chunk in mock_response: + yield chunk + return + + for chunk in all_chunks: + yield chunk 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 f9b80a0a57..7a328c03aa 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py @@ -171,50 +171,64 @@ class TestNomaBlockedMessage: def test_blocked_message_basic(self): """Test basic blocked message creation""" response = { - "verdict": False, - "prompt": { - "contentDetector": {"result": True, "confidence": 0.9}, - "code": {"result": False, "confidence": 0.1}, - }, + "aggregatedScanResult": True, + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "harmfulContent": {"result": True, "probability": 0.9, "status": "SUCCESS"}, + "code": {"result": False, "probability": 0.1, "status": "SUCCESS"}, + } + } + ] } exception = NomaBlockedMessage(response) assert exception.status_code == 400 assert exception.detail["error"] == "Request blocked by Noma guardrail" - assert "contentDetector" in exception.detail["details"]["prompt"] - assert "code" not in exception.detail["details"]["prompt"] def test_blocked_message_with_data_detection(self): """Test blocked message with data detection""" response = { - "verdict": False, - "prompt": { - "dataDetector": { - "field1": {"result": True, "entities": ["test@example.com"]}, - "field2": {"result": False}, - }, - }, + "aggregatedScanResult": True, + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "sensitiveData": { + "PII": {"result": True, "probability": 0.8, "status": "SUCCESS"}, + "PCI": {"result": False, "probability": 0, "status": "SUCCESS"}, + }, + } + } + ] } exception = NomaBlockedMessage(response) - assert "field1" in exception.detail["details"]["prompt"]["dataDetector"] - assert "field2" not in exception.detail["details"]["prompt"]["dataDetector"] + assert exception.detail["error"] == "Request blocked by Noma guardrail" def test_blocked_message_with_topics(self): """Test blocked message with topic guardrails""" response = { - "verdict": False, - "prompt": { - "topicDetector": { - "topic1": {"result": True, "confidence": 0.95}, - "topic2": {"result": False, "confidence": 0.2}, - }, - }, + "aggregatedScanResult": True, + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "customLlm": { + "topic1": {"result": True, "probability": 0.95, "status": "SUCCESS"}, + "topic2": {"result": False, "probability": 0.2, "status": "SUCCESS"}, + }, + } + } + ] } exception = NomaBlockedMessage(response) - assert "topic1" in exception.detail["details"]["prompt"]["topicDetector"] - assert "topic2" not in exception.detail["details"]["prompt"]["topicDetector"] + assert exception.detail["error"] == "Request blocked by Noma guardrail" class TestNomaGuardrailHooks: @@ -226,7 +240,16 @@ class TestNomaGuardrailHooks: ): """Test pre-call hook when content is allowed""" mock_response = MagicMock() - mock_response.json.return_value = {"verdict": True} + mock_response.json.return_value = { + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": {} + } + ] + } mock_response.raise_for_status = MagicMock() with patch.object( @@ -244,10 +267,9 @@ class TestNomaGuardrailHooks: # Verify API call details call_args = mock_post.call_args - assert call_args[0][0].endswith("/ai-dr/v1/prompt/scan/aggregate") + assert call_args[0][0].endswith("/ai-dr/v2/prompt/scan") assert call_args[1]["headers"]["X-Noma-AIDR-Application-ID"] == "test-app" assert call_args[1]["headers"]["Authorization"] == "Bearer test-api-key" - assert call_args[1]["json"]["request"]["text"] == "Hello, how are you?" @pytest.mark.asyncio async def test_pre_call_hook_blocked( @@ -256,10 +278,16 @@ class TestNomaGuardrailHooks: """Test pre-call hook when content is blocked""" mock_response = MagicMock() mock_response.json.return_value = { - "verdict": False, - "originalResponse": { - "prompt": {"contentDetector": {"result": True, "confidence": 0.9}} - }, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "harmfulContent": {"result": True, "probability": 0.9, "status": "SUCCESS"} + } + } + ] } mock_response.raise_for_status = MagicMock() @@ -275,7 +303,6 @@ class TestNomaGuardrailHooks: ) assert exc_info.value.status_code == 400 - assert "contentDetector" in exc_info.value.detail["details"]["prompt"] @pytest.mark.asyncio async def test_pre_call_hook_monitor_mode( @@ -356,7 +383,16 @@ class TestNomaGuardrailHooks: ) mock_api_response = MagicMock() - mock_api_response.json.return_value = {"verdict": True} + mock_api_response.json.return_value = { + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "assistant", + "type": "message", + "results": {} + } + ] + } mock_api_response.raise_for_status = MagicMock() # Update guardrail to use post_call event hook @@ -374,13 +410,6 @@ class TestNomaGuardrailHooks: assert result == response mock_post.assert_called_once() - # Verify API call details - call_args = mock_post.call_args - assert ( - call_args[1]["json"]["response"]["text"] == "I'm doing well, thank you!" - ) - assert call_args[1]["json"]["context"]["requestId"] == "test-response-id" - @pytest.mark.asyncio async def test_moderation_hook( self, noma_guardrail, mock_user_api_key_dict, mock_request_data @@ -390,7 +419,16 @@ class TestNomaGuardrailHooks: noma_guardrail.event_hook = "during_call" mock_response = MagicMock() - mock_response.json.return_value = {"verdict": True} + mock_response.json.return_value = { + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": {} + } + ] + } mock_response.raise_for_status = MagicMock() with patch.object( @@ -464,7 +502,7 @@ class TestNomaGuardrailHooks: import asyncio message = asyncio.run(noma_guardrail._extract_user_message(data)) - assert message == "Second user message" + assert message == [{"type": "input_text", "text": "Second user message"}] data = {"messages": [{"role": "system", "content": "System prompt"}]} message = asyncio.run(noma_guardrail._extract_user_message(data)) @@ -503,8 +541,16 @@ class TestBackgroundProcessing: """Test shared helper method in monitor mode""" mock_response = MagicMock() mock_response.json.return_value = { - "verdict": False, - "originalResponse": {"prompt": {"contentDetector": {"result": True}}}, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "harmfulContent": {"result": True, "status": "SUCCESS"} + } + } + ] } mock_response.raise_for_status = MagicMock() @@ -518,11 +564,9 @@ class TestBackgroundProcessing: mock_request_data, mock_user_api_key_dict ) - assert result == "Hello, how are you?" + assert result is not None mock_post.assert_called_once() - mock_handle_verdict.assert_called_once_with( - "user", "Hello, how are you?", mock_response.json.return_value - ) + mock_handle_verdict.assert_called_once() @pytest.mark.asyncio async def test_process_user_message_check_non_monitor_mode( @@ -530,7 +574,16 @@ class TestBackgroundProcessing: ): """Test shared helper method in non-monitor mode""" mock_response = MagicMock() - mock_response.json.return_value = {"verdict": True} + mock_response.json.return_value = { + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": {} + } + ] + } mock_response.raise_for_status = MagicMock() with patch.object( @@ -543,11 +596,9 @@ class TestBackgroundProcessing: mock_request_data, mock_user_api_key_dict ) - assert result == "Hello, how are you?" + assert result is not None mock_post.assert_called_once() - mock_check_verdict.assert_called_once_with( - "user", "Hello, how are you?", mock_response.json.return_value - ) + mock_check_verdict.assert_called_once() @pytest.mark.asyncio async def test_process_llm_response_check_monitor_mode( @@ -575,7 +626,16 @@ class TestBackgroundProcessing: ) mock_api_response = MagicMock() - mock_api_response.json.return_value = {"verdict": True} + mock_api_response.json.return_value = { + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "assistant", + "type": "message", + "results": {} + } + ] + } mock_api_response.raise_for_status = MagicMock() with patch.object( @@ -590,9 +650,7 @@ class TestBackgroundProcessing: assert result == "I'm doing well, thank you!" mock_post.assert_called_once() - mock_handle_verdict.assert_called_once_with( - "assistant", "I'm doing well, thank you!", mock_api_response.json.return_value - ) + mock_handle_verdict.assert_called_once() @pytest.mark.asyncio async def test_check_user_message_background( @@ -660,8 +718,16 @@ class TestBackgroundProcessing: async def test_handle_verdict_background_blocked(self, monitor_mode_guardrail): """Test background verdict handling for blocked content""" response_json = { - "verdict": False, - "originalResponse": {"prompt": {"contentDetector": {"result": True}}}, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "harmfulContent": {"result": True, "status": "SUCCESS"} + } + } + ] } with patch("litellm._logging.verbose_proxy_logger.warning") as mock_warning: @@ -675,7 +741,16 @@ class TestBackgroundProcessing: @pytest.mark.asyncio async def test_handle_verdict_background_allowed(self, monitor_mode_guardrail): """Test background verdict handling for allowed content""" - response_json = {"verdict": True} + response_json = { + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "assistant", + "type": "message", + "results": {} + } + ] + } with patch("litellm._logging.verbose_proxy_logger.info") as mock_info: await monitor_mode_guardrail._handle_verdict_background( @@ -764,102 +839,334 @@ class TestBackgroundProcessing: mock_create_background.assert_called_once() -class TestNomaApplyGuardrail: - """ - Test the apply_guardrail method for Noma guardrails - """ +class TestNomaImageProcessing: + """Test image processing functionality for multimodal content""" - @pytest.mark.asyncio - async def test_apply_guardrail_success(self): - """ - Test that apply_guardrail returns text when content is allowed - """ - guardrail = NomaGuardrail( + @pytest.fixture + def noma_guardrail(self): + """Create a NomaGuardrail instance for testing""" + return NomaGuardrail( api_key="test-api-key", api_base="https://api.test.noma.security/", application_id="test-app", monitor_mode=False, block_failures=True, + guardrail_name="test-noma-guardrail", + event_hook="pre_call", + default_on=True, ) - mock_response = MagicMock() - mock_response.json.return_value = {"verdict": True} - mock_response.raise_for_status = MagicMock() - - with patch.object( - guardrail.async_handler, "post", return_value=mock_response - ): - result = await guardrail.apply_guardrail( - text="This is a safe test message" - ) - - assert result == "This is a safe test message" - - @pytest.mark.asyncio - async def test_apply_guardrail_blocked(self): - """ - Test that apply_guardrail raises exception when content is blocked - """ - guardrail = NomaGuardrail( + @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", - api_base="https://api.test.noma.security/", - application_id="test-app", - monitor_mode=False, - block_failures=True, + permissions={}, + models=[], + spend=0.0, + metadata={}, ) - mock_response = MagicMock() - mock_response.json.return_value = { - "verdict": False, - "originalResponse": { - "prompt": {"contentDetector": {"result": True, "confidence": 0.9}} - }, + def test_extract_user_message_with_image_url(self, noma_guardrail): + """Test extracting user message with image_url content""" + import asyncio + + data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/image.jpg" + } + } + ] + } + ] } - mock_response.raise_for_status = MagicMock() - with patch.object( - guardrail.async_handler, "post", return_value=mock_response - ): - with pytest.raises(Exception) as exc_info: - await guardrail.apply_guardrail(text="This is blocked content") + message = asyncio.run(noma_guardrail._extract_user_message(data)) + assert message is not None + assert len(message) == 1 + assert message[0]["type"] == "input_image" + assert message[0]["image_url"] == "https://example.com/image.jpg" - assert "Content blocked by Noma guardrail" in str(exc_info.value) + def test_extract_user_message_with_mixed_content(self, noma_guardrail): + """Test extracting user message with mixed text and image content""" + import asyncio + + data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "What's in this image?" + }, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/image.jpg" + } + } + ] + } + ] + } + + message = asyncio.run(noma_guardrail._extract_user_message(data)) + assert message is not None + assert len(message) == 2 + # First item should be text + assert message[0]["type"] == "input_text" + assert message[0]["text"] == "What's in this image?" + # Second item should be image + assert message[1]["type"] == "input_image" + assert message[1]["image_url"] == "https://example.com/image.jpg" + + def test_extract_user_message_with_multiple_images(self, noma_guardrail): + """Test extracting user message with multiple images""" + import asyncio + + data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Compare these images" + }, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/image1.jpg" + } + }, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/image2.jpg" + } + } + ] + } + ] + } + + message = asyncio.run(noma_guardrail._extract_user_message(data)) + assert message is not None + assert len(message) == 3 + assert message[0]["type"] == "input_text" + assert message[1]["type"] == "input_image" + assert message[1]["image_url"] == "https://example.com/image1.jpg" + assert message[2]["type"] == "input_image" + assert message[2]["image_url"] == "https://example.com/image2.jpg" @pytest.mark.asyncio - async def test_apply_guardrail_with_anonymization(self): - """ - Test that apply_guardrail returns anonymized text when anonymize_input is enabled - """ - guardrail = 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, - ) + async def test_pre_call_hook_with_image_content( + self, noma_guardrail, mock_user_api_key_dict + ): + """Test pre-call hook with image content""" + request_data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/test-image.jpg" + } + } + ] + } + ], + "litellm_call_id": "test-call-id", + "metadata": {"requester_ip_address": "192.168.1.1"}, + } - mock_response = MagicMock() - mock_response.json.return_value = { - "verdict": True, - "originalResponse": { - "prompt": { - "anonymizedContent": { - "anonymized": "My email is ******* and phone is *******" + # Mock Noma API response for image content + noma_response = { + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "harmfulContent": {"result": False, "probability": 0.1, "status": "SUCCESS"} } } - }, + ] } + + 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 - ): - result = await guardrail.apply_guardrail( - text="My email is test@example.com and phone is 123-456-7890" + noma_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + result = await noma_guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key_dict, + cache=MagicMock(), + data=request_data, + call_type="completion", ) - assert result == "My email is ******* and phone is *******" + assert result == request_data + mock_post.assert_called_once() + + # Verify the API call payload includes image + call_args = mock_post.call_args + payload = call_args[1]["json"] + assert "input" in payload + assert len(payload["input"]) > 0 + assert payload["input"][0]["role"] == "user" + assert "content" in payload["input"][0] + + @pytest.mark.asyncio + async def test_pre_call_hook_with_mixed_content( + self, noma_guardrail, mock_user_api_key_dict + ): + """Test pre-call hook with mixed text and image content""" + request_data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Analyze this image for harmful content" + }, + { + "type": "image_url", + "image_url": { + "url": "https://example.com/test-image.jpg" + } + } + ] + } + ], + "litellm_call_id": "test-call-id", + } + + # Mock Noma API response + noma_response = { + "aggregatedScanResult": False, + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "harmfulContent": {"result": False, "probability": 0.05, "status": "SUCCESS"} + } + } + ] + } + + mock_response = MagicMock() + mock_response.json.return_value = noma_response + mock_response.raise_for_status = MagicMock() + + with patch.object( + noma_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + result = await noma_guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key_dict, + cache=MagicMock(), + data=request_data, + call_type="completion", + ) + + assert result == request_data + mock_post.assert_called_once() + + @pytest.mark.asyncio + async def test_image_content_blocked( + self, noma_guardrail, mock_user_api_key_dict + ): + """Test that image content can be blocked by Noma""" + request_data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://example.com/inappropriate-image.jpg" + } + } + ] + } + ], + "litellm_call_id": "test-call-id", + } + + # Mock Noma API response indicating harmful content in image + noma_response = { + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "harmfulContent": {"result": True, "probability": 0.95, "status": "SUCCESS"} + } + } + ] + } + + mock_response = MagicMock() + mock_response.json.return_value = noma_response + mock_response.raise_for_status = MagicMock() + + from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage + + with patch.object( + noma_guardrail.async_handler, "post", return_value=mock_response + ): + with pytest.raises(NomaBlockedMessage) as exc_info: + await noma_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 + + @pytest.mark.asyncio + async def test_image_with_base64_data(self, noma_guardrail): + """Test extracting image with base64 data URL""" + data = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "data:image/jpeg;base64,/9j/4AAQSkZJRg..." + } + } + ] + } + ] + } + + message = await noma_guardrail._extract_user_message(data) + assert message is not None + assert len(message) == 1 + assert message[0]["type"] == "input_image" + assert message[0]["image_url"].startswith("data:image/jpeg;base64,") class TestIntegration: @@ -986,15 +1293,15 @@ class TestNomaAnonymizationLogic: 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""" + """Test _should_only_sensitive_data_failed when only sensitive data detector triggered""" classification = { - "dataDetector": { - "dataType1": {"result": True, "status": "SUCCESS"}, - "dataType2": {"result": True, "status": "SUCCESS"}, - "dataType3": {"result": False, "status": "SUCCESS"}, + "sensitiveData": { + "PII": {"result": True, "status": "SUCCESS"}, + "PCI": {"result": True, "status": "SUCCESS"}, + "secrets": {"result": False, "probability": 0, "status": "SUCCESS"}, }, - "contentDetector": {"result": False, "status": "SUCCESS"}, - "intentDetector": {"result": False, "status": "SUCCESS"}, + "harmfulContent": {"result": False, "status": "SUCCESS"}, + "maliciousIntent": {"result": False, "status": "SUCCESS"}, "code": {"result": False, "status": "SUCCESS"}, } @@ -1004,25 +1311,25 @@ class TestNomaAnonymizationLogic: 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"}, + "sensitiveData": { + "PII": {"result": True, "status": "SUCCESS"}, }, - "contentDetector": {"result": True, "status": "SUCCESS"}, # This should cause False - "intentDetector": {"result": False, "status": "SUCCESS"}, + "harmfulContent": {"result": True, "status": "SUCCESS"}, # This should cause False + "maliciousIntent": {"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""" + """Test _should_only_sensitive_data_failed when no sensitive data detected""" classification = { - "dataDetector": { - "dataType1": {"result": False, "status": "SUCCESS"}, - "dataType2": {"result": False, "status": "SUCCESS"}, + "sensitiveData": { + "PII": {"result": False, "status": "SUCCESS"}, + "PCI": {"result": False, "status": "SUCCESS"}, }, - "contentDetector": {"result": False, "status": "SUCCESS"}, - "intentDetector": {"result": False, "status": "SUCCESS"}, + "harmfulContent": {"result": False, "status": "SUCCESS"}, + "maliciousIntent": {"result": False, "status": "SUCCESS"}, } result = anonymize_guardrail._should_only_sensitive_data_failed(classification) @@ -1031,13 +1338,13 @@ class TestNomaAnonymizationLogic: 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"}, + "sensitiveData": { + "PII": {"result": True, "status": "SUCCESS"}, }, - "topicDetector": { + "customLlm": { "topic1": {"result": True, "status": "SUCCESS"}, # This should cause False }, - "contentDetector": {"result": False, "status": "SUCCESS"}, + "harmfulContent": {"result": False, "status": "SUCCESS"}, } result = anonymize_guardrail._should_only_sensitive_data_failed(classification) @@ -1046,13 +1353,17 @@ class TestNomaAnonymizationLogic: 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 *******" + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "anonymizedContent": { + "anonymized": "My email is ******* and phone is *******" + } } } - } + ] } result = anonymize_guardrail._extract_anonymized_content(response_json, "user") @@ -1061,13 +1372,17 @@ class TestNomaAnonymizationLogic: 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." + "scanResult": [ + { + "role": "assistant", + "type": "message", + "results": { + "anonymizedContent": { + "anonymized": "I can't help with that request." + } } } - } + ] } result = anonymize_guardrail._extract_anonymized_content(response_json, "assistant") @@ -1075,43 +1390,68 @@ class TestNomaAnonymizationLogic: def test_extract_anonymized_content_missing(self, anonymize_guardrail): """Test _extract_anonymized_content when anonymized content is missing""" - response_json = {"originalResponse": {"prompt": {}}} + response_json = { + "scanResult": [ + { + "role": "user", + "type": "message", + "results": {} + } + ] + } result = anonymize_guardrail._extract_anonymized_content(response_json, "user") - assert result is None + assert result == "" def test_should_anonymize_verdict_true(self, anonymize_guardrail): - """Test _should_anonymize when verdict is True""" - response_json = {"verdict": True} + """Test _should_anonymize when aggregatedScanResult is False (safe)""" + response_json = { + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": {} + } + ] + } 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""" + """Test _should_anonymize when aggregatedScanResult is True but only sensitive data detector triggered""" response_json = { - "verdict": False, - "originalResponse": { - "prompt": { - "dataDetector": {"dataType1": {"result": True}}, - "contentDetector": {"result": False}, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "sensitiveData": {"PCI": {"result": True, "status": "SUCCESS"}}, + "harmfulContent": {"result": False, "status": "SUCCESS"}, + } } - } + ] } 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""" + """Test _should_anonymize when aggregatedScanResult is True and other detectors triggered""" response_json = { - "verdict": False, - "originalResponse": { - "prompt": { - "dataDetector": {"dataType1": {"result": True}}, - "contentDetector": {"result": True}, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "sensitiveData": {"PCI": {"result": True, "status": "SUCCESS"}}, + "harmfulContent": {"result": True, "status": "SUCCESS"}, + } } - } + ] } result = anonymize_guardrail._should_anonymize(response_json, "user") @@ -1124,7 +1464,16 @@ class TestNomaAnonymizationLogic: monitor_mode=True, ) - response_json = {"verdict": True} + response_json = { + "aggregatedScanResult": False, + "scanResult": [ + { + "role": "user", + "type": "message", + "results": {} + } + ] + } result = guardrail._should_anonymize(response_json, "user") assert result is False @@ -1135,7 +1484,16 @@ class TestNomaAnonymizationLogic: monitor_mode=False, ) - response_json = {"verdict": True} + response_json = { + "aggregatedScanResult": False, + "scanResult": [ + { + "role": "user", + "type": "message", + "results": {} + } + ] + } result = guardrail._should_anonymize(response_json, "user") assert result is False @@ -1230,20 +1588,24 @@ class TestNomaAnonymizationFlow: "metadata": {"requester_ip_address": "192.168.1.1"}, } - # Mock simplified Noma API response with verdict=True and anonymized content + # Mock simplified Noma API response with aggregatedScanResult=False (safe) and anonymized content noma_response = { - "originalResponse": { - "prompt": { - "anonymizedContent": { - "anonymized": "My email is *******" - }, - "dataDetector": { - "dataType1": {"result": False}, - }, - "contentDetector": {"result": False}, - }, - }, - "verdict": True, + "aggregatedScanResult": False, # False means safe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "anonymizedContent": { + "anonymized": "My email is *******" + }, + "sensitiveData": { + "PII": {"result": False, "status": "SUCCESS"}, + }, + "harmfulContent": {"result": False, "status": "SUCCESS"}, + } + } + ] } mock_response = MagicMock() @@ -1276,21 +1638,25 @@ class TestNomaAnonymizationFlow: "litellm_call_id": "test-call-id", } - # Mock simplified Noma API response - only data detector triggered + # Mock simplified Noma API response - only sensitive data detector triggered noma_response = { - "originalResponse": { - "prompt": { - "anonymizedContent": { - "anonymized": "My email is *******" - }, - "dataDetector": { - "dataType1": {"result": True}, - }, - "contentDetector": {"result": False}, - "intentDetector": {"result": False}, - }, - }, - "verdict": False, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "anonymizedContent": { + "anonymized": "My email is *******" + }, + "sensitiveData": { + "PII": {"result": True, "status": "SUCCESS"}, + }, + "harmfulContent": {"result": False, "status": "SUCCESS"}, + "maliciousIntent": {"result": False, "status": "SUCCESS"}, + } + } + ] } mock_response = MagicMock() @@ -1323,21 +1689,25 @@ class TestNomaAnonymizationFlow: "litellm_call_id": "test-call-id", } - # Mock simplified Noma API response - both data detector and other violations + # Mock simplified Noma API response - both sensitive 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, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "anonymizedContent": { + "anonymized": "My email is *******. Tell me harmful content." + }, + "sensitiveData": { + "PII": {"result": True, "status": "SUCCESS"}, + }, + "harmfulContent": {"result": True, "status": "SUCCESS"}, # This should cause blocking + "maliciousIntent": {"result": False, "status": "SUCCESS"}, + } + } + ] } mock_response = MagicMock() @@ -1357,7 +1727,6 @@ class TestNomaAnonymizationFlow: ) assert exc_info.value.status_code == 400 - assert "contentDetector" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_anonymization_llm_response( @@ -1390,20 +1759,26 @@ class TestNomaAnonymizationFlow: # Mock simplified Noma API response for LLM response check noma_response = { - "originalResponse": { - "response": { - "anonymizedContent": { - "anonymized": "My email is *******" + "aggregatedScanResult": True, + "scanResult": [ + { + "role": "assistant", + "type": "message", + "results": { + "anonymizedContent": { + "anonymized": "My email is *******" + }, + "sensitiveData": { + "PCI": { + "probability": 0.8, + "result": True, + "status": "SUCCESS" + }, + }, }, - "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() @@ -1443,18 +1818,22 @@ class TestNomaAnonymizationFlow: } noma_response = { - "originalResponse": { - "prompt": { - "anonymizedContent": { - "anonymized": "My email is *******" - }, - "dataDetector": { - "dataType1": {"result": True}, - }, - "contentDetector": {"result": False}, - }, - }, - "verdict": False, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "anonymizedContent": { + "anonymized": "My email is *******" + }, + "sensitiveData": { + "PII": {"result": True, "status": "SUCCESS"}, + }, + "harmfulContent": {"result": False, "status": "SUCCESS"}, + } + } + ] } mock_response = MagicMock() @@ -1518,15 +1897,19 @@ class TestNomaAnonymizationFlow: } noma_response = { - "originalResponse": { - "prompt": { - "dataDetector": { - "dataType1": {"result": True}, - }, - "contentDetector": {"result": False}, - }, - }, - "verdict": False, + "aggregatedScanResult": True, # True means unsafe + "scanResult": [ + { + "role": "user", + "type": "message", + "results": { + "sensitiveData": { + "PII": {"result": True, "status": "SUCCESS"}, + }, + "harmfulContent": {"result": False, "status": "SUCCESS"}, + } + } + ] } mock_response = MagicMock() @@ -1576,17 +1959,23 @@ class TestNomaAnonymizationFlow: # Mock Noma API response with no anonymized content available noma_response = { - "originalResponse": { - "response": { - "dataDetector": { - "dataType1": {"result": True}, + "aggregatedScanResult": True, + "scanResult": [ + { + "role": "assistant", + "type": "message", + "results": { + "sensitiveData": { + "PCI": { + "probability": 0.8, + "result": True, + "status": "SUCCESS" + }, + }, }, - "contentDetector": {"result": False}, }, - }, - "verdict": False, + ], } - mock_response = MagicMock() mock_response.json.return_value = noma_response mock_response.raise_for_status = MagicMock()