ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Noma Security guardrail support

Browse Source
This commit is contained in:
Dor Zion 2025-08-11 12:05:28 +03:00
parent 5ae44e3275
commit 6a1f5bdc2f
No known key found for this signature in database
6 changed files with 1256 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

299
docs/my-website/docs/proxy/guardrails/noma_security.md Normal file
View File

@ -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
<Tabs>
<TabItem label="Unsuccessful call" value="not-allowed">
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"
}
}
```
</TabItem>
<TabItem label="Successful Call" value="allowed">
```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
}
}
```
</TabItem>
</Tabs>
## 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.
<Tabs>
<TabItem value="openai" label="OpenAI Python">
```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"
}
}
}
}
)
```
</TabItem>
<TabItem value="curl" label="Curl">
```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"
}
}
}
}'
```
</TabItem>
</Tabs>
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
}
}
}
}
}
```

1
docs/my-website/sidebars.js
View File

@ -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",

36
litellm/proxy/guardrails/guardrail_hooks/noma/__init__.py Normal file
View File

@ -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,
}

403
litellm/proxy/guardrails/guardrail_hooks/noma/noma.py Normal file
View File

@ -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)

19
litellm/types/guardrails.py
View File

@ -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")

498
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py Normal file
View File

@ -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