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