mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-21 08:26:34 +00:00
[Feat] New Guardrail - Dynamo AI Guardrail (#15920)
* add dynamo types * fix Dynamo guard * add dynamo guardrail * add dynamo ai docs guard * docs fix * test dynamo * test LASSO
This commit is contained in:
Ishaan Jaff
@@ -440,6 +440,10 @@ router_settings:
|
||||
| DAYS_IN_A_MONTH | Days in a month for calculation purposes. Default is 28
|
||||
| DAYS_IN_A_WEEK | Days in a week for calculation purposes. Default is 7
|
||||
| DAYS_IN_A_YEAR | Days in a year for calculation purposes. Default is 365
|
||||
| DYNAMOAI_API_KEY | API key for DynamoAI Guardrails service
|
||||
| DYNAMOAI_API_BASE | Base URL for DynamoAI API. Default is https://api.dynamo.ai
|
||||
| DYNAMOAI_MODEL_ID | Model ID for DynamoAI tracking/logging purposes
|
||||
| DYNAMOAI_POLICY_IDS | Comma-separated list of DynamoAI policy IDs to apply
|
||||
| DD_BASE_URL | Base URL for Datadog integration
|
||||
| DATADOG_BASE_URL | (Alternative to DD_BASE_URL) Base URL for Datadog integration
|
||||
| _DATADOG_BASE_URL | (Alternative to DD_BASE_URL) Base URL for Datadog integration
|
||||
|
||||
@@ -0,0 +1,214 @@
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# DynamoAI Guardrails
|
||||
|
||||
LiteLLM supports DynamoAI guardrails for content moderation and policy enforcement on LLM inputs and outputs.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Define Guardrails on your LiteLLM config.yaml
|
||||
|
||||
Define your guardrails under the `guardrails` section:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: openai/gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
guardrails:
|
||||
- guardrail_name: "dynamoai-guard"
|
||||
litellm_params:
|
||||
guardrail: dynamoai
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/DYNAMOAI_API_KEY
|
||||
```
|
||||
|
||||
#### Supported values for `mode`
|
||||
|
||||
- `pre_call` - Run **before** LLM call, on **input**
|
||||
- `post_call` - Run **after** LLM call, on **output**
|
||||
- `during_call` - Run **during** LLM call, on **input**. Same as `pre_call` but runs in parallel as LLM call
|
||||
|
||||
### 2. Set Environment Variables
|
||||
|
||||
```bash
|
||||
export DYNAMOAI_API_KEY="your-api-key"
|
||||
# Optional: Set policy IDs via environment variable (comma-separated)
|
||||
export DYNAMOAI_POLICY_IDS="policy-id-1,policy-id-2,policy-id-3"
|
||||
```
|
||||
|
||||
### 3. Start LiteLLM Gateway
|
||||
|
||||
```shell
|
||||
litellm --config config.yaml --detailed_debug
|
||||
```
|
||||
|
||||
### 4. Test Request
|
||||
|
||||
**[Langchain, OpenAI SDK Usage Examples](../proxy/user_keys#request-format)**
|
||||
|
||||
<Tabs>
|
||||
<TabItem label="Successful Call" value="allowed">
|
||||
|
||||
```shell showLineNumbers title="Successful Request"
|
||||
curl -i http://localhost:4000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is the capital of France?"}
|
||||
],
|
||||
"guardrails": ["dynamoai-guard"]
|
||||
}'
|
||||
```
|
||||
|
||||
**Response: HTTP 200 Success**
|
||||
|
||||
Content passes all policy checks and is allowed through.
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem label="Blocked Call" value="not-allowed">
|
||||
|
||||
```shell showLineNumbers title="Blocked Request"
|
||||
curl -i http://localhost:4000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Content that violates policy"}
|
||||
],
|
||||
"guardrails": ["dynamoai-guard"]
|
||||
}'
|
||||
```
|
||||
|
||||
**Expected Response on Block: HTTP 400 Error**
|
||||
|
||||
```json showLineNumbers
|
||||
{
|
||||
"error": {
|
||||
"message": "Guardrail failed: 1 violation(s) detected\n\n- POLICY NAME:\n Action: BLOCK\n Method: TOXICITY\n Description: Policy description\n Policy ID: policy-id-123",
|
||||
"type": "None",
|
||||
"param": "None",
|
||||
"code": "400"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Advanced Configuration
|
||||
|
||||
### Specify Policy IDs
|
||||
|
||||
Configure specific DynamoAI policies to apply:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "dynamoai-policies"
|
||||
litellm_params:
|
||||
guardrail: dynamoai
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/DYNAMOAI_API_KEY
|
||||
policy_ids:
|
||||
- "policy-id-1"
|
||||
- "policy-id-2"
|
||||
- "policy-id-3"
|
||||
```
|
||||
|
||||
### Custom API Base
|
||||
|
||||
Specify a custom DynamoAI API endpoint:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "dynamoai-custom"
|
||||
litellm_params:
|
||||
guardrail: dynamoai
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/DYNAMOAI_API_KEY
|
||||
api_base: "https://custom.dynamo.ai"
|
||||
```
|
||||
|
||||
### Model ID for Tracking
|
||||
|
||||
Add a model ID for tracking and logging purposes:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "dynamoai-tracked"
|
||||
litellm_params:
|
||||
guardrail: dynamoai
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/DYNAMOAI_API_KEY
|
||||
model_id: "gpt-4-production"
|
||||
```
|
||||
|
||||
### Input and Output Guardrails
|
||||
|
||||
Configure separate guardrails for input and output:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
# Input guardrail
|
||||
- guardrail_name: "dynamoai-input"
|
||||
litellm_params:
|
||||
guardrail: dynamoai
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/DYNAMOAI_API_KEY
|
||||
|
||||
# Output guardrail
|
||||
- guardrail_name: "dynamoai-output"
|
||||
litellm_params:
|
||||
guardrail: dynamoai
|
||||
mode: "post_call"
|
||||
api_key: os.environ/DYNAMOAI_API_KEY
|
||||
```
|
||||
|
||||
## Configuration Options
|
||||
|
||||
| Parameter | Type | Description | Default |
|
||||
|-----------|------|-------------|---------|
|
||||
| `api_key` | string | DynamoAI API key (required) | `DYNAMOAI_API_KEY` env var |
|
||||
| `api_base` | string | DynamoAI API base URL | `https://api.dynamo.ai` |
|
||||
| `policy_ids` | array | List of DynamoAI policy IDs to apply (optional) | `DYNAMOAI_POLICY_IDS` env var (comma-separated) |
|
||||
| `model_id` | string | Model ID for tracking/logging | `DYNAMOAI_MODEL_ID` env var |
|
||||
| `mode` | string | When to run: `pre_call`, `post_call`, or `during_call` | Required |
|
||||
|
||||
## Observability
|
||||
|
||||
DynamoAI guardrail logs include:
|
||||
|
||||
- **guardrail_status**: `success`, `guardrail_intervened`, or `guardrail_failed_to_respond`
|
||||
- **guardrail_provider**: `dynamoai`
|
||||
- **guardrail_json_response**: Full API response with policy details
|
||||
- **duration**: Time taken for guardrail check
|
||||
- **start_time** and **end_time**: Timestamps
|
||||
|
||||
These logs are available through your configured LiteLLM logging callbacks.
|
||||
|
||||
## Error Handling
|
||||
|
||||
The guardrail handles errors gracefully:
|
||||
|
||||
- **API Failures**: Logs error and raises exception with status `guardrail_failed_to_respond`
|
||||
- **Policy Violations**: Raises `ValueError` with detailed violation information
|
||||
- **Invalid Configuration**: Raises `ValueError` on initialization if API key is missing
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- Only the `BLOCK` action is currently supported
|
||||
- `WARN`, `REDACT`, and `SANITIZE` actions are treated as success (pass through)
|
||||
|
||||
## Support
|
||||
|
||||
For more information about DynamoAI:
|
||||
- Website: [https://dynamo.ai](https://dynamo.ai)
|
||||
- Documentation: Contact DynamoAI for API documentation
|
||||
|
||||
@@ -43,6 +43,7 @@ const sidebars = {
|
||||
"proxy/guardrails/lakera_ai",
|
||||
"proxy/guardrails/model_armor",
|
||||
"proxy/guardrails/noma_security",
|
||||
"proxy/guardrails/dynamoai",
|
||||
"proxy/guardrails/openai_moderation",
|
||||
"proxy/guardrails/pangea",
|
||||
"proxy/guardrails/pillar_security",
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
from .dynamoai import DynamoAIGuardrails
|
||||
|
||||
__all__ = ["DynamoAIGuardrails"]
|
||||
|
||||
@@ -0,0 +1,520 @@
|
||||
# +-------------------------------------------------------------+
|
||||
#
|
||||
# Use DynamoAI Guardrails for your LLM calls
|
||||
# https://dynamo.ai
|
||||
#
|
||||
# +-------------------------------------------------------------+
|
||||
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import (
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
Dict,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Type,
|
||||
Union,
|
||||
)
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.caching import DualCache
|
||||
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.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.dynamoai import (
|
||||
DynamoAIProcessedResult,
|
||||
DynamoAIRequest,
|
||||
DynamoAIResponse,
|
||||
)
|
||||
from litellm.types.utils import GuardrailStatus, ModelResponseStream
|
||||
|
||||
GUARDRAIL_NAME = "dynamoai"
|
||||
|
||||
|
||||
class DynamoAIGuardrails(CustomGuardrail):
|
||||
"""
|
||||
DynamoAI Guardrails integration for LiteLLM.
|
||||
|
||||
Provides content moderation and policy enforcement using DynamoAI's guardrail API.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: str = "litellm_test",
|
||||
api_key: Optional[str] = None,
|
||||
api_base: Optional[str] = None,
|
||||
model_id: str = "",
|
||||
policy_ids: List[str] = [],
|
||||
**kwargs,
|
||||
):
|
||||
self.async_handler = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback
|
||||
)
|
||||
|
||||
# Set API configuration
|
||||
self.api_key = api_key or os.getenv("DYNAMOAI_API_KEY")
|
||||
if not self.api_key:
|
||||
raise ValueError(
|
||||
"DynamoAI API key is required. Set DYNAMOAI_API_KEY environment variable or pass api_key parameter."
|
||||
)
|
||||
|
||||
self.api_base = api_base or os.getenv(
|
||||
"DYNAMOAI_API_BASE", "https://api.dynamo.ai"
|
||||
)
|
||||
self.api_url = f"{self.api_base}/v1/moderation/analyze/"
|
||||
|
||||
# Model ID for tracking/logging purposes
|
||||
self.model_id = model_id or os.getenv("DYNAMOAI_MODEL_ID", "")
|
||||
|
||||
# Policy IDs - get from parameter, env var, or use empty list
|
||||
env_policy_ids = os.getenv("DYNAMOAI_POLICY_IDS", "")
|
||||
self.policy_ids = policy_ids or (env_policy_ids.split(",") if env_policy_ids else [])
|
||||
self.guardrail_name = guardrail_name
|
||||
self.guardrail_provider = "dynamoai"
|
||||
|
||||
# store kwargs as optional_params
|
||||
self.optional_params = kwargs
|
||||
|
||||
# Set supported event hooks
|
||||
if "supported_event_hooks" not in kwargs:
|
||||
kwargs["supported_event_hooks"] = [
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
GuardrailEventHooks.during_call,
|
||||
]
|
||||
|
||||
super().__init__(guardrail_name=guardrail_name, **kwargs)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"DynamoAI Guardrail initialized with guardrail_name=%s, model_id=%s",
|
||||
self.guardrail_name,
|
||||
self.model_id,
|
||||
)
|
||||
|
||||
async def _call_dynamoai_guardrails(
|
||||
self,
|
||||
messages: List[Dict[str, Any]],
|
||||
text_type: str = "input",
|
||||
request_data: Optional[dict] = None,
|
||||
) -> DynamoAIResponse:
|
||||
"""
|
||||
Call DynamoAI Guardrails API to analyze messages for policy violations.
|
||||
|
||||
Args:
|
||||
messages: List of messages to analyze
|
||||
text_type: Type of text being analyzed ("input" or "output")
|
||||
request_data: Optional request data for logging purposes
|
||||
|
||||
Returns:
|
||||
DynamoAIResponse: Response from the DynamoAI Guardrails API
|
||||
"""
|
||||
start_time = datetime.now()
|
||||
|
||||
payload: DynamoAIRequest = {
|
||||
"messages": messages,
|
||||
}
|
||||
|
||||
# Add optional fields if provided
|
||||
if self.policy_ids:
|
||||
payload["policyIds"] = self.policy_ids
|
||||
if self.model_id:
|
||||
payload["modelId"] = self.model_id
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {self.api_key}"
|
||||
}
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"DynamoAI request to %s with payload=%s",
|
||||
self.api_url,
|
||||
payload,
|
||||
)
|
||||
|
||||
try:
|
||||
response = await self.async_handler.post(
|
||||
url=self.api_url,
|
||||
json=dict(payload),
|
||||
headers=headers,
|
||||
)
|
||||
response.raise_for_status()
|
||||
response_json = response.json()
|
||||
|
||||
end_time = datetime.now()
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
# Add guardrail information to request trace
|
||||
if request_data:
|
||||
guardrail_status = self._determine_guardrail_status(response_json)
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response=response_json,
|
||||
request_data=request_data,
|
||||
guardrail_status=guardrail_status,
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=end_time.timestamp(),
|
||||
duration=duration,
|
||||
)
|
||||
|
||||
return response_json
|
||||
|
||||
except httpx.HTTPError as e:
|
||||
end_time = datetime.now()
|
||||
duration = (end_time - start_time).total_seconds()
|
||||
|
||||
verbose_proxy_logger.error(
|
||||
"DynamoAI API request failed: %s", str(e)
|
||||
)
|
||||
|
||||
# Add guardrail information with failure status
|
||||
if request_data:
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response={"error": str(e)},
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_failed_to_respond",
|
||||
start_time=start_time.timestamp(),
|
||||
end_time=end_time.timestamp(),
|
||||
duration=duration,
|
||||
)
|
||||
|
||||
raise
|
||||
|
||||
def _process_dynamoai_guardrails_response(
|
||||
self, response: DynamoAIResponse
|
||||
) -> DynamoAIProcessedResult:
|
||||
"""
|
||||
Process the response from the DynamoAI Guardrails API
|
||||
|
||||
Args:
|
||||
response: The response from the API with 'finalAction' and 'appliedPolicies' keys
|
||||
|
||||
Returns:
|
||||
DynamoAIProcessedResult: Processed response with detected violations
|
||||
"""
|
||||
final_action = response.get("finalAction", "NONE")
|
||||
applied_policies = response.get("appliedPolicies", [])
|
||||
|
||||
violations_detected: List[str] = []
|
||||
violation_details: Dict[str, Any] = {}
|
||||
|
||||
# For now, only handle BLOCK action
|
||||
if final_action == "BLOCK":
|
||||
for applied_policy in applied_policies:
|
||||
policy_info = applied_policy.get("policy", {})
|
||||
policy_outputs = applied_policy.get("outputs", {})
|
||||
|
||||
# Get policy name and action
|
||||
policy_name = policy_info.get("name", "unknown")
|
||||
|
||||
# Check for action in multiple places
|
||||
policy_action = (
|
||||
applied_policy.get("action") or
|
||||
(policy_outputs.get("action") if policy_outputs else None) or
|
||||
"NONE"
|
||||
)
|
||||
|
||||
# Only include policies with BLOCK action
|
||||
if policy_action == "BLOCK":
|
||||
violations_detected.append(policy_name)
|
||||
violation_details[policy_name] = {
|
||||
"policyId": policy_info.get("id"),
|
||||
"action": policy_action,
|
||||
"method": policy_info.get("method"),
|
||||
"description": policy_info.get("description"),
|
||||
"message": policy_outputs.get("message") if policy_outputs else None,
|
||||
}
|
||||
|
||||
return {
|
||||
"violations_detected": violations_detected,
|
||||
"violation_details": violation_details
|
||||
}
|
||||
|
||||
def _determine_guardrail_status(
|
||||
self, response_json: DynamoAIResponse
|
||||
) -> GuardrailStatus:
|
||||
"""
|
||||
Determine the guardrail status based on DynamoAI API response.
|
||||
|
||||
Returns:
|
||||
"success": Content allowed through with no violations (finalAction is NONE)
|
||||
"guardrail_intervened": Content blocked (finalAction is BLOCK)
|
||||
"guardrail_failed_to_respond": Technical error or API failure
|
||||
"""
|
||||
try:
|
||||
if not isinstance(response_json, dict):
|
||||
return "guardrail_failed_to_respond"
|
||||
|
||||
# Check for error in response
|
||||
if response_json.get("error"):
|
||||
return "guardrail_failed_to_respond"
|
||||
|
||||
final_action = response_json.get("finalAction", "NONE")
|
||||
|
||||
if final_action == "NONE":
|
||||
return "success"
|
||||
elif final_action == "BLOCK":
|
||||
return "guardrail_intervened"
|
||||
|
||||
# For now, treat other actions as success (WARN, REDACT, SANITIZE not implemented yet)
|
||||
return "success"
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
"Error determining DynamoAI guardrail status: %s", str(e)
|
||||
)
|
||||
return "guardrail_failed_to_respond"
|
||||
|
||||
def _create_error_message(self, processed_result: DynamoAIProcessedResult) -> str:
|
||||
"""
|
||||
Create a detailed error message from processed guardrail results.
|
||||
|
||||
Args:
|
||||
processed_result: Processed response with detected violations
|
||||
|
||||
Returns:
|
||||
Formatted error message string
|
||||
"""
|
||||
violations_detected = processed_result["violations_detected"]
|
||||
violation_details = processed_result["violation_details"]
|
||||
|
||||
error_message = f"Guardrail failed: {len(violations_detected)} violation(s) detected\n\n"
|
||||
|
||||
for policy_name in violations_detected:
|
||||
error_message += f"- {policy_name.upper()}:\n"
|
||||
details = violation_details.get(policy_name, {})
|
||||
|
||||
# Format violation details
|
||||
if details.get("action"):
|
||||
error_message += f" Action: {details['action']}\n"
|
||||
if details.get("method"):
|
||||
error_message += f" Method: {details['method']}\n"
|
||||
if details.get("description"):
|
||||
error_message += f" Description: {details['description']}\n"
|
||||
if details.get("message"):
|
||||
error_message += f" Message: {details['message']}\n"
|
||||
if details.get("policyId"):
|
||||
error_message += f" Policy ID: {details['policyId']}\n"
|
||||
error_message += "\n"
|
||||
|
||||
return error_message.strip()
|
||||
|
||||
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",
|
||||
"anthropic_messages",
|
||||
],
|
||||
) -> Union[Exception, str, dict, None]:
|
||||
"""
|
||||
Runs before the LLM API call
|
||||
Runs on only Input
|
||||
Use this if you want to MODIFY the input
|
||||
"""
|
||||
verbose_proxy_logger.debug("Running DynamoAI pre-call hook")
|
||||
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
|
||||
event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call
|
||||
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
|
||||
return data
|
||||
|
||||
_messages = data.get("messages")
|
||||
if _messages:
|
||||
result = await self._call_dynamoai_guardrails(
|
||||
messages=_messages,
|
||||
text_type="input",
|
||||
request_data=data,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Guardrails async_pre_call_hook result=%s", result)
|
||||
|
||||
# Process the guardrails response
|
||||
processed_result = self._process_dynamoai_guardrails_response(result)
|
||||
violations_detected = processed_result["violations_detected"]
|
||||
|
||||
# If any violations are detected, raise an error
|
||||
if violations_detected:
|
||||
error_message = self._create_error_message(processed_result)
|
||||
raise ValueError(error_message)
|
||||
|
||||
# Add guardrail to applied guardrails header
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
|
||||
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",
|
||||
"anthropic_messages",
|
||||
],
|
||||
):
|
||||
"""
|
||||
Runs in parallel to LLM API call
|
||||
Runs on only Input
|
||||
|
||||
This can NOT modify the input, only used to reject or accept a call before going to LLM API
|
||||
"""
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
|
||||
event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
|
||||
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
|
||||
return
|
||||
|
||||
_messages = data.get("messages")
|
||||
if _messages:
|
||||
result = await self._call_dynamoai_guardrails(
|
||||
messages=_messages,
|
||||
text_type="input",
|
||||
request_data=data,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Guardrails async_moderation_hook result=%s", result)
|
||||
|
||||
# Process the guardrails response
|
||||
processed_result = self._process_dynamoai_guardrails_response(result)
|
||||
violations_detected = processed_result["violations_detected"]
|
||||
|
||||
# If any violations are detected, raise an error
|
||||
if violations_detected:
|
||||
error_message = self._create_error_message(processed_result)
|
||||
raise ValueError(error_message)
|
||||
|
||||
# Add guardrail to applied guardrails header
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
async def async_post_call_success_hook(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response,
|
||||
):
|
||||
"""
|
||||
Runs on response from LLM API call
|
||||
|
||||
It can be used to reject a response
|
||||
|
||||
Uses DynamoAI guardrails to check the response for policy violations
|
||||
"""
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
if (
|
||||
self.should_run_guardrail(
|
||||
data=data, event_type=GuardrailEventHooks.post_call
|
||||
)
|
||||
is not True
|
||||
):
|
||||
return
|
||||
|
||||
verbose_proxy_logger.debug("async_post_call_success_hook response=%s", response)
|
||||
|
||||
# Check if the ModelResponse has text content in its choices
|
||||
# to avoid sending empty content to DynamoAI (e.g., during tool calls)
|
||||
if isinstance(response, litellm.ModelResponse):
|
||||
has_text_content = False
|
||||
dynamoai_messages: List[Dict[str, Any]] = []
|
||||
|
||||
for choice in response.choices:
|
||||
if isinstance(choice, litellm.Choices):
|
||||
if choice.message.content and isinstance(choice.message.content, str):
|
||||
has_text_content = True
|
||||
dynamoai_messages.append({
|
||||
"role": choice.message.role or "assistant",
|
||||
"content": choice.message.content
|
||||
})
|
||||
|
||||
if not has_text_content:
|
||||
verbose_proxy_logger.warning(
|
||||
"DynamoAI: not running guardrail. No output text in response"
|
||||
)
|
||||
return
|
||||
|
||||
if dynamoai_messages:
|
||||
result = await self._call_dynamoai_guardrails(
|
||||
messages=dynamoai_messages,
|
||||
text_type="output",
|
||||
request_data=data,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Guardrails async_post_call_success_hook result=%s", result)
|
||||
|
||||
# Process the guardrails response
|
||||
processed_result = self._process_dynamoai_guardrails_response(result)
|
||||
violations_detected = processed_result["violations_detected"]
|
||||
|
||||
# If any violations are detected, raise an error
|
||||
if violations_detected:
|
||||
error_message = self._create_error_message(processed_result)
|
||||
raise ValueError(error_message)
|
||||
|
||||
# Add guardrail to applied guardrails header
|
||||
add_guardrail_to_applied_guardrails_header(
|
||||
request_data=data, guardrail_name=self.guardrail_name
|
||||
)
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: Any,
|
||||
request_data: dict,
|
||||
) -> AsyncGenerator[ModelResponseStream, None]:
|
||||
"""
|
||||
Passes the entire stream to the guardrail
|
||||
|
||||
This is useful for guardrails that need to see the entire response, such as PII masking.
|
||||
|
||||
Triggered by mode: 'post_call'
|
||||
"""
|
||||
async for item in response:
|
||||
yield item
|
||||
|
||||
@staticmethod
|
||||
def get_config_model() -> Optional[Type[GuardrailConfigModel]]:
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.dynamoai import (
|
||||
DynamoAIGuardrailConfigModel,
|
||||
)
|
||||
|
||||
return DynamoAIGuardrailConfigModel
|
||||
|
||||
@@ -0,0 +1,118 @@
|
||||
# Type definitions for DynamoAI Guardrails API
|
||||
|
||||
import enum
|
||||
from typing import Any, Dict, List, Literal, Optional, TypedDict
|
||||
|
||||
from pydantic import Field
|
||||
|
||||
from .base import GuardrailConfigModel
|
||||
|
||||
|
||||
class DynamoAIMessage(TypedDict):
|
||||
"""Message structure for DynamoAI API"""
|
||||
role: str
|
||||
content: str
|
||||
|
||||
class DynamoRequestMetadata(TypedDict):
|
||||
endUserId: Optional[str]
|
||||
|
||||
class DynamoTextType(str, enum.Enum):
|
||||
MODEL_INPUT = "MODEL_INPUT"
|
||||
MODEL_RESPONSE = "MODEL_RESPONSE"
|
||||
|
||||
|
||||
class PolicyMethod(str, enum.Enum):
|
||||
PII = "PII"
|
||||
TOXICITY = "TOXICITY"
|
||||
ALIGNMENT = "ALIGNMENT"
|
||||
HALLUCINATION = "HALLUCINATION"
|
||||
RAG_HALLUCINATION = "RAG_HALLUCINATION"
|
||||
|
||||
|
||||
class PolicyApplicableTo(str, enum.Enum):
|
||||
INPUT = "INPUT"
|
||||
OUTPUT = "OUTPUT"
|
||||
ALL = "ALL"
|
||||
|
||||
|
||||
class DynamoAIRequest(TypedDict, total=False):
|
||||
"""Request structure for DynamoAI /moderation/analyze endpoint"""
|
||||
messages: List[Dict[str, Any]]
|
||||
textType: Optional[DynamoTextType]
|
||||
policyIds: List[str]
|
||||
modelId: Optional[str]
|
||||
clientId: Optional[str]
|
||||
metadata: Optional[DynamoRequestMetadata]
|
||||
|
||||
|
||||
class PolicyInfo(TypedDict, total=False):
|
||||
"""Policy information from DynamoAI response"""
|
||||
id: str
|
||||
name: str
|
||||
description: str
|
||||
method: PolicyMethod
|
||||
action: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"]
|
||||
methodParams: Dict[str, Any]
|
||||
decisionParams: Dict[str, Any]
|
||||
applicableTo: PolicyApplicableTo
|
||||
created_at: str
|
||||
creatorId: str
|
||||
|
||||
|
||||
class PolicyOutputs(TypedDict, total=False):
|
||||
"""Outputs from the policy"""
|
||||
action: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"]
|
||||
message: Optional[str]
|
||||
|
||||
|
||||
class AppliedPolicyDto(TypedDict, total=False):
|
||||
"""Applied policy details from DynamoAI response"""
|
||||
policy: PolicyInfo
|
||||
outputs: Optional[Dict[str, Any]]
|
||||
action: Optional[str]
|
||||
|
||||
|
||||
class DynamoAIResponse(TypedDict, total=False):
|
||||
"""Response structure from DynamoAI /moderation/analyze endpoint"""
|
||||
text: str
|
||||
textType: DynamoTextType
|
||||
finalAction: Literal["BLOCK", "WARN", "REDACT", "SANITIZE", "NONE"]
|
||||
appliedPolicies: List[AppliedPolicyDto]
|
||||
error: Optional[str]
|
||||
|
||||
|
||||
class DynamoAIProcessedResult(TypedDict):
|
||||
"""Processed result from DynamoAI guardrail check"""
|
||||
violations_detected: List[str]
|
||||
violation_details: Dict[str, Any]
|
||||
|
||||
|
||||
|
||||
class DynamoAIGuardrailConfigModel(GuardrailConfigModel):
|
||||
"""Configuration model for DynamoAI Guardrails"""
|
||||
|
||||
api_key: Optional[str] = Field(
|
||||
default=None,
|
||||
description="API key for DynamoAI Guardrails. If not provided, the `DYNAMOAI_API_KEY` environment variable is checked.",
|
||||
)
|
||||
api_base: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Base URL for DynamoAI API. If not provided, the `DYNAMOAI_API_BASE` environment variable is checked, defaults to https://api.dynamo.ai",
|
||||
)
|
||||
policy_ids: Optional[List[str]] = Field(
|
||||
default=None,
|
||||
description="List of DynamoAI policy IDs to apply. If not provided, the `DYNAMOAI_POLICY_IDS` environment variable is checked (comma-separated).",
|
||||
)
|
||||
model_id: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Model ID for tracking/logging purposes. If not provided, the `DYNAMOAI_MODEL_ID` environment variable is checked.",
|
||||
)
|
||||
guardrail_name: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Name of the guardrail for identification in logs and traces.",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "DynamoAI Guardrails"
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""
|
||||
Test DynamoAI Guardrails integration
|
||||
"""
|
||||
import sys
|
||||
import os
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
from litellm.proxy.guardrails.guardrail_hooks.dynamoai import DynamoAIGuardrails
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.caching.caching import DualCache
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dynamoai_blocks_content_with_block_action():
|
||||
"""
|
||||
Test that DynamoAI guardrail blocks content when finalAction is BLOCK.
|
||||
"""
|
||||
# Create guardrail instance
|
||||
guardrail = DynamoAIGuardrails(
|
||||
guardrail_name="test-dynamoai",
|
||||
api_key="test-api-key",
|
||||
api_base="https://api.dynamo.ai",
|
||||
)
|
||||
|
||||
# Mock the DynamoAI API response with BLOCK action
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"text": "This is harmful content",
|
||||
"textType": "MODEL_INPUT",
|
||||
"finalAction": "BLOCK",
|
||||
"appliedPolicies": [
|
||||
{
|
||||
"policy": {
|
||||
"id": "policy-123",
|
||||
"name": "Toxicity Policy",
|
||||
"description": "Blocks toxic content",
|
||||
"method": "TOXICITY",
|
||||
},
|
||||
"outputs": {
|
||||
"action": "BLOCK",
|
||||
"message": "Content contains toxic language"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
guardrail.async_handler.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "This is harmful content"}
|
||||
],
|
||||
}
|
||||
|
||||
# Mock should_run_guardrail to return True
|
||||
guardrail.should_run_guardrail = MagicMock(return_value=True)
|
||||
|
||||
# Test that the guardrail raises ValueError for blocked content
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
await guardrail.async_pre_call_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
call_type="completion",
|
||||
cache=MagicMock(spec=DualCache),
|
||||
)
|
||||
|
||||
# Verify the error message contains policy information
|
||||
error_message = str(exc_info.value)
|
||||
assert "Guardrail failed" in error_message
|
||||
assert "TOXICITY POLICY" in error_message.upper()
|
||||
assert "BLOCK" in error_message.upper()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_dynamoai_allows_content_with_none_action():
|
||||
"""
|
||||
Test that DynamoAI guardrail allows content when finalAction is NONE.
|
||||
"""
|
||||
# Create guardrail instance
|
||||
guardrail = DynamoAIGuardrails(
|
||||
guardrail_name="test-dynamoai",
|
||||
api_key="test-api-key",
|
||||
api_base="https://api.dynamo.ai",
|
||||
)
|
||||
|
||||
# Mock the DynamoAI API response with NONE action (no violations)
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"text": "Hello, how are you?",
|
||||
"textType": "MODEL_INPUT",
|
||||
"finalAction": "NONE",
|
||||
"appliedPolicies": []
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
guardrail.async_handler.post = AsyncMock(return_value=mock_response)
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello, how are you?"}
|
||||
],
|
||||
}
|
||||
|
||||
# Mock should_run_guardrail to return True
|
||||
guardrail.should_run_guardrail = MagicMock(return_value=True)
|
||||
|
||||
# Test that the guardrail allows the content (no exception raised)
|
||||
result = await guardrail.async_pre_call_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
call_type="completion",
|
||||
cache=MagicMock(spec=DualCache),
|
||||
)
|
||||
|
||||
# Should return the request data unchanged
|
||||
assert result == request_data
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -103,37 +103,40 @@ async def test_callback():
|
||||
}
|
||||
|
||||
# Test violation detection
|
||||
mock_response = Response(
|
||||
json={
|
||||
"violations_detected": True,
|
||||
"deputies": {
|
||||
"jailbreak": True,
|
||||
"custom-policies": False,
|
||||
"sexual": False,
|
||||
"hate": False,
|
||||
"illegality": False,
|
||||
"violence": False,
|
||||
"pattern-detection": False,
|
||||
},
|
||||
"deputies_predictions": {
|
||||
"jailbreak": 0.923,
|
||||
"custom-policies": 0.234,
|
||||
"sexual": 0.145,
|
||||
"hate": 0.156,
|
||||
"illegality": 0.167,
|
||||
"violence": 0.178,
|
||||
"pattern-detection": 0.189,
|
||||
},
|
||||
"findings": {
|
||||
"jailbreak": [{"action": "BLOCK", "severity": "HIGH"}]
|
||||
}
|
||||
},
|
||||
status_code=200,
|
||||
request=Request(
|
||||
method="POST", url="https://server.lasso.security/gateway/v2/classify"
|
||||
),
|
||||
)
|
||||
mock_response.raise_for_status = lambda: None
|
||||
|
||||
with pytest.raises(HTTPException) as excinfo:
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=Response(
|
||||
json={
|
||||
"deputies": {
|
||||
"jailbreak": True,
|
||||
"custom-policies": False,
|
||||
"sexual": False,
|
||||
"hate": False,
|
||||
"illegality": False,
|
||||
"violence": False,
|
||||
"pattern-detection": False,
|
||||
},
|
||||
"deputies_predictions": {
|
||||
"jailbreak": 0.923,
|
||||
"custom-policies": 0.234,
|
||||
"sexual": 0.145,
|
||||
"hate": 0.156,
|
||||
"illegality": 0.167,
|
||||
"violence": 0.178,
|
||||
"pattern-detection": 0.189,
|
||||
},
|
||||
"violations_detected": True,
|
||||
},
|
||||
status_code=200,
|
||||
request=Request(
|
||||
method="POST", url="https://server.lasso.security/gateway/v1/chat"
|
||||
),
|
||||
),
|
||||
):
|
||||
with patch.object(lasso_guardrail.async_handler, "post", return_value=mock_response):
|
||||
await lasso_guardrail.async_pre_call_hook(
|
||||
data=data,
|
||||
cache=DualCache(),
|
||||
@@ -146,36 +149,37 @@ async def test_callback():
|
||||
assert "jailbreak" in str(excinfo.value.detail)
|
||||
|
||||
# Test no violation
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
return_value=Response(
|
||||
json={
|
||||
"deputies": {
|
||||
"jailbreak": False,
|
||||
"custom-policies": False,
|
||||
"sexual": False,
|
||||
"hate": False,
|
||||
"illegality": False,
|
||||
"violence": False,
|
||||
"pattern-detection": False,
|
||||
},
|
||||
"deputies_predictions": {
|
||||
"jailbreak": 0.123,
|
||||
"custom-policies": 0.234,
|
||||
"sexual": 0.145,
|
||||
"hate": 0.156,
|
||||
"illegality": 0.167,
|
||||
"violence": 0.178,
|
||||
"pattern-detection": 0.189,
|
||||
},
|
||||
"violations_detected": False,
|
||||
mock_response_no_violation = Response(
|
||||
json={
|
||||
"violations_detected": False,
|
||||
"deputies": {
|
||||
"jailbreak": False,
|
||||
"custom-policies": False,
|
||||
"sexual": False,
|
||||
"hate": False,
|
||||
"illegality": False,
|
||||
"violence": False,
|
||||
"pattern-detection": False,
|
||||
},
|
||||
status_code=200,
|
||||
request=Request(
|
||||
method="POST", url="https://server.lasso.security/gateway/v1/chat"
|
||||
),
|
||||
"deputies_predictions": {
|
||||
"jailbreak": 0.123,
|
||||
"custom-policies": 0.234,
|
||||
"sexual": 0.145,
|
||||
"hate": 0.156,
|
||||
"illegality": 0.167,
|
||||
"violence": 0.178,
|
||||
"pattern-detection": 0.189,
|
||||
},
|
||||
"findings": {}
|
||||
},
|
||||
status_code=200,
|
||||
request=Request(
|
||||
method="POST", url="https://server.lasso.security/gateway/v2/classify"
|
||||
),
|
||||
):
|
||||
)
|
||||
mock_response_no_violation.raise_for_status = lambda: None
|
||||
|
||||
with patch.object(lasso_guardrail.async_handler, "post", return_value=mock_response_no_violation):
|
||||
result = await lasso_guardrail.async_pre_call_hook(
|
||||
data=data,
|
||||
cache=DualCache(),
|
||||
@@ -231,10 +235,7 @@ async def test_api_error_handling():
|
||||
}
|
||||
|
||||
# Test handling of connection error
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
side_effect=Exception("Connection error"),
|
||||
):
|
||||
with patch.object(lasso_guardrail.async_handler, "post", side_effect=Exception("Connection error")):
|
||||
# Expect the guardrail to raise a LassoGuardrailAPIError
|
||||
with pytest.raises(LassoGuardrailAPIError) as excinfo:
|
||||
await lasso_guardrail.async_pre_call_hook(
|
||||
@@ -249,10 +250,7 @@ async def test_api_error_handling():
|
||||
assert "Connection error" in str(excinfo.value)
|
||||
|
||||
# Test with a different error message
|
||||
with patch(
|
||||
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
||||
side_effect=Exception("API timeout"),
|
||||
):
|
||||
with patch.object(lasso_guardrail.async_handler, "post", side_effect=Exception("API timeout")):
|
||||
# Expect the guardrail to raise a LassoGuardrailAPIError
|
||||
with pytest.raises(LassoGuardrailAPIError) as excinfo:
|
||||
await lasso_guardrail.async_pre_call_hook(
|
||||
|
||||
Reference in New Issue
Block a user