tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-19 20:24:27 +00:00
Code Issues Packages Projects Releases Wiki Activity

Litellm noma guardrail support images (#16199)

Browse Source 
* noma support v2 api and images with during call

* supporting streams and images with texts

* Supporting text now

* annonymization works

* removing function

* fixing noma.py

* all old tests pass

* adding new tests

* removing changes

* Fixing application id headers

* fix whitespace

* deleting unused imports
This commit is contained in:
dean-zavad
2025-11-05 11:23:39 -08:00
committed by GitHub
parent 9b925c7e47
commit f19356db60
2 changed files with 831 additions and 388 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/guardrails/guardrail_hooks/noma/noma.py
+177 -123
View File
@@ -6,11 +6,11 @@
# +-------------------------------------------------------------+
import asyncio
import copy
import os
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, Final, Literal, Optional, Type, Union
from urllib.parse import urljoin
import json
from fastapi import HTTPException
@@ -26,6 +26,18 @@ from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import EmbeddingResponse, GuardrailStatus, ImageResponse
from litellm.types.utils import (
ModelResponseStream,
)
from typing import (
List,
AsyncGenerator
)
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.main import stream_chunk_builder
from litellm.types.utils import TextCompletionResponse
# Constants
USER_ROLE: Final[Literal["user"]] = "user"
ASSISTANT_ROLE: Final[Literal["assistant"]] = "assistant"
@@ -43,79 +55,14 @@ class NomaBlockedMessage(HTTPException):
"""Exception raised when Noma guardrail blocks a message"""
def __init__(self, classification_response: dict):
classification = self._filter_triggered_classifications(classification_response)
super().__init__(
status_code=400,
detail={
"error": "Request blocked by Noma guardrail",
"details": classification,
"details": classification_response,
},
)
def _filter_triggered_classifications(
self,
response_dict: dict,
) -> dict:
"""Filter and return only triggered classifications"""
filtered_response = copy.deepcopy(response_dict)
# Filter prompt classifications if present
if filtered_response.get("prompt"):
filtered_response["prompt"] = self.filter_classification_object(
filtered_response["prompt"]
)
# Filter response classifications if present
if filtered_response.get("response"):
filtered_response["response"] = self.filter_classification_object(
filtered_response["response"]
)
return filtered_response
def filter_classification_object(
self,
classification_obj: dict,
) -> dict:
"""Filter classification object to only include triggered items"""
if not classification_obj:
return {}
result = {}
for key, value in classification_obj.items():
if value is None:
continue
if key in [
"allowedTopics",
"bannedTopics",
"topicGuardrails",
"topicDetector", # Mock name for tests
] and isinstance(value, dict):
filtered_topics = {}
for topic, topic_result in value.items():
if self._is_result_true(topic_result):
filtered_topics[topic] = topic_result
if filtered_topics:
result[key] = filtered_topics
elif key in SENSITIVE_DATA_DETECTOR_KEYS and isinstance(value, dict):
filtered_sensitive = {}
for data_type, data_result in value.items():
if self._is_result_true(data_result):
filtered_sensitive[data_type] = data_result
if filtered_sensitive:
result[key] = filtered_sensitive
elif isinstance(value, dict) and "result" in value:
if self._is_result_true(value):
result[key] = value
return result
def _is_result_true(self, result_obj: Optional[Dict[str, Any]]) -> bool:
"""
Check if a result object has a "result" field that is True.
@@ -141,7 +88,7 @@ class NomaGuardrail(CustomGuardrail):
"""
_DEFAULT_API_BASE = "https://api.noma.security/"
_AIDR_ENDPOINT = "/ai-dr/v1/prompt/scan/aggregate"
_AIDR_ENDPOINT = "/ai-dr/v2/prompt/scan"
def __init__(
self,
@@ -212,7 +159,15 @@ class NomaGuardrail(CustomGuardrail):
if not user_message:
return None
payload = {"request": {"text": user_message}}
payload = {
"input": [
{
"type": "message",
"role": "user",
"content": user_message
}
]
}
response_json = await self._call_noma_api(
payload=payload,
llm_request_id=None,
@@ -240,9 +195,9 @@ class NomaGuardrail(CustomGuardrail):
if self.monitor_mode:
await self._handle_verdict_background(
USER_ROLE, user_message, response_json
USER_ROLE, json.dumps(user_message), response_json
)
return user_message
return json.dumps(user_message)
# Check if we should anonymize content
if self._should_anonymize(response_json, USER_ROLE):
@@ -257,8 +212,8 @@ class NomaGuardrail(CustomGuardrail):
)
return anonymized_content
await self._check_verdict(USER_ROLE, user_message, response_json)
return user_message
await self._check_verdict(USER_ROLE, json.dumps(user_message), response_json)
return json.dumps(user_message)
async def _process_llm_response_check(
self,
@@ -283,7 +238,20 @@ class NomaGuardrail(CustomGuardrail):
if not content or not isinstance(content, str):
return None
payload = {"response": {"text": content}}
payload = {
"input": [
{
"type": "message",
"role": "assistant",
"content": [
{
"type": "input_text",
"text": content
}
]
}
]
}
response_json = await self._call_noma_api(
payload=payload,
@@ -312,7 +280,7 @@ class NomaGuardrail(CustomGuardrail):
if self.monitor_mode:
await self._handle_verdict_background(
ASSISTANT_ROLE, content, response_json
ASSISTANT_ROLE, json.dumps(content), response_json
)
return content
@@ -349,18 +317,19 @@ class NomaGuardrail(CustomGuardrail):
if not isinstance(response_json, dict):
return "guardrail_failed_to_respond"
# Get the verdict from the response
verdict = response_json.get("verdict", True)
# Get the aggregatedScanResult from the response
# aggregatedScanResult=True means unsafe (block), False means safe (allow)
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
# If verdict is True, content is allowed
if verdict is True:
# If aggregatedScanResult is False, content is safe/allowed
if aggregated_scan_result is False:
return "success"
# If verdict is False, content is blocked/flagged
if verdict is False:
# If aggregatedScanResult is True, content is blocked/flagged
if aggregated_scan_result is True:
return "guardrail_intervened"
# If verdict is missing or invalid, treat as failure
# If aggregatedScanResult is missing or invalid, treat as failure
return "guardrail_failed_to_respond"
except Exception as e:
@@ -419,17 +388,16 @@ class NomaGuardrail(CustomGuardrail):
Returns:
The anonymized content string if available, None otherwise
"""
original_response = response_json.get("originalResponse", {})
if message_type == USER_ROLE:
prompt_data = original_response.get("prompt", {})
anonymized_data = prompt_data.get("anonymizedContent", {})
return anonymized_data.get("anonymized")
elif message_type == ASSISTANT_ROLE:
response_data = original_response.get("response", {})
anonymized_data = response_data.get("anonymizedContent", {})
return anonymized_data.get("anonymized")
# Extract from new scanResult structure
scan_result = response_json.get("scanResult", [])
if not scan_result:
return None
# Find the scan result matching the message type (role)
for result_item in scan_result:
if result_item.get("role") == message_type:
return result_item.get("results", {}).get("anonymizedContent", {}).get("anonymized", "")
return None
def _should_anonymize(self, response_json: dict, message_type: MessageRole) -> bool:
@@ -437,8 +405,8 @@ class NomaGuardrail(CustomGuardrail):
Determine if content should be anonymized based on Noma API response.
Logic:
- If verdict=True: Content is safe, anonymize if anonymized version exists
- If verdict=False: Check if only sensitiveData detectors have result=True
- If aggregatedScanResult=False: Content is safe, anonymize if anonymized version exists
- If aggregatedScanResult=True: Check if only sensitiveData detectors have result=True
- If yes: Anonymize
- If no: Block (other violations detected)
@@ -453,23 +421,26 @@ class NomaGuardrail(CustomGuardrail):
if self.monitor_mode or not self.anonymize_input:
return False
verdict = response_json.get("verdict", True)
# If verdict is True, anonymize (content is considered safe)
if verdict:
# aggregatedScanResult=False means safe, True means unsafe
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
# If aggregatedScanResult is False, content is safe - anonymize if available
if not aggregated_scan_result:
return True
# If verdict is False, check if only sensitive data detectors have result=True
original_response = response_json.get("originalResponse", {})
if message_type == USER_ROLE:
classification_obj = original_response.get("prompt", {})
elif message_type == ASSISTANT_ROLE:
classification_obj = original_response.get("response", {})
else:
# If aggregatedScanResult is True (unsafe), check if only sensitive data detectors triggered
scan_result = response_json.get("scanResult", [])
if not scan_result:
return False
if not isinstance(scan_result, list) or len(scan_result) == 0:
return False
for result_item in scan_result:
if result_item.get("role") == message_type:
return self._should_only_sensitive_data_failed(result_item.get("results", {}))
# Anonymize only if solely sensitive data (PII/PCI/secrets) was detected
return self._should_only_sensitive_data_failed(classification_obj)
return False
def _is_result_true(self, result_obj: Optional[Dict[str, Any]]) -> bool:
"""
@@ -557,12 +528,17 @@ class NomaGuardrail(CustomGuardrail):
message: str,
response_json: dict,
) -> None:
"""Handle verdict from Noma API in background - logging only, never blocks"""
"""Handle aggregatedScanResult from Noma API in background - logging only, never blocks
aggregatedScanResult=True means unsafe, False means safe
"""
try:
if not response_json.get("verdict", True):
# aggregatedScanResult=True means blocked, False means allowed
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
if aggregated_scan_result: # True = unsafe
msg = f"Noma guardrail blocked {type} message: {message}"
verbose_proxy_logger.warning(msg)
else:
else: # False = safe
msg = f"Noma guardrail allowed {type} message: {message}"
verbose_proxy_logger.info(msg)
except Exception as e:
@@ -588,6 +564,7 @@ class NomaGuardrail(CustomGuardrail):
"anthropic_messages",
],
) -> Optional[Union[Exception, str, dict]]:
verbose_proxy_logger.debug("Running Noma pre-call hook")
if (
@@ -756,7 +733,7 @@ class NomaGuardrail(CustomGuardrail):
request_data: dict,
response: LLMResponse,
user_auth: UserAPIKeyAuth,
) -> Union[Exception, ModelResponse, Any]:
) -> Any:
"""Check LLM response for policy violations"""
content = await self._process_llm_response_check(
request_data, response, user_auth
@@ -766,7 +743,7 @@ class NomaGuardrail(CustomGuardrail):
return response
async def _extract_user_message(self, data: dict) -> Optional[str]:
async def _extract_user_message(self, data: dict) -> Optional[List[dict]]:
"""Extract the last user message from request data"""
messages = data.get("messages", [])
if not messages:
@@ -778,10 +755,40 @@ class NomaGuardrail(CustomGuardrail):
return None
last_user_message = user_messages[-1].get("content", "")
if not last_user_message or not isinstance(last_user_message, str):
if isinstance(last_user_message, str):
return [{
"type": "input_text",
"text": last_user_message
}]
elif isinstance(last_user_message, list):
converted_messages = []
for message in last_user_message:
converted_message = self._convert_single_user_message_to_payload(message)
if converted_message is not None:
converted_messages.append(converted_message)
return converted_messages
else:
return None
return last_user_message
def _convert_single_user_message_to_payload(self, user_message: Any) -> Optional[dict]:
if isinstance(user_message, str):
return {
"type": "input_text",
"text": user_message
}
elif user_message.get("type", "") == "image_url":
return {
"type": "input_image",
"image_url": user_message.get("image_url", {}).get("url", "")
}
elif user_message.get("type", "") == "text":
return {
"type": "input_text",
"text": user_message.get("text", "")
}
else:
return None
async def _call_noma_api(
self,
@@ -793,7 +800,6 @@ class NomaGuardrail(CustomGuardrail):
) -> dict:
call_id = request_data.get("litellm_call_id")
headers = {
"X-Noma-AIDR-Application-ID": self.application_id,
**({"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}),
**({"X-Noma-Request-ID": call_id} if call_id else {}),
}
@@ -806,11 +812,11 @@ class NomaGuardrail(CustomGuardrail):
headers=headers,
json={
**payload,
"context": {
"x-noma-context": {
"applicationId": extra_data.get("application_id")
or request_data.get("metadata", {})
.get("headers", {})
.get("x-noma-application-id"),
.get("x-noma-application-id") or self.application_id,
"ipAddress": request_data.get("metadata", {}).get(
"requester_ip_address", None
),
@@ -833,24 +839,29 @@ class NomaGuardrail(CustomGuardrail):
response_json: dict,
) -> None:
"""
Check the verdict from the Noma API and raise an exception if needed
Check the aggregatedScanResult from the Noma API and raise an exception if needed.
aggregatedScanResult=True means unsafe (block), False means safe (allow)
"""
if not response_json.get("verdict", True):
# aggregatedScanResult=True means blocked, False means allowed
aggregated_scan_result = response_json.get("aggregatedScanResult", False)
if aggregated_scan_result: # True = unsafe, block it
msg = f"Noma guardrail blocked {type} message: {message}"
if self.monitor_mode:
verbose_proxy_logger.warning(msg)
else:
verbose_proxy_logger.debug(msg)
original_response = response_json.get("originalResponse", {})
original_response = response_json.get("scanResult", {})
# Use the full response as the original response for error details
raise NomaBlockedMessage(original_response)
else:
else: # False = safe, allow it
msg = f"Noma guardrail allowed {type} message: {message}"
if self.monitor_mode:
verbose_proxy_logger.info(msg)
else:
verbose_proxy_logger.debug(msg)
@staticmethod
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
from litellm.types.proxy.guardrails.guardrail_hooks.noma import (
@@ -859,3 +870,46 @@ class NomaGuardrail(CustomGuardrail):
return NomaGuardrailConfigModel
async def async_post_call_streaming_iterator_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
response: Any,
request_data: dict,
) -> AsyncGenerator[ModelResponseStream, None]:
"""Process streaming response chunks with Noma guardrail."""
all_chunks: List[ModelResponseStream] = []
async for chunk in response:
all_chunks.append(chunk)
if not all_chunks:
return
assembled_model_response: Optional[
Union[ModelResponse, TextCompletionResponse]
] = stream_chunk_builder(chunks=all_chunks)
if isinstance(assembled_model_response, ModelResponse):
try:
processed_response = await self._check_llm_response(
request_data, assembled_model_response, user_api_key_dict
)
except NomaBlockedMessage:
raise
except Exception as e:
if self.block_failures:
raise
verbose_proxy_logger.error(
f"Noma streaming post-call hook failed: {str(e)}"
)
for chunk in all_chunks:
yield chunk
return
mock_response = MockResponseIterator(model_response=processed_response)
async for chunk in mock_response:
yield chunk
return
for chunk in all_chunks:
yield chunk
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_noma.py
+654 -265
View File
File diff suppressed because it is too large Load Diff