(feat) Generic Guardrail API - allows guardrail providers to add INSTANT support for LiteLLM w/out PR to repo (#17175)
* feat(generic_guardrail_api.py): new generic api for guardrails Allows guardrail providers to work with litellm for guardrails without needing to make a PR to LiteLLM * docs(generic_guardrail_api.md): document new generic guardrail api * Fix: Improve PII detection and guardrail API integration Co-authored-by: krrishdholakia <krrishdholakia@gmail.com> * feat: correctly extract raw request from guardrail api * docs(generic_guardrail_api.md): document this is a beta feature --------- Co-authored-by: Cursor Agent <cursoragent@cursor.com>
This commit is contained in:
Krish Dholakia
GitHub
parent
f434ca61ec
commit
b6d6f834e0
564
cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py
Normal file
564
cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py
Normal file
@ -0,0 +1,564 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Mock Bedrock Guardrail API Server
|
||||
|
||||
This is a FastAPI server that mimics the AWS Bedrock Guardrail API for testing purposes.
|
||||
It follows the same API spec as the real Bedrock guardrail endpoint.
|
||||
|
||||
Usage:
|
||||
python mock_bedrock_guardrail_server.py
|
||||
|
||||
The server will start on http://localhost:8080
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import Any, Dict, List, Literal, Optional
|
||||
|
||||
from fastapi import Depends, FastAPI, Header, HTTPException, status
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# ============================================================================
|
||||
# Request/Response Models (matching Bedrock API spec)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class BedrockTextContent(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class BedrockContentItem(BaseModel):
|
||||
text: BedrockTextContent
|
||||
|
||||
|
||||
class BedrockRequest(BaseModel):
|
||||
source: Literal["INPUT", "OUTPUT"]
|
||||
content: List[BedrockContentItem] = Field(default_factory=list)
|
||||
|
||||
|
||||
class BedrockGuardrailOutput(BaseModel):
|
||||
text: Optional[str] = None
|
||||
|
||||
|
||||
class TopicPolicyItem(BaseModel):
|
||||
name: str
|
||||
type: str
|
||||
action: Literal["BLOCKED", "NONE"]
|
||||
|
||||
|
||||
class TopicPolicy(BaseModel):
|
||||
topics: List[TopicPolicyItem] = Field(default_factory=list)
|
||||
|
||||
|
||||
class ContentFilterItem(BaseModel):
|
||||
type: str
|
||||
confidence: str
|
||||
action: Literal["BLOCKED", "NONE"]
|
||||
|
||||
|
||||
class ContentPolicy(BaseModel):
|
||||
filters: List[ContentFilterItem] = Field(default_factory=list)
|
||||
|
||||
|
||||
class CustomWord(BaseModel):
|
||||
match: str
|
||||
action: Literal["BLOCKED", "NONE"]
|
||||
|
||||
|
||||
class WordPolicy(BaseModel):
|
||||
customWords: List[CustomWord] = Field(default_factory=list)
|
||||
managedWordLists: List[Dict[str, Any]] = Field(default_factory=list)
|
||||
|
||||
|
||||
class PiiEntity(BaseModel):
|
||||
type: str
|
||||
match: str
|
||||
action: Literal["BLOCKED", "ANONYMIZED", "NONE"]
|
||||
|
||||
|
||||
class RegexMatch(BaseModel):
|
||||
name: str
|
||||
match: str
|
||||
regex: str
|
||||
action: Literal["BLOCKED", "ANONYMIZED", "NONE"]
|
||||
|
||||
|
||||
class SensitiveInformationPolicy(BaseModel):
|
||||
piiEntities: List[PiiEntity] = Field(default_factory=list)
|
||||
regexes: List[RegexMatch] = Field(default_factory=list)
|
||||
|
||||
|
||||
class ContextualGroundingFilter(BaseModel):
|
||||
type: str
|
||||
threshold: float
|
||||
score: float
|
||||
action: Literal["BLOCKED", "NONE"]
|
||||
|
||||
|
||||
class ContextualGroundingPolicy(BaseModel):
|
||||
filters: List[ContextualGroundingFilter] = Field(default_factory=list)
|
||||
|
||||
|
||||
class Assessment(BaseModel):
|
||||
topicPolicy: Optional[TopicPolicy] = None
|
||||
contentPolicy: Optional[ContentPolicy] = None
|
||||
wordPolicy: Optional[WordPolicy] = None
|
||||
sensitiveInformationPolicy: Optional[SensitiveInformationPolicy] = None
|
||||
contextualGroundingPolicy: Optional[ContextualGroundingPolicy] = None
|
||||
|
||||
|
||||
class BedrockGuardrailResponse(BaseModel):
|
||||
usage: Dict[str, int] = Field(
|
||||
default_factory=lambda: {"topicPolicyUnits": 1, "contentPolicyUnits": 1}
|
||||
)
|
||||
action: Literal["NONE", "GUARDRAIL_INTERVENED"] = "NONE"
|
||||
outputs: List[BedrockGuardrailOutput] = Field(default_factory=list)
|
||||
assessments: List[Assessment] = Field(default_factory=list)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Mock Guardrail Configuration
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class GuardrailConfig(BaseModel):
|
||||
"""Configuration for mock guardrail behavior"""
|
||||
|
||||
blocked_words: List[str] = Field(
|
||||
default_factory=lambda: ["offensive", "inappropriate", "badword"]
|
||||
)
|
||||
blocked_topics: List[str] = Field(default_factory=lambda: ["violence", "illegal"])
|
||||
pii_patterns: Dict[str, str] = Field(
|
||||
default_factory=lambda: {
|
||||
"EMAIL": r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b",
|
||||
"PHONE": r"\b\d{3}[-.]?\d{3}[-.]?\d{4}\b",
|
||||
"SSN": r"\b\d{3}-\d{2}-\d{4}\b",
|
||||
"CREDIT_CARD": r"\b\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b",
|
||||
}
|
||||
)
|
||||
anonymize_pii: bool = True # If True, ANONYMIZE PII; if False, BLOCK it
|
||||
bearer_token: str = "mock-bedrock-token-12345"
|
||||
|
||||
|
||||
# Global config
|
||||
GUARDRAIL_CONFIG = GuardrailConfig()
|
||||
|
||||
# ============================================================================
|
||||
# FastAPI App Setup
|
||||
# ============================================================================
|
||||
|
||||
app = FastAPI(
|
||||
title="Mock Bedrock Guardrail API",
|
||||
description="Mock server mimicking AWS Bedrock Guardrail API",
|
||||
version="1.0.0",
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Authentication
|
||||
# ============================================================================
|
||||
|
||||
|
||||
async def verify_bearer_token(authorization: Optional[str] = Header(None)) -> str:
|
||||
"""
|
||||
Verify the Bearer token from the Authorization header.
|
||||
|
||||
Args:
|
||||
authorization: The Authorization header value
|
||||
|
||||
Returns:
|
||||
The token if valid
|
||||
|
||||
Raises:
|
||||
HTTPException: If token is missing or invalid
|
||||
"""
|
||||
if authorization is None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Missing Authorization header",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
# Check if it's a Bearer token
|
||||
parts = authorization.split()
|
||||
print(f"parts: {parts}")
|
||||
if len(parts) != 2 or parts[0].lower() != "bearer":
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
detail="Invalid Authorization header format. Expected: Bearer <token>",
|
||||
headers={"WWW-Authenticate": "Bearer"},
|
||||
)
|
||||
|
||||
token = parts[1]
|
||||
|
||||
# Verify token
|
||||
if token != GUARDRAIL_CONFIG.bearer_token:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail="Invalid bearer token",
|
||||
)
|
||||
|
||||
return token
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Guardrail Logic
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def check_blocked_words(text: str) -> Optional[WordPolicy]:
|
||||
"""Check if text contains blocked words"""
|
||||
found_words = []
|
||||
text_lower = text.lower()
|
||||
|
||||
for word in GUARDRAIL_CONFIG.blocked_words:
|
||||
if word.lower() in text_lower:
|
||||
found_words.append(CustomWord(match=word, action="BLOCKED"))
|
||||
|
||||
if found_words:
|
||||
return WordPolicy(customWords=found_words)
|
||||
return None
|
||||
|
||||
|
||||
def check_blocked_topics(text: str) -> Optional[TopicPolicy]:
|
||||
"""Check if text contains blocked topics"""
|
||||
found_topics = []
|
||||
text_lower = text.lower()
|
||||
|
||||
for topic in GUARDRAIL_CONFIG.blocked_topics:
|
||||
if topic.lower() in text_lower:
|
||||
found_topics.append(
|
||||
TopicPolicyItem(name=topic, type=topic.upper(), action="BLOCKED")
|
||||
)
|
||||
|
||||
if found_topics:
|
||||
return TopicPolicy(topics=found_topics)
|
||||
return None
|
||||
|
||||
|
||||
def check_pii(text: str) -> tuple[Optional[SensitiveInformationPolicy], str]:
|
||||
"""
|
||||
Check for PII in text and return policy + anonymized text
|
||||
|
||||
Returns:
|
||||
Tuple of (SensitiveInformationPolicy or None, anonymized_text)
|
||||
"""
|
||||
pii_entities = []
|
||||
anonymized_text = text
|
||||
action = "ANONYMIZED" if GUARDRAIL_CONFIG.anonymize_pii else "BLOCKED"
|
||||
|
||||
for pii_type, pattern in GUARDRAIL_CONFIG.pii_patterns.items():
|
||||
try:
|
||||
# Compile the regex pattern with a timeout to prevent ReDoS attacks
|
||||
compiled_pattern = re.compile(pattern)
|
||||
matches = compiled_pattern.finditer(text)
|
||||
for match in matches:
|
||||
matched_text = match.group()
|
||||
pii_entities.append(
|
||||
PiiEntity(type=pii_type, match=matched_text, action=action)
|
||||
)
|
||||
|
||||
# Anonymize the text if configured
|
||||
if GUARDRAIL_CONFIG.anonymize_pii:
|
||||
anonymized_text = anonymized_text.replace(
|
||||
matched_text, f"[{pii_type}_REDACTED]"
|
||||
)
|
||||
except re.error:
|
||||
# Invalid regex pattern - skip it and log a warning
|
||||
print(f"Warning: Invalid regex pattern for PII type {pii_type}: {pattern}")
|
||||
continue
|
||||
|
||||
if pii_entities:
|
||||
return SensitiveInformationPolicy(piiEntities=pii_entities), anonymized_text
|
||||
|
||||
return None, text
|
||||
|
||||
|
||||
def process_guardrail_request(
|
||||
request: BedrockRequest,
|
||||
) -> tuple[BedrockGuardrailResponse, List[str]]:
|
||||
"""
|
||||
Process a guardrail request and return the response.
|
||||
|
||||
Returns:
|
||||
Tuple of (response, list of output texts)
|
||||
"""
|
||||
all_text_content = []
|
||||
output_texts = []
|
||||
|
||||
# Extract all text from content items
|
||||
for content_item in request.content:
|
||||
if content_item.text and content_item.text.text:
|
||||
all_text_content.append(content_item.text.text)
|
||||
|
||||
# Combine all text for analysis
|
||||
combined_text = " ".join(all_text_content)
|
||||
|
||||
# Initialize response
|
||||
response = BedrockGuardrailResponse()
|
||||
assessment = Assessment()
|
||||
has_intervention = False
|
||||
|
||||
# Check for blocked words
|
||||
word_policy = check_blocked_words(combined_text)
|
||||
if word_policy:
|
||||
assessment.wordPolicy = word_policy
|
||||
has_intervention = True
|
||||
|
||||
# Check for blocked topics
|
||||
topic_policy = check_blocked_topics(combined_text)
|
||||
if topic_policy:
|
||||
assessment.topicPolicy = topic_policy
|
||||
has_intervention = True
|
||||
|
||||
# Check for PII
|
||||
for text in all_text_content:
|
||||
pii_policy, anonymized_text = check_pii(text)
|
||||
if pii_policy:
|
||||
assessment.sensitiveInformationPolicy = pii_policy
|
||||
if GUARDRAIL_CONFIG.anonymize_pii:
|
||||
# If anonymizing, we don't block, we modify the text
|
||||
output_texts.append(anonymized_text)
|
||||
has_intervention = True
|
||||
else:
|
||||
# If not anonymizing PII, we block it
|
||||
output_texts.append(text)
|
||||
has_intervention = True
|
||||
else:
|
||||
output_texts.append(text)
|
||||
|
||||
# Build response
|
||||
if has_intervention:
|
||||
response.action = "GUARDRAIL_INTERVENED"
|
||||
# Only add assessment if there were interventions
|
||||
response.assessments = [assessment]
|
||||
|
||||
# Add outputs (modified or original text)
|
||||
response.outputs = [BedrockGuardrailOutput(text=txt) for txt in output_texts]
|
||||
|
||||
return response, output_texts
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# API Endpoints
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@app.get("/")
|
||||
async def root():
|
||||
"""Health check endpoint"""
|
||||
return {
|
||||
"service": "Mock Bedrock Guardrail API",
|
||||
"status": "running",
|
||||
"endpoint_format": "/guardrail/{guardrailIdentifier}/version/{guardrailVersion}/apply",
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
async def health():
|
||||
"""Health check endpoint"""
|
||||
return {"status": "healthy"}
|
||||
|
||||
|
||||
@app.post(
|
||||
"/guardrail/{guardrailIdentifier}/version/{guardrailVersion}/apply",
|
||||
response_model=BedrockGuardrailResponse,
|
||||
)
|
||||
async def apply_guardrail(
|
||||
guardrailIdentifier: str,
|
||||
guardrailVersion: str,
|
||||
request: BedrockRequest,
|
||||
token: str = Depends(verify_bearer_token),
|
||||
) -> BedrockGuardrailResponse:
|
||||
"""
|
||||
Apply guardrail to input or output content.
|
||||
|
||||
This endpoint mimics the AWS Bedrock ApplyGuardrail API.
|
||||
|
||||
Args:
|
||||
guardrailIdentifier: The guardrail ID
|
||||
guardrailVersion: The guardrail version
|
||||
request: The guardrail request containing content to analyze
|
||||
token: Bearer token (verified by dependency)
|
||||
|
||||
Returns:
|
||||
BedrockGuardrailResponse with analysis results
|
||||
"""
|
||||
# Process the request
|
||||
response, output_texts = process_guardrail_request(request)
|
||||
|
||||
# Log the request (optional, for debugging)
|
||||
print(f"Guardrail applied: {guardrailIdentifier} v{guardrailVersion}")
|
||||
print(f"Source: {request.source}")
|
||||
print(f"Action: {response.action}")
|
||||
|
||||
return response
|
||||
|
||||
|
||||
"""
|
||||
LiteLLM exposes a basic guardrail API with the text extracted from the request and sent to the guardrail API, as well as the received request body for any further processing.
|
||||
|
||||
This works across all LiteLLM endpoints (completion, anthropic /v1/messages, responses api, image generation, embedding, etc.)
|
||||
|
||||
This makes it easy to support your own guardrail API without having to make a PR to LiteLLM.
|
||||
|
||||
LiteLLM supports passing any provider specific params from LiteLLM config.yaml to the guardrail API.
|
||||
|
||||
Example:
|
||||
|
||||
```yaml
|
||||
guardrails:
|
||||
- guardrail_name: "bedrock-content-guard"
|
||||
litellm_params:
|
||||
guardrail: generic_guardrail_api
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/GUARDRAIL_API_KEY
|
||||
api_base: os.environ/GUARDRAIL_API_BASE
|
||||
additional_provider_specific_params:
|
||||
api_version: os.environ/GUARDRAIL_API_VERSION # additional provider specific params
|
||||
```
|
||||
|
||||
This is a beta API. Please help us improve it.
|
||||
"""
|
||||
|
||||
|
||||
class LitellmBasicGuardrailRequest(BaseModel):
|
||||
text: str
|
||||
request_body: Dict[str, Any] = Field(default_factory=dict)
|
||||
additional_provider_specific_params: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class LitellmBasicGuardrailResponse(BaseModel):
|
||||
action: Literal[
|
||||
"BLOCKED", "NONE", "GUARDRAIL_INTERVENED"
|
||||
] # BLOCKED = litellm will raise an error, NONE = litellm will continue, GUARDRAIL_INTERVENED = litellm will continue, but the text was modified by the guardrail
|
||||
blocked_reason: Optional[str] = None # only if action is BLOCKED, otherwise None
|
||||
text: Optional[str] = None
|
||||
|
||||
|
||||
@app.post(
|
||||
"/beta/litellm_basic_guardrail_api",
|
||||
response_model=LitellmBasicGuardrailResponse,
|
||||
)
|
||||
async def beta_litellm_basic_guardrail_api(
|
||||
request: LitellmBasicGuardrailRequest,
|
||||
) -> LitellmBasicGuardrailResponse:
|
||||
"""
|
||||
Apply guardrail to input or output content.
|
||||
|
||||
This endpoint mimics the AWS Bedrock ApplyGuardrail API.
|
||||
|
||||
Args:
|
||||
request: The guardrail request containing content to analyze
|
||||
token: Bearer token (verified by dependency)
|
||||
|
||||
Returns:
|
||||
LitellmBasicGuardrailResponse with analysis results
|
||||
"""
|
||||
print(f"request: {request}")
|
||||
if "ishaan" in request.text.lower():
|
||||
return LitellmBasicGuardrailResponse(
|
||||
action="BLOCKED", blocked_reason="Ishaan is not allowed"
|
||||
)
|
||||
elif "pii_value" in request.text:
|
||||
return LitellmBasicGuardrailResponse(
|
||||
action="GUARDRAIL_INTERVENED",
|
||||
text=request.text.replace("pii_value", "pii_value_redacted"),
|
||||
)
|
||||
return LitellmBasicGuardrailResponse(action="NONE")
|
||||
|
||||
|
||||
@app.post("/config/update")
|
||||
async def update_config(
|
||||
config: GuardrailConfig, token: str = Depends(verify_bearer_token)
|
||||
):
|
||||
"""
|
||||
Update the guardrail configuration.
|
||||
|
||||
This is a testing endpoint to modify the mock guardrail behavior.
|
||||
|
||||
Args:
|
||||
config: New guardrail configuration
|
||||
token: Bearer token (verified by dependency)
|
||||
|
||||
Returns:
|
||||
Updated configuration
|
||||
"""
|
||||
global GUARDRAIL_CONFIG
|
||||
GUARDRAIL_CONFIG = config
|
||||
return {"status": "updated", "config": GUARDRAIL_CONFIG}
|
||||
|
||||
|
||||
@app.get("/config")
|
||||
async def get_config(token: str = Depends(verify_bearer_token)):
|
||||
"""
|
||||
Get the current guardrail configuration.
|
||||
|
||||
Args:
|
||||
token: Bearer token (verified by dependency)
|
||||
|
||||
Returns:
|
||||
Current configuration
|
||||
"""
|
||||
return GUARDRAIL_CONFIG
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Error Handlers
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@app.exception_handler(HTTPException)
|
||||
async def http_exception_handler(request, exc: HTTPException):
|
||||
"""Custom error handler for HTTP exceptions"""
|
||||
return JSONResponse(
|
||||
status_code=exc.status_code,
|
||||
content={"error": exc.detail},
|
||||
headers=exc.headers,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Main
|
||||
# ============================================================================
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
# Get configuration from environment
|
||||
host = os.getenv("MOCK_BEDROCK_HOST", "0.0.0.0")
|
||||
port = int(os.getenv("MOCK_BEDROCK_PORT", "8080"))
|
||||
bearer_token = os.getenv("MOCK_BEDROCK_TOKEN", "mock-bedrock-token-12345")
|
||||
|
||||
# Update config with environment token
|
||||
GUARDRAIL_CONFIG.bearer_token = bearer_token
|
||||
|
||||
print("=" * 80)
|
||||
print("Mock Bedrock Guardrail API Server")
|
||||
print("=" * 80)
|
||||
print(f"Server starting on: http://{host}:{port}")
|
||||
print(f"Bearer Token: {bearer_token}")
|
||||
print(f"Endpoint: POST /guardrail/{{id}}/version/{{version}}/apply")
|
||||
print("=" * 80)
|
||||
print("\nExample curl command:")
|
||||
print(
|
||||
f"""
|
||||
curl -X POST "http://{host}:{port}/guardrail/test-guardrail/version/1/apply" \\
|
||||
-H "Authorization: Bearer {bearer_token}" \\
|
||||
-H "Content-Type: application/json" \\
|
||||
-d '{{
|
||||
"source": "INPUT",
|
||||
"content": [
|
||||
{{
|
||||
"text": {{
|
||||
"text": "Hello, my email is test@example.com"
|
||||
}}
|
||||
}}
|
||||
]
|
||||
}}'
|
||||
"""
|
||||
)
|
||||
print("=" * 80)
|
||||
|
||||
uvicorn.run(app, host=host, port=port)
|
||||
160
docs/my-website/docs/adding_provider/generic_guardrail_api.md
Normal file
160
docs/my-website/docs/adding_provider/generic_guardrail_api.md
Normal file
@ -0,0 +1,160 @@
|
||||
# [BETA] Generic Guardrail API - Integrate Without a PR
|
||||
|
||||
## The Problem
|
||||
|
||||
As a guardrail provider, integrating with LiteLLM traditionally requires:
|
||||
- Making a PR to the LiteLLM repository
|
||||
- Waiting for review and merge
|
||||
- Maintaining provider-specific code in LiteLLM's codebase
|
||||
- Updating the integration for changes to your API
|
||||
|
||||
## The Solution
|
||||
|
||||
The **Generic Guardrail API** lets you integrate with LiteLLM **instantly** by implementing a simple API endpoint. No PR required.
|
||||
|
||||
### Key Benefits
|
||||
|
||||
1. **No PR Needed** - Deploy and integrate immediately
|
||||
2. **Universal Support** - Works across ALL LiteLLM endpoints (chat, embeddings, image generation, etc.)
|
||||
3. **Simple Contract** - One endpoint, three response types
|
||||
4. **Custom Parameters** - Pass provider-specific params via config
|
||||
5. **Full Control** - You own and maintain your guardrail API
|
||||
|
||||
## How It Works
|
||||
|
||||
1. LiteLLM extracts text from any request (chat messages, embeddings, image prompts, etc.)
|
||||
2. Sends extracted text + original request to your API endpoint
|
||||
3. Your API responds with: `BLOCKED`, `NONE`, or `GUARDRAIL_INTERVENED`
|
||||
4. LiteLLM enforces the decision
|
||||
|
||||
## API Contract
|
||||
|
||||
### Endpoint
|
||||
|
||||
Implement `POST /beta/litellm_basic_guardrail_api`
|
||||
|
||||
### Request Format
|
||||
|
||||
```json
|
||||
{
|
||||
"text": "extracted text from the request",
|
||||
"request_body": {}, // full original request for context
|
||||
"additional_provider_specific_params": {
|
||||
// your custom params from config
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### Response Format
|
||||
|
||||
```json
|
||||
{
|
||||
"action": "BLOCKED" | "NONE" | "GUARDRAIL_INTERVENED",
|
||||
"blocked_reason": "why content was blocked", // required if action=BLOCKED
|
||||
"text": "modified text" // required if action=GUARDRAIL_INTERVENED
|
||||
}
|
||||
```
|
||||
|
||||
**Actions:**
|
||||
- `BLOCKED` - LiteLLM raises error and blocks request
|
||||
- `NONE` - Request proceeds unchanged
|
||||
- `GUARDRAIL_INTERVENED` - Request proceeds with modified text
|
||||
|
||||
## LiteLLM Configuration
|
||||
|
||||
Add to `config.yaml`:
|
||||
|
||||
```yaml
|
||||
litellm_settings:
|
||||
guardrails:
|
||||
- guardrail_name: "my-guardrail"
|
||||
litellm_params:
|
||||
guardrail: generic_guardrail_api
|
||||
mode: pre_call # or post_call, during_call
|
||||
api_base: https://your-guardrail-api.com
|
||||
api_key: os.environ/YOUR_GUARDRAIL_API_KEY # optional
|
||||
additional_provider_specific_params:
|
||||
# your custom parameters
|
||||
threshold: 0.8
|
||||
language: "en"
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
Users apply your guardrail by name:
|
||||
|
||||
```python
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
guardrails=["my-guardrail"]
|
||||
)
|
||||
```
|
||||
|
||||
Or with dynamic parameters:
|
||||
|
||||
```python
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-4",
|
||||
messages=[{"role": "user", "content": "hello"}],
|
||||
guardrails=[{
|
||||
"my-guardrail": {
|
||||
"extra_body": {
|
||||
"custom_threshold": 0.9
|
||||
}
|
||||
}
|
||||
}]
|
||||
)
|
||||
```
|
||||
|
||||
## Implementation Example
|
||||
|
||||
See [mock_bedrock_guardrail_server.py](https://github.com/BerriAI/litellm/blob/main/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py) for a complete reference implementation.
|
||||
|
||||
**Minimal FastAPI example:**
|
||||
|
||||
```python
|
||||
from fastapi import FastAPI
|
||||
from pydantic import BaseModel
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
class GuardrailRequest(BaseModel):
|
||||
text: str
|
||||
request_body: dict
|
||||
additional_provider_specific_params: dict
|
||||
|
||||
class GuardrailResponse(BaseModel):
|
||||
action: str # BLOCKED, NONE, or GUARDRAIL_INTERVENED
|
||||
blocked_reason: str | None = None
|
||||
text: str | None = None
|
||||
|
||||
@app.post("/beta/litellm_basic_guardrail_api")
|
||||
async def apply_guardrail(request: GuardrailRequest):
|
||||
# Your guardrail logic here
|
||||
if "badword" in request.text.lower():
|
||||
return GuardrailResponse(
|
||||
action="BLOCKED",
|
||||
blocked_reason="Content contains prohibited terms"
|
||||
)
|
||||
|
||||
return GuardrailResponse(action="NONE")
|
||||
```
|
||||
|
||||
## When to Use This
|
||||
|
||||
✅ **Use Generic Guardrail API when:**
|
||||
- You want instant integration without waiting for PRs
|
||||
- You maintain your own guardrail service
|
||||
- You need full control over updates and features
|
||||
- You want to support all LiteLLM endpoints automatically
|
||||
|
||||
❌ **Make a PR when:**
|
||||
- You want deeper integration with LiteLLM internals
|
||||
- Your guardrail requires complex LiteLLM-specific logic
|
||||
- You want to be featured as a built-in provider
|
||||
|
||||
## Questions?
|
||||
|
||||
This is a **beta API**. We're actively improving it based on feedback. Open an issue or PR if you need additional capabilities.
|
||||
|
||||
@ -45,6 +45,7 @@ const sidebars = {
|
||||
type: "category",
|
||||
"label": "Contributing to Guardrails",
|
||||
items: [
|
||||
"adding_provider/generic_guardrail_api",
|
||||
"adding_provider/simple_guardrail_tutorial",
|
||||
"adding_provider/adding_guardrail_support",
|
||||
]
|
||||
|
||||
@ -16,4 +16,4 @@ callback_settings:
|
||||
callback_type: generic_api
|
||||
endpoint: https://webhook.site/efc57707-9018-478c-bdf1-2ffaabb2b315
|
||||
headers:
|
||||
Authorization: Bearer sk-1234
|
||||
Authorization: Bearer sk-1234
|
||||
|
||||
@ -0,0 +1,37 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
from .generic_guardrail_api import GenericGuardrailAPI
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
|
||||
import litellm
|
||||
|
||||
_generic_guardrail_api_callback = GenericGuardrailAPI(
|
||||
api_base=litellm_params.api_base,
|
||||
headers=getattr(litellm_params, "headers", None),
|
||||
additional_provider_specific_params=getattr(
|
||||
litellm_params, "additional_provider_specific_params", {}
|
||||
),
|
||||
guardrail_name=guardrail.get("guardrail_name", ""),
|
||||
event_hook=litellm_params.mode,
|
||||
default_on=litellm_params.default_on,
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(
|
||||
_generic_guardrail_api_callback
|
||||
)
|
||||
return _generic_guardrail_api_callback
|
||||
|
||||
|
||||
guardrail_initializer_registry = {
|
||||
SupportedGuardrailIntegrations.GENERIC_GUARDRAIL_API.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.GENERIC_GUARDRAIL_API.value: GenericGuardrailAPI,
|
||||
}
|
||||
@ -0,0 +1,52 @@
|
||||
# Example configuration for Generic Guardrail API
|
||||
|
||||
model_list:
|
||||
- model_name: gpt-4
|
||||
litellm_params:
|
||||
model: gpt-4
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
litellm_settings:
|
||||
guardrails:
|
||||
- guardrail_name: "my-generic-guardrail"
|
||||
litellm_params:
|
||||
guardrail: generic_guardrail_api
|
||||
mode: pre_call # Options: pre_call, post_call, during_call, [pre_call, post_call]
|
||||
api_key: os.environ/GENERIC_GUARDRAIL_API_KEY # Optional if using Bearer auth
|
||||
api_base: http://localhost:8080 # Required. Endpoint /beta/litellm_basic_guardrail_api is automatically appended
|
||||
default_on: false # Set to true to apply to all requests by default
|
||||
additional_provider_specific_params:
|
||||
# Any additional parameters your guardrail API needs
|
||||
api_version: "v1"
|
||||
custom_param: "value"
|
||||
|
||||
# Usage examples:
|
||||
|
||||
# 1. Apply guardrail to a specific request:
|
||||
# curl --location 'http://localhost:4000/chat/completions' \
|
||||
# --header 'Authorization: Bearer sk-1234' \
|
||||
# --header 'Content-Type: application/json' \
|
||||
# --data '{
|
||||
# "model": "gpt-4",
|
||||
# "messages": [{"role": "user", "content": "Test message"}],
|
||||
# "guardrails": ["my-generic-guardrail"]
|
||||
# }'
|
||||
|
||||
# 2. Apply guardrail with dynamic parameters:
|
||||
# curl --location 'http://localhost:4000/chat/completions' \
|
||||
# --header 'Authorization: Bearer sk-1234' \
|
||||
# --header 'Content-Type: application/json' \
|
||||
# --data '{
|
||||
# "model": "gpt-4",
|
||||
# "messages": [{"role": "user", "content": "Test message"}],
|
||||
# "guardrails": [
|
||||
# {
|
||||
# "my-generic-guardrail": {
|
||||
# "extra_body": {
|
||||
# "custom_threshold": 0.8
|
||||
# }
|
||||
# }
|
||||
# }
|
||||
# ]
|
||||
# }'
|
||||
|
||||
@ -0,0 +1,235 @@
|
||||
# +-------------------------------------------------------------+
|
||||
#
|
||||
# Use Generic Guardrail API for your LLM calls
|
||||
#
|
||||
# +-------------------------------------------------------------+
|
||||
# Thank you users! We ❤️ you! - Krrish & Ishaan
|
||||
|
||||
import os
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
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.types.guardrails import GuardrailEventHooks
|
||||
|
||||
GUARDRAIL_NAME = "generic_guardrail_api"
|
||||
|
||||
|
||||
class GenericGuardrailAPIRequest:
|
||||
"""Request model for the Generic Guardrail API"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
text: str,
|
||||
request_body: Dict[str, Any],
|
||||
additional_provider_specific_params: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
self.text = text
|
||||
self.request_body = request_body
|
||||
self.additional_provider_specific_params = (
|
||||
additional_provider_specific_params or {}
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"text": self.text,
|
||||
"request_body": self.request_body,
|
||||
"additional_provider_specific_params": self.additional_provider_specific_params,
|
||||
}
|
||||
|
||||
|
||||
class GenericGuardrailAPIResponse:
|
||||
"""Response model for the Generic Guardrail API"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
action: str,
|
||||
blocked_reason: Optional[str] = None,
|
||||
text: Optional[str] = None,
|
||||
):
|
||||
self.action = action
|
||||
self.blocked_reason = blocked_reason
|
||||
self.text = text
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse":
|
||||
return cls(
|
||||
action=data.get("action", "NONE"),
|
||||
blocked_reason=data.get("blocked_reason"),
|
||||
text=data.get("text"),
|
||||
)
|
||||
|
||||
|
||||
class GenericGuardrailAPI(CustomGuardrail):
|
||||
"""
|
||||
Generic Guardrail API integration for LiteLLM.
|
||||
|
||||
This integration allows you to use any guardrail API that follows the
|
||||
LiteLLM Basic Guardrail API spec without needing to write custom integration code.
|
||||
|
||||
The API should accept a POST request with:
|
||||
{
|
||||
"text": str,
|
||||
"request_body": dict,
|
||||
"additional_provider_specific_params": dict
|
||||
}
|
||||
|
||||
And return:
|
||||
{
|
||||
"action": "BLOCKED" | "NONE" | "GUARDRAIL_INTERVENED",
|
||||
"blocked_reason": str (optional, only if action is BLOCKED),
|
||||
"text": str (optional, modified text if action is GUARDRAIL_INTERVENED)
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
headers: Optional[Dict[str, Any]] = None,
|
||||
api_base: Optional[str] = None,
|
||||
additional_provider_specific_params: Optional[Dict[str, Any]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
self.async_handler = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.GuardrailCallback
|
||||
)
|
||||
self.headers = headers or {}
|
||||
base_url = api_base or os.environ.get("GENERIC_GUARDRAIL_API_BASE")
|
||||
|
||||
if not base_url:
|
||||
raise ValueError(
|
||||
"api_base is required for Generic Guardrail API. "
|
||||
"Set GENERIC_GUARDRAIL_API_BASE environment variable or pass it in litellm_params"
|
||||
)
|
||||
|
||||
# Append the endpoint path if not already present
|
||||
if not base_url.endswith("/beta/litellm_basic_guardrail_api"):
|
||||
base_url = base_url.rstrip("/")
|
||||
self.api_base = f"{base_url}/beta/litellm_basic_guardrail_api"
|
||||
else:
|
||||
self.api_base = base_url
|
||||
|
||||
self.additional_provider_specific_params = (
|
||||
additional_provider_specific_params or {}
|
||||
)
|
||||
|
||||
# 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__(**kwargs)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Generic Guardrail API initialized with api_base: %s", self.api_base
|
||||
)
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
text: str,
|
||||
language: Optional[str] = None,
|
||||
entities: Optional[List] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Apply the Generic Guardrail API to the given text.
|
||||
|
||||
This is the main method that gets called by the framework.
|
||||
|
||||
Args:
|
||||
text: The text to check
|
||||
language: Optional language parameter (not used by Generic API)
|
||||
entities: Optional entities parameter (not used by Generic API)
|
||||
request_data: Optional request data dictionary for logging metadata
|
||||
|
||||
Returns:
|
||||
The processed text (original or modified)
|
||||
|
||||
Raises:
|
||||
Exception: If the guardrail blocks the request
|
||||
"""
|
||||
verbose_proxy_logger.debug("Generic Guardrail API: Applying guardrail to text")
|
||||
|
||||
# Use provided request_data or create an empty dict
|
||||
if request_data is None:
|
||||
request_data = {}
|
||||
|
||||
request_body = request_data.get("body") or {}
|
||||
|
||||
# Merge additional provider specific params from config and dynamic params
|
||||
additional_params = {**self.additional_provider_specific_params}
|
||||
|
||||
# Get dynamic params from request if available
|
||||
dynamic_params = self.get_guardrail_dynamic_request_body_params(request_body)
|
||||
if dynamic_params:
|
||||
additional_params.update(dynamic_params)
|
||||
|
||||
# Create request payload
|
||||
guardrail_request = GenericGuardrailAPIRequest(
|
||||
text=text,
|
||||
request_body=request_body,
|
||||
additional_provider_specific_params=additional_params,
|
||||
)
|
||||
|
||||
# Prepare headers
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self.headers:
|
||||
headers.update(self.headers)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Generic Guardrail API request to %s: %s",
|
||||
self.api_base,
|
||||
{"text_length": len(text), "has_request_body": bool(request_data)},
|
||||
)
|
||||
|
||||
try:
|
||||
# Make the API request
|
||||
response = await self.async_handler.post(
|
||||
url=self.api_base,
|
||||
json=guardrail_request.to_dict(),
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
response.raise_for_status()
|
||||
response_json = response.json()
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Generic Guardrail API response: %s", response_json
|
||||
)
|
||||
|
||||
guardrail_response = GenericGuardrailAPIResponse.from_dict(response_json)
|
||||
|
||||
# Handle the response
|
||||
if guardrail_response.action == "BLOCKED":
|
||||
# Block the request
|
||||
error_message = (
|
||||
guardrail_response.blocked_reason or "Content violates policy"
|
||||
)
|
||||
verbose_proxy_logger.warning(
|
||||
"Generic Guardrail API blocked request: %s", error_message
|
||||
)
|
||||
raise Exception(f"Content blocked by guardrail: {error_message}")
|
||||
|
||||
elif guardrail_response.action == "GUARDRAIL_INTERVENED":
|
||||
# Content was modified by the guardrail
|
||||
if guardrail_response.text:
|
||||
verbose_proxy_logger.debug("Generic Guardrail API modified text")
|
||||
return guardrail_response.text
|
||||
|
||||
# Action is NONE or no modifications needed
|
||||
return text
|
||||
|
||||
except Exception as e:
|
||||
# Check if it's already an exception we raised
|
||||
if "Content blocked by guardrail" in str(e):
|
||||
raise
|
||||
verbose_proxy_logger.error(
|
||||
"Generic Guardrail API: failed to make request: %s", str(e)
|
||||
)
|
||||
raise Exception(f"Generic Guardrail API failed: {str(e)}")
|
||||
@ -8,6 +8,9 @@ from typing_extensions import Required, TypedDict
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import (
|
||||
EnkryptAIGuardrailConfigs,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
|
||||
GenericGuardrailAPIOptionalParams,
|
||||
)
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import (
|
||||
GraySwanGuardrailConfigModel,
|
||||
)
|
||||
@ -18,7 +21,6 @@ from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
|
||||
ToolPermissionGuardrailConfigModel,
|
||||
)
|
||||
|
||||
|
||||
"""
|
||||
Pydantic object defining how to set guardrails on litellm proxy
|
||||
|
||||
@ -59,6 +61,7 @@ class SupportedGuardrailIntegrations(Enum):
|
||||
IBM_GUARDRAILS = "ibm_guardrails"
|
||||
LITELLM_CONTENT_FILTER = "litellm_content_filter"
|
||||
PROMPT_SECURITY = "prompt_security"
|
||||
GENERIC_GUARDRAIL_API = "generic_guardrail_api"
|
||||
|
||||
|
||||
class Role(Enum):
|
||||
@ -590,6 +593,12 @@ class BaseLitellmParams(BaseModel): # works for new and patch update guardrails
|
||||
description="Whether to fail the request if Model Armor encounters an error",
|
||||
)
|
||||
|
||||
# Generic Guardrail API params
|
||||
additional_provider_specific_params: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="Additional provider-specific parameters for generic guardrail APIs",
|
||||
)
|
||||
|
||||
model_config = ConfigDict(extra="allow", protected_namespaces=())
|
||||
|
||||
|
||||
|
||||
@ -0,0 +1,29 @@
|
||||
from typing import Any, Dict, Literal, Optional
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
|
||||
|
||||
class GenericGuardrailAPIOptionalParams(BaseModel):
|
||||
"""Optional parameters for the Generic Guardrail API"""
|
||||
|
||||
additional_provider_specific_params: Optional[Dict[str, Any]] = Field(
|
||||
default=None,
|
||||
description="Additional provider-specific parameters to send with the guardrail request",
|
||||
)
|
||||
|
||||
|
||||
class GenericGuardrailAPIConfigModel(
|
||||
GuardrailConfigModel[GenericGuardrailAPIOptionalParams],
|
||||
):
|
||||
"""Configuration parameters for the Generic Guardrail API guardrail"""
|
||||
|
||||
optional_params: Optional[GenericGuardrailAPIOptionalParams] = Field(
|
||||
default_factory=GenericGuardrailAPIOptionalParams,
|
||||
description="Optional parameters for the Generic Guardrail API guardrail",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def ui_friendly_name() -> str:
|
||||
return "Generic Guardrail API"
|
||||
Loading…
Reference in New Issue
Block a user