From 6a1f5bdc2f1be4deccc0db0280c1789014c6224c Mon Sep 17 00:00:00 2001 From: Dor Zion Date: Mon, 11 Aug 2025 12:05:28 +0300 Subject: [PATCH] Add Noma Security guardrail support --- .../docs/proxy/guardrails/noma_security.md | 299 +++++++++++ docs/my-website/sidebars.js | 1 + .../guardrail_hooks/noma/__init__.py | 36 ++ .../guardrails/guardrail_hooks/noma/noma.py | 403 ++++++++++++++ litellm/types/guardrails.py | 19 + .../guardrails/guardrail_hooks/test_noma.py | 498 ++++++++++++++++++ 6 files changed, 1256 insertions(+) create mode 100644 docs/my-website/docs/proxy/guardrails/noma_security.md create mode 100644 litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/noma/noma.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py diff --git a/docs/my-website/docs/proxy/guardrails/noma_security.md b/docs/my-website/docs/proxy/guardrails/noma_security.md new file mode 100644 index 0000000000..3a50841d65 --- /dev/null +++ b/docs/my-website/docs/proxy/guardrails/noma_security.md @@ -0,0 +1,299 @@ +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Noma Security + +Use [Noma Security](https://noma.security/) to protect your LLM applications with comprehensive AI content moderation and safety guardrails. + +## Quick Start + +### 1. Define Guardrails on your LiteLLM config.yaml + +Define your guardrails under the `guardrails` section: + +```yaml showLineNumbers title="litellm config.yaml" +model_list: + - model_name: gpt-4o-mini + litellm_params: + model: openai/gpt-4o-mini + api_key: os.environ/OPENAI_API_KEY + +guardrails: + - guardrail_name: "noma-guard" + litellm_params: + guardrail: noma + mode: "during_call" + api_key: os.environ/NOMA_API_KEY + api_base: os.environ/NOMA_API_BASE + - guardrail_name: "noma-pre-guard" + litellm_params: + guardrail: noma + mode: "pre_call" + api_key: os.environ/NOMA_API_KEY + api_base: os.environ/NOMA_API_BASE +``` + +#### Supported values for `mode` + +- `pre_call` Run **before** LLM call, on **input** +- `post_call` Run **after** LLM call, on **input & output** +- `during_call` Run **during** LLM call, on **input**. Same as `pre_call` but runs in parallel with the LLM call. Response not returned until guardrail check completes + +### 2. Start LiteLLM Gateway + +```shell +litellm --config config.yaml --detailed_debug +``` + +### 3. Test request + + + + +Expect this to fail since the request contains harmful content: + +```shell showLineNumbers title="Curl Request" +curl -i http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4o-mini", + "messages": [ + {"role": "user", "content": "Tell me how to hack into someone's email account"} + ] + }' +``` + +Expected response on failure: + +```json +{ + "error": { + "message": "{\n \"error\": \"Request blocked by Noma guardrail\",\n \"details\": {\n \"prompt\": {\n \"harmfulContent\": {\n \"result\": true,\n \"confidence\": 0.95\n }\n }\n }\n }", + "type": "None", + "param": "None", + "code": "400" + } +} +``` + + + + + +```shell showLineNumbers title="Curl Request" +curl -i http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "gpt-4o-mini", + "messages": [ + {"role": "user", "content": "What is the capital of France?"} + ] + }' +``` + +Expected response: + +```json +{ + "id": "chatcmpl-123", + "object": "chat.completion", + "created": 1677652288, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "The capital of France is Paris." + }, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21 + } +} +``` + + + + +## Supported Params + +```yaml +guardrails: + - guardrail_name: "noma-guard" + litellm_params: + guardrail: noma + mode: "pre_call" + api_key: os.environ/NOMA_API_KEY + api_base: os.environ/NOMA_API_BASE + ### OPTIONAL ### + # application_id: "my-app" + # monitor_mode: false + # block_failures: true +``` + +### Required Parameters + +- **`api_key`**: Your Noma Security API key (set as `os.environ/NOMA_API_KEY` in YAML config) + +### Optional Parameters + +- **`api_base`**: Noma API base URL (defaults to `https://api.noma.security/`) +- **`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`) + +## Environment Variables + +You can set these environment variables instead of hardcoding values in your config: + +```shell +export NOMA_API_KEY="your-api-key-here" +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 +``` + +## Advanced Configuration + +### Monitor Mode + +Use monitor mode to test your guardrails without blocking requests: + +```yaml +guardrails: + - guardrail_name: "noma-monitor" + litellm_params: + guardrail: noma + mode: "pre_call" + api_key: os.environ/NOMA_API_KEY + monitor_mode: true # Log violations but don't block +``` + +### Handling API Failures + +Control behavior when the Noma API is unavailable: + +```yaml +guardrails: + - guardrail_name: "noma-failopen" + litellm_params: + guardrail: noma + mode: "pre_call" + api_key: os.environ/NOMA_API_KEY + block_failures: false # Allow requests to proceed if guardrail API fails +``` + +### Multiple Guardrails + +Apply different configurations for input and output: + +```yaml +guardrails: + - guardrail_name: "noma-strict-input" + litellm_params: + guardrail: noma + mode: "pre_call" + api_key: os.environ/NOMA_API_KEY + block_failures: true + + - guardrail_name: "noma-monitor-output" + litellm_params: + guardrail: noma + mode: "post_call" + api_key: os.environ/NOMA_API_KEY + monitor_mode: true +``` + +## ✨ Pass Additional Parameters + +Use `extra_body` to pass additional parameters to the Noma Security API call, such as dynamically setting the application ID for specific requests. + + + + +```python +import openai +client = openai.OpenAI( + api_key="your-api-key", + base_url="http://0.0.0.0:4000" +) + +response = client.chat.completions.create( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Hello, how are you?"}], + extra_body={ + "guardrails": { + "noma-guard": { + "extra_body": { + "application_id": "my-specific-app-id" + } + } + } + } +) +``` + + + + +```shell +curl 'http://0.0.0.0:4000/v1/chat/completions' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "gpt-4o-mini", + "messages": [ + { + "role": "user", + "content": "Hello, how are you?" + } + ], + "guardrails": { + "noma-guard": { + "extra_body": { + "application_id": "my-specific-app-id" + } + } + } +}' +``` + + + +This allows you to override the default `application_id` parameter for specific requests, which is useful for tracking usage across different applications or components. + +## Response Details + +When content is blocked, Noma provides detailed information about the violations as JSON inside the `message` field, with the following structure: + +```json +{ + "error": "Request blocked by Noma guardrail", + "details": { + "prompt": { + "harmfulContent": { + "result": true, + "confidence": 0.95 + }, + "sensitiveData": { + "email": { + "result": true, + "entities": ["user@example.com"] + } + }, + "bannedTopics": { + "violence": { + "result": true, + "confidence": 0.88 + } + } + } + } +} +``` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 419afcd546..7d55525919 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -40,6 +40,7 @@ const sidebars = { "proxy/guardrails/guardrails_ai", "proxy/guardrails/lakera_ai", "proxy/guardrails/model_armor", + "proxy/guardrails/noma_security", "proxy/guardrails/openai_moderation", "proxy/guardrails/pangea", "proxy/guardrails/pillar_security", diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py new file mode 100644 index 0000000000..dc3e4d9768 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py @@ -0,0 +1,36 @@ +from typing import TYPE_CHECKING + +from litellm.types.guardrails import SupportedGuardrailIntegrations + +from .noma import NomaGuardrail + +if TYPE_CHECKING: + from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): + import litellm + + _noma_callback = NomaGuardrail( + guardrail_name=guardrail.get("guardrail_name", ""), + api_key=litellm_params.api_key, + api_base=litellm_params.api_base, + application_id=litellm_params.application_id, + monitor_mode=litellm_params.monitor_mode, + block_failures=litellm_params.block_failures, + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_noma_callback) + + return _noma_callback + + +guardrail_initializer_registry = { + SupportedGuardrailIntegrations.NOMA.value: initialize_guardrail, +} + + +guardrail_class_registry = { + SupportedGuardrailIntegrations.NOMA.value: NomaGuardrail, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py new file mode 100644 index 0000000000..ed5929f056 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -0,0 +1,403 @@ +# +-------------------------------------------------------------+ +# +# Noma Security Guardrail Integration for LiteLLM +# https://noma.security +# +# +-------------------------------------------------------------+ + +import copy +import os +from typing import Any, Dict, Literal, Optional, Union +from urllib.parse import urljoin + +from fastapi import HTTPException + +import litellm +from litellm import DualCache, ModelResponse +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import EmbeddingResponse, ImageResponse + + +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, + }, + ) + + 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", + ] 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 == "sensitiveData" 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. + + 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 + + +class NomaGuardrail(CustomGuardrail): + """ + Noma Security Guardrail for LiteLLM + + This guardrail integrates with Noma Security's AI-DR API to provide + content moderation and safety checks for LLM inputs and outputs. + """ + + _DEFAULT_API_BASE = "https://api.noma.security/" + _AIDR_ENDPOINT = "/ai-dr/v1/prompt/scan/aggregate" + + def __init__( + self, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + application_id: Optional[str] = None, + monitor_mode: Optional[bool] = None, + block_failures: Optional[bool] = None, + **kwargs, + ): + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) + self.api_key = api_key or os.environ.get("NOMA_API_KEY") + self.api_base = api_base or os.environ.get( + "NOMA_API_BASE", NomaGuardrail._DEFAULT_API_BASE + ) + self.application_id = application_id or os.environ.get( + "NOMA_APPLICATION_ID", "litellm" + ) + + if monitor_mode is None: + self.monitor_mode = ( + os.environ.get("NOMA_MONITOR_MODE", "false").lower() == "true" + ) + else: + self.monitor_mode = monitor_mode + + if block_failures is None: + self.block_failures = ( + os.environ.get("NOMA_BLOCK_FAILURES", "true").lower() == "true" + ) + else: + self.block_failures = block_failures + + super().__init__(**kwargs) + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: Literal[ + "completion", + "text_completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + "pass_through_endpoint", + "rerank", + "mcp_call", + ], + ) -> Optional[Union[Exception, str, dict]]: + verbose_proxy_logger.debug("Running Noma pre-call hook") + + if ( + self.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + is False + ): + return data + + try: + return await self._check_user_message(data, user_api_key_dict) + except NomaBlockedMessage: + raise + except Exception as e: + verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}") + + if self.block_failures and not self.monitor_mode: + raise + return data + + async def async_moderation_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + call_type: Literal[ + "completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + "responses", + "mcp_call", + ], + ) -> Union[Exception, str, dict, None]: + event_type: GuardrailEventHooks = GuardrailEventHooks.during_call + if self.should_run_guardrail(data=data, event_type=event_type) is not True: + return data + + try: + return await self._check_user_message(data, user_api_key_dict) + except NomaBlockedMessage: + raise + except Exception as e: + verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}") + + if self.block_failures and not self.monitor_mode: + raise + return data + + async def async_post_call_success_hook( + self, + data: dict, + user_api_key_dict: UserAPIKeyAuth, + response: Union[Any, ModelResponse, EmbeddingResponse, ImageResponse], + ): + event_type: GuardrailEventHooks = GuardrailEventHooks.post_call + if self.should_run_guardrail(data=data, event_type=event_type) is not True: + return response + + try: + return await self._check_llm_response(data, response, user_api_key_dict) + except NomaBlockedMessage: + raise + except Exception as e: + verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}") + if self.block_failures and not self.monitor_mode: + raise + return response + + async def _check_user_message( + self, + request_data: dict, + user_auth: UserAPIKeyAuth, + ) -> Union[Exception, str, dict, None]: + """Check user message for policy violations""" + extra_data = self.get_guardrail_dynamic_request_body_params(request_data) + + user_message = await self._extract_user_message(request_data) + if not user_message: + return request_data + + payload = {"request": {"text": user_message}} + response_json = await self._call_noma_api( + payload=payload, + llm_request_id=None, + request_data=request_data, + user_auth=user_auth, + extra_data=extra_data, + ) + await self._check_verdict("user", user_message, response_json) + + return request_data + + async def _check_llm_response( + self, + request_data: dict, + response: Union[Any, ModelResponse, EmbeddingResponse, ImageResponse], + user_auth: UserAPIKeyAuth, + ) -> Union[Exception, ModelResponse, Any]: + """Check LLM response for policy violations""" + extra_data = self.get_guardrail_dynamic_request_body_params(request_data) + + if not isinstance(response, litellm.ModelResponse): + return response + + content = None + for choice in response.choices: + if isinstance(choice, litellm.Choices) and choice.message.content: + content = choice.message.content + break + + if not content or not isinstance(content, str): + return response + + payload = {"response": {"text": content}} + + response_json = await self._call_noma_api( + payload=payload, + llm_request_id=response.id, + request_data=request_data, + user_auth=user_auth, + extra_data=extra_data, + ) + await self._check_verdict("assistant", content, response_json) + + return response + + async def _extract_user_message(self, data: dict) -> Optional[str]: + """Extract the last user message from request data""" + messages = data.get("messages", []) + if not messages: + return None + + # Get the last user message + user_messages = [msg for msg in messages if msg.get("role") == "user"] + if not user_messages: + return None + + last_user_message = user_messages[-1].get("content", "") + if not last_user_message or not isinstance(last_user_message, str): + return None + + return last_user_message + + async def _call_noma_api( + self, + payload: dict, + llm_request_id: Optional[str], + request_data: dict, + user_auth: UserAPIKeyAuth, + extra_data: dict, + ) -> 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 {}), + } + endpoint = urljoin( + self.api_base or "https://api.noma.security/", NomaGuardrail._AIDR_ENDPOINT + ) + + response = await self.async_handler.post( + endpoint, + headers=headers, + json={ + **payload, + "context": { + "applicationId": extra_data.get("application_id") + or request_data.get("metadata", {}) + .get("headers", {}) + .get("x-noma-application-id"), + "ipAddress": request_data.get("metadata", {}).get( + "requester_ip_address", None + ), + "userId": user_auth.user_email + if user_auth.user_email + else user_auth.user_id, + "sessionId": call_id, + "requestId": llm_request_id, + }, + }, + ) + response.raise_for_status() + + return response.json() + + async def _check_verdict( + self, + type: Literal["user", "assistant"], + message: str, + response_json: dict, + ) -> None: + """ + Check the verdict from the Noma API and raise an exception if needed + """ + if not response_json.get("verdict", True): + msg = str.format( + "Noma guardrail blocked {type} message: {message}", + type=type, + message=message, + ) + + if self.monitor_mode: + verbose_proxy_logger.warning(msg) + else: + verbose_proxy_logger.debug(msg) + original_response = response_json.get("originalResponse", {}) + raise NomaBlockedMessage(original_response) + else: + msg = str.format( + "Noma guardrail allowed {type} message: {message}", + type=type, + message=message, + ) + if self.monitor_mode: + verbose_proxy_logger.info(msg) + else: + verbose_proxy_logger.debug(msg) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index fd18484a89..f31f304bda 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -40,6 +40,7 @@ class SupportedGuardrailIntegrations(Enum): AZURE_TEXT_MODERATIONS = "azure/text_moderations" MODEL_ARMOR = "model_armor" OPENAI_MODERATION = "openai_moderation" + NOMA = "noma" class Role(Enum): SYSTEM = "system" @@ -359,6 +360,23 @@ class PillarGuardrailConfigModel(BaseModel): ) +class NomaGuardrailConfigModel(BaseModel): + """Configuration parameters for the Noma Security guardrail""" + + application_id: Optional[str] = Field( + default=None, + description="Application ID for Noma Security. Defaults to 'litellm' if not provided", + ) + monitor_mode: Optional[bool] = Field( + default=None, + description="If True, logs violations without blocking. Defaults to False if not provided", + ) + block_failures: Optional[bool] = Field( + default=None, + description="If True, blocks requests on API failures. Defaults to True if not provided", + ) + + class BaseLitellmParams(BaseModel): # works for new and patch update guardrails api_key: Optional[str] = Field( default=None, description="API key for the guardrail service" @@ -445,6 +463,7 @@ class LitellmParams( LakeraV2GuardrailConfigModel, LassoGuardrailConfigModel, PillarGuardrailConfigModel, + NomaGuardrailConfigModel, BaseLitellmParams, ): guardrail: str = Field(description="The type of guardrail integration to use") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py new file mode 100644 index 0000000000..aeea5f81b1 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py @@ -0,0 +1,498 @@ +import os +from unittest.mock import MagicMock, patch + +import httpx +import pytest + +import litellm +from litellm import ModelResponse +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.noma import ( + NomaGuardrail, + initialize_guardrail, +) +from litellm.proxy.guardrails.guardrail_hooks.noma.noma import NomaBlockedMessage +from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.types.utils import Choices, Message + + +@pytest.fixture +def noma_guardrail(): + """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, + ) + + +@pytest.fixture +def mock_user_api_key_dict(): + """Create a mock UserAPIKeyAuth object""" + return UserAPIKeyAuth( + user_id="test-user-id", + user_email="test@example.com", + key_name="test-key", + key_alias=None, + team_id=None, + team_alias=None, + user_role=None, + api_key="test-api-key", + permissions={}, + models=[], + spend=0.0, + max_budget=None, + soft_budget=None, + tpm_limit=None, + rpm_limit=None, + parallel_request_limit=None, + metadata={}, + max_parallel_requests=None, + allowed_cache_controls=[], + model_spend={}, + model_max_budget={}, + ) + + +@pytest.fixture +def mock_request_data(): + """Create mock request data""" + return { + "messages": [ + {"role": "system", "content": "You are a helpful assistant"}, + {"role": "user", "content": "Hello, how are you?"}, + ], + "litellm_call_id": "test-call-id", + "metadata": {"requester_ip_address": "192.168.1.1"}, + } + + +class TestNomaGuardrailConfiguration: + """Test configuration and initialization of Noma guardrail""" + + def test_init_with_config(self): + """Test initializing Noma guardrail via init_guardrails_v2""" + with patch.dict( + os.environ, + { + "NOMA_API_KEY": "test-api-key", + "NOMA_API_BASE": "https://api.test.noma.security/", + }, + ): + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "noma-pre-guard", + "litellm_params": { + "guardrail": "noma", + "mode": "pre_call", + "application_id": "test-app", + "monitor_mode": False, + "block_failures": True, + }, + } + ], + config_file_path="", + ) + + def test_init_with_env_vars(self): + """Test initialization with environment variables""" + with patch.dict( + os.environ, + { + "NOMA_API_KEY": "env-api-key", + "NOMA_API_BASE": "https://env.api.noma.security/", + "NOMA_APPLICATION_ID": "env-app-id", + "NOMA_MONITOR_MODE": "true", + "NOMA_BLOCK_FAILURES": "false", + }, + ): + guardrail = NomaGuardrail() + assert guardrail.api_key == "env-api-key" + assert guardrail.api_base == "https://env.api.noma.security/" + assert guardrail.application_id == "env-app-id" + assert guardrail.monitor_mode is True + assert guardrail.block_failures is False + + def test_init_with_params_override_env(self): + """Test that constructor params override environment variables""" + with patch.dict( + os.environ, + { + "NOMA_API_KEY": "env-api-key", + "NOMA_MONITOR_MODE": "true", + }, + ): + guardrail = NomaGuardrail( + api_key="param-api-key", + monitor_mode=False, + ) + assert guardrail.api_key == "param-api-key" + assert guardrail.monitor_mode is False + + def test_initialize_guardrail_function(self): + """Test the initialize_guardrail function""" + from litellm.types.guardrails import Guardrail, LitellmParams + + litellm_params = LitellmParams( + guardrail="noma", + mode="pre_call", + api_key="test-key", + api_base="https://test.api/", + application_id="test-app", + monitor_mode=True, + block_failures=False, + ) + + guardrail = Guardrail( + guardrail_name="test-guardrail", + litellm_params=litellm_params, + ) + + with patch("litellm.logging_callback_manager.add_litellm_callback") as mock_add: + result = initialize_guardrail(litellm_params, guardrail) + + assert isinstance(result, NomaGuardrail) + assert result.api_key == "test-key" + assert result.api_base == "https://test.api/" + assert result.application_id == "test-app" + assert result.monitor_mode is True + assert result.block_failures is False + mock_add.assert_called_once_with(result) + + +class TestNomaBlockedMessage: + """Test the NomaBlockedMessage exception class""" + + def test_blocked_message_basic(self): + """Test basic blocked message creation""" + response = { + "verdict": False, + "prompt": { + "harmfulContent": {"result": True, "confidence": 0.9}, + "code": {"result": False, "confidence": 0.1}, + }, + } + + 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 "code" not in exception.detail["details"]["prompt"] + + def test_blocked_message_with_sensitive_data(self): + """Test blocked message with sensitive data detection""" + response = { + "verdict": False, + "prompt": { + "sensitiveData": { + "email": {"result": True, "entities": ["test@example.com"]}, + "phone": {"result": False}, + }, + }, + } + + exception = NomaBlockedMessage(response) + assert "email" in exception.detail["details"]["prompt"]["sensitiveData"] + assert "phone" not in exception.detail["details"]["prompt"]["sensitiveData"] + + 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}, + }, + }, + } + + exception = NomaBlockedMessage(response) + assert "violence" in exception.detail["details"]["prompt"]["bannedTopics"] + assert "politics" not in exception.detail["details"]["prompt"]["bannedTopics"] + + +class TestNomaGuardrailHooks: + """Test the guardrail hook methods""" + + @pytest.mark.asyncio + async def test_pre_call_hook_allowed( + self, noma_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test pre-call hook when content is allowed""" + mock_response = MagicMock() + mock_response.json.return_value = {"verdict": True} + 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=mock_request_data, + call_type="completion", + ) + + assert result == mock_request_data + mock_post.assert_called_once() + + # 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[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( + self, noma_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test pre-call hook when content is blocked""" + mock_response = MagicMock() + mock_response.json.return_value = { + "verdict": False, + "originalResponse": { + "prompt": {"harmfulContent": {"result": True, "confidence": 0.9}} + }, + } + mock_response.raise_for_status = MagicMock() + + 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=mock_request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "harmfulContent" in exc_info.value.detail["details"]["prompt"] + + @pytest.mark.asyncio + async def test_pre_call_hook_monitor_mode( + self, mock_user_api_key_dict, mock_request_data + ): + """Test pre-call hook in monitor mode (logs but doesn't block)""" + guardrail = NomaGuardrail( + api_key="test-key", + monitor_mode=True, + guardrail_name="test-guardrail", + event_hook="pre_call", + default_on=True, + ) + + mock_response = MagicMock() + mock_response.json.return_value = { + "verdict": False, + "originalResponse": {"prompt": {"harmfulContent": {"result": True}}}, + } + mock_response.raise_for_status = MagicMock() + + with patch.object(guardrail.async_handler, "post", return_value=mock_response): + # Should not raise exception in monitor mode + result = await guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key_dict, + cache=MagicMock(), + data=mock_request_data, + call_type="completion", + ) + + assert result == mock_request_data + + @pytest.mark.asyncio + async def test_post_call_success_hook( + self, noma_guardrail, mock_user_api_key_dict, mock_request_data + ): + """Test post-call success hook""" + # Create a mock ModelResponse + response = ModelResponse( + id="test-response-id", + choices=[ + Choices( + finish_reason="stop", + index=0, + message=Message( + content="I'm doing well, thank you!", 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_api_response = MagicMock() + mock_api_response.json.return_value = {"verdict": True} + mock_api_response.raise_for_status = MagicMock() + + # Update guardrail to use post_call event hook + noma_guardrail.event_hook = "post_call" + + with patch.object( + noma_guardrail.async_handler, "post", return_value=mock_api_response + ) as mock_post: + result = await noma_guardrail.async_post_call_success_hook( + data=mock_request_data, + user_api_key_dict=mock_user_api_key_dict, + response=response, + ) + + 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 + ): + """Test moderation hook (during_call)""" + # Update guardrail to use during_call event hook + noma_guardrail.event_hook = "during_call" + + mock_response = MagicMock() + mock_response.json.return_value = {"verdict": True} + mock_response.raise_for_status = MagicMock() + + with patch.object( + noma_guardrail.async_handler, "post", return_value=mock_response + ): + result = await noma_guardrail.async_moderation_hook( + data=mock_request_data, + user_api_key_dict=mock_user_api_key_dict, + call_type="completion", + ) + + assert result == mock_request_data + + @pytest.mark.asyncio + async def test_api_failure_handling( + self, noma_guardrail, mock_user_api_key_dict, mock_request_data + ): + with patch.object( + noma_guardrail.async_handler, + "post", + side_effect=httpx.HTTPStatusError( + "API Error", request=MagicMock(), response=MagicMock(status_code=500) + ), + ): + with pytest.raises(httpx.HTTPStatusError): + await noma_guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key_dict, + cache=MagicMock(), + data=mock_request_data, + call_type="completion", + ) + + @pytest.mark.asyncio + async def test_api_failure_no_block( + self, mock_user_api_key_dict, mock_request_data + ): + guardrail = NomaGuardrail( + api_key="test-key", + block_failures=False, + guardrail_name="test-guardrail", + event_hook="pre_call", + default_on=True, + ) + + with patch.object( + guardrail.async_handler, + "post", + side_effect=httpx.HTTPStatusError( + "API Error", request=MagicMock(), response=MagicMock(status_code=500) + ), + ): + result = await guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key_dict, + cache=MagicMock(), + data=mock_request_data, + call_type="completion", + ) + + assert result == mock_request_data + + def test_extract_user_message(self, noma_guardrail): + data = { + "messages": [ + {"role": "system", "content": "System prompt"}, + {"role": "user", "content": "First user message"}, + {"role": "assistant", "content": "Assistant response"}, + {"role": "user", "content": "Second user message"}, + ] + } + + import asyncio + + message = asyncio.run(noma_guardrail._extract_user_message(data)) + assert message == "Second user message" + + data = {"messages": [{"role": "system", "content": "System prompt"}]} + message = asyncio.run(noma_guardrail._extract_user_message(data)) + assert message is None + + data = {"messages": []} + message = asyncio.run(noma_guardrail._extract_user_message(data)) + assert message is None + + data = {} + message = asyncio.run(noma_guardrail._extract_user_message(data)) + assert message is None + + +class TestIntegration: + @pytest.mark.asyncio + async def test_full_guardrail_flow(self): + """Test full guardrail flow with multiple hooks""" + with patch.dict( + os.environ, + { + "NOMA_API_KEY": "test-api-key", + "NOMA_API_BASE": "https://api.test.noma.security/", + }, + ): + init_guardrails_v2( + all_guardrails=[ + { + "guardrail_name": "noma-pre-guard", + "litellm_params": { + "guardrail": "noma", + "mode": "pre_call", + "application_id": "test-app", + }, + }, + { + "guardrail_name": "noma-post-guard", + "litellm_params": { + "guardrail": "noma", + "mode": "post_call", + "application_id": "test-app", + }, + }, + ], + config_file_path="", + ) + + custom_loggers = ( + litellm.logging_callback_manager.get_custom_loggers_for_type( + callback_type=litellm.integrations.custom_guardrail.CustomGuardrail + ) + ) + assert len(custom_loggers) >= 2