mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-13 16:23:06 +00:00
[Feat] Add LiteLLM Gateway built in guardrail (#16338)
* add ContentFilterAction * store pre-built regex patterns * add v0 of content filter guard * add _filter_messages * test content filter guard * init ContentFilterGuardrail * fix ContentFilterGuardrail enums * rename folder * fix litellm_content_filter * refactor content filter guard * test content filter * add streaming for ContentFilterGuardrail * test_streaming_hook_mask * add litellm_content_filter * docs show litellm content filter * docs litellm content filter * fix lnting * Potential fix for code scanning alert no. 3675: Clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
This commit is contained in:
Ishaan Jaff
co-authored by
Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>
parent
1fec48499f
commit
18a5c4f75a
@@ -0,0 +1,419 @@
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# LiteLLM Content Filter
|
||||
|
||||
**Built-in guardrail** for detecting and filtering sensitive information using regex patterns and keyword matching. No external dependencies required.
|
||||
|
||||
## Overview
|
||||
|
||||
| Property | Details |
|
||||
|----------|---------|
|
||||
| Description | On-device guardrail for detecting and filtering sensitive information using regex patterns and keyword matching. Built into LiteLLM with no external dependencies. |
|
||||
| Guardrail Name | `litellm_content_filter` |
|
||||
| Detection Methods | Prebuilt regex patterns, custom regex, keyword matching |
|
||||
| Actions | `BLOCK` (reject request), `MASK` (redact content) |
|
||||
| Supported Modes | `pre_call`, `post_call`, `during_call` (streaming) |
|
||||
| Performance | Fast - runs locally, no external API calls |
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Define Guardrails in config.yaml
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
model_list:
|
||||
- model_name: gpt-3.5-turbo
|
||||
litellm_params:
|
||||
model: openai/gpt-3.5-turbo
|
||||
api_key: os.environ/OPENAI_API_KEY
|
||||
|
||||
guardrails:
|
||||
- guardrail_name: "content-filter-pre"
|
||||
litellm_params:
|
||||
guardrail: litellm_content_filter
|
||||
mode: "pre_call"
|
||||
|
||||
# Prebuilt patterns for common PII
|
||||
patterns:
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "us_ssn"
|
||||
action: "BLOCK"
|
||||
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "email"
|
||||
action: "MASK"
|
||||
|
||||
# Custom blocked keywords
|
||||
blocked_words:
|
||||
- keyword: "confidential"
|
||||
action: "BLOCK"
|
||||
description: "Sensitive internal information"
|
||||
```
|
||||
|
||||
### 2. Start LiteLLM Gateway
|
||||
|
||||
```shell
|
||||
litellm --config config.yaml
|
||||
```
|
||||
|
||||
### 3. Test Request
|
||||
|
||||
<Tabs>
|
||||
<TabItem label="SSN Blocked" value="ssn-blocked">
|
||||
|
||||
```shell
|
||||
curl -i http://localhost:4000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [
|
||||
{"role": "user", "content": "My SSN is 123-45-6789"}
|
||||
],
|
||||
"guardrails": ["content-filter-pre"]
|
||||
}'
|
||||
```
|
||||
|
||||
**Response: HTTP 400 Error**
|
||||
```json
|
||||
{
|
||||
"error": {
|
||||
"message": {
|
||||
"error": "Content blocked: us_ssn pattern detected",
|
||||
"pattern": "us_ssn"
|
||||
},
|
||||
"code": "400"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem label="Email Masked" value="email-masked">
|
||||
|
||||
```shell
|
||||
curl -i http://localhost:4000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "gpt-3.5-turbo",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Contact me at john@example.com"}
|
||||
],
|
||||
"guardrails": ["content-filter-pre"]
|
||||
}'
|
||||
```
|
||||
|
||||
The request is sent to the LLM with the email masked:
|
||||
```
|
||||
Contact me at [EMAIL_REDACTED]
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Configuration
|
||||
|
||||
### Supported Modes
|
||||
|
||||
- **`pre_call`** - Run before LLM call, filters input messages
|
||||
- **`post_call`** - Run after LLM call, filters output responses
|
||||
- **`during_call`** - Run during streaming, filters each chunk in real-time
|
||||
|
||||
### Actions
|
||||
|
||||
- **`BLOCK`** - Reject the request with HTTP 400 error
|
||||
- **`MASK`** - Replace sensitive content with redaction tags (e.g., `[EMAIL_REDACTED]`)
|
||||
|
||||
## Prebuilt Patterns
|
||||
|
||||
### Available Patterns
|
||||
|
||||
| Pattern Name | Description | Example |
|
||||
|-------------|-------------|---------|
|
||||
| `us_ssn` | US Social Security Numbers | `123-45-6789` |
|
||||
| `email` | Email addresses | `user@example.com` |
|
||||
| `phone` | Phone numbers | `+1-555-123-4567` |
|
||||
| `visa` | Visa credit cards | `4532-1234-5678-9010` |
|
||||
| `mastercard` | Mastercard credit cards | `5425-2334-3010-9903` |
|
||||
| `amex` | American Express cards | `3782-822463-10005` |
|
||||
| `aws_access_key` | AWS access keys | `AKIAIOSFODNN7EXAMPLE` |
|
||||
| `aws_secret_key` | AWS secret keys | `wJalrXUtnFEMI/K7MDENG/bPxRfi...` |
|
||||
| `github_token` | GitHub tokens | `ghp_16C7e42F292c6912E7710c838347Ae178B4a` |
|
||||
|
||||
### Using Prebuilt Patterns
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "pii-filter"
|
||||
litellm_params:
|
||||
guardrail: litellm_content_filter
|
||||
mode: "pre_call"
|
||||
patterns:
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "us_ssn"
|
||||
action: "BLOCK"
|
||||
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "email"
|
||||
action: "MASK"
|
||||
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "aws_access_key"
|
||||
action: "BLOCK"
|
||||
```
|
||||
|
||||
## Custom Regex Patterns
|
||||
|
||||
Define your own regex patterns for domain-specific sensitive data:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "custom-patterns"
|
||||
litellm_params:
|
||||
guardrail: litellm_content_filter
|
||||
mode: "pre_call"
|
||||
patterns:
|
||||
# Custom employee ID format
|
||||
- pattern_type: "regex"
|
||||
pattern: '\b[A-Z]{3}-\d{4}\b'
|
||||
name: "employee_id"
|
||||
action: "MASK"
|
||||
|
||||
# Custom project code format
|
||||
- pattern_type: "regex"
|
||||
pattern: 'PROJECT-\d{6}'
|
||||
name: "project_code"
|
||||
action: "BLOCK"
|
||||
```
|
||||
|
||||
## Keyword Filtering
|
||||
|
||||
Block or mask specific keywords:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "keyword-filter"
|
||||
litellm_params:
|
||||
guardrail: litellm_content_filter
|
||||
mode: "pre_call"
|
||||
blocked_words:
|
||||
- keyword: "confidential"
|
||||
action: "BLOCK"
|
||||
description: "Internal confidential information"
|
||||
|
||||
- keyword: "proprietary"
|
||||
action: "MASK"
|
||||
description: "Proprietary company data"
|
||||
|
||||
- keyword: "secret_project"
|
||||
action: "BLOCK"
|
||||
```
|
||||
|
||||
### Loading Keywords from File
|
||||
|
||||
For large keyword lists, use a YAML file:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "keyword-file-filter"
|
||||
litellm_params:
|
||||
guardrail: litellm_content_filter
|
||||
mode: "pre_call"
|
||||
blocked_words_file: "/path/to/sensitive_keywords.yaml"
|
||||
```
|
||||
|
||||
```yaml showLineNumbers title="sensitive_keywords.yaml"
|
||||
blocked_words:
|
||||
- keyword: "project_apollo"
|
||||
action: "BLOCK"
|
||||
description: "Confidential project codename"
|
||||
|
||||
- keyword: "internal_api"
|
||||
action: "MASK"
|
||||
description: "Internal API references"
|
||||
|
||||
- keyword: "customer_database"
|
||||
action: "BLOCK"
|
||||
description: "Protected database name"
|
||||
```
|
||||
|
||||
## Streaming Support
|
||||
|
||||
Content filter works with streaming responses by checking each chunk:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "streaming-filter"
|
||||
litellm_params:
|
||||
guardrail: litellm_content_filter
|
||||
mode: "during_call" # Check each streaming chunk
|
||||
patterns:
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "email"
|
||||
action: "MASK"
|
||||
```
|
||||
|
||||
```python
|
||||
import openai
|
||||
|
||||
client = openai.OpenAI(
|
||||
api_key="sk-1234",
|
||||
base_url="http://localhost:4000"
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "Tell me about yourself"}],
|
||||
stream=True,
|
||||
extra_body={"guardrails": ["streaming-filter"]}
|
||||
)
|
||||
|
||||
for chunk in response:
|
||||
print(chunk.choices[0].delta.content)
|
||||
# Emails automatically masked in real-time
|
||||
```
|
||||
|
||||
## Customizing Redaction Tags
|
||||
|
||||
When using the `MASK` action, sensitive content is replaced with redaction tags. You can customize how these tags appear.
|
||||
|
||||
### Default Behavior
|
||||
|
||||
**Patterns:** Each pattern type gets its own tag based on the pattern name
|
||||
```
|
||||
Input: "My email is john@example.com and SSN is 123-45-6789"
|
||||
Output: "My email is [EMAIL_REDACTED] and SSN is [US_SSN_REDACTED]"
|
||||
```
|
||||
|
||||
**Keywords:** All keywords use the same generic tag
|
||||
```
|
||||
Input: "This is confidential and proprietary information"
|
||||
Output: "This is [KEYWORD_REDACTED] and [KEYWORD_REDACTED] information"
|
||||
```
|
||||
|
||||
### Customizing Tags
|
||||
|
||||
Use `pattern_redaction_format` and `keyword_redaction_tag` to change the redaction format:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "custom-redaction"
|
||||
litellm_params:
|
||||
guardrail: litellm_content_filter
|
||||
mode: "pre_call"
|
||||
pattern_redaction_format: "***{pattern_name}***" # Use {pattern_name} placeholder
|
||||
keyword_redaction_tag: "***REDACTED***"
|
||||
patterns:
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "email"
|
||||
action: "MASK"
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "us_ssn"
|
||||
action: "MASK"
|
||||
blocked_words:
|
||||
- keyword: "confidential"
|
||||
action: "MASK"
|
||||
```
|
||||
|
||||
**Output:**
|
||||
```
|
||||
Input: "Email john@example.com, SSN 123-45-6789, confidential data"
|
||||
Output: "Email ***EMAIL***, SSN ***US_SSN***, ***REDACTED*** data"
|
||||
```
|
||||
|
||||
**Key Points:**
|
||||
- `pattern_redaction_format` must include `{pattern_name}` placeholder
|
||||
- Pattern names are automatically uppercased (e.g., `email` → `EMAIL`)
|
||||
- `keyword_redaction_tag` is a fixed string (no placeholders)
|
||||
|
||||
## Use Cases
|
||||
|
||||
### 1. PII Protection
|
||||
Block or mask personally identifiable information before sending to LLMs:
|
||||
|
||||
```yaml
|
||||
patterns:
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "us_ssn"
|
||||
action: "BLOCK"
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "email"
|
||||
action: "MASK"
|
||||
```
|
||||
|
||||
### 2. Credential Detection
|
||||
Prevent API keys and secrets from being exposed:
|
||||
|
||||
```yaml
|
||||
patterns:
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "aws_access_key"
|
||||
action: "BLOCK"
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "github_token"
|
||||
action: "BLOCK"
|
||||
```
|
||||
|
||||
### 3. Sensitive Internal Data Protection
|
||||
Block or mask references to confidential internal projects, codenames, or proprietary information:
|
||||
|
||||
```yaml
|
||||
blocked_words:
|
||||
- keyword: "project_titan"
|
||||
action: "BLOCK"
|
||||
description: "Confidential project codename"
|
||||
- keyword: "internal_api"
|
||||
action: "MASK"
|
||||
description: "Internal system references"
|
||||
```
|
||||
|
||||
For large lists of sensitive terms, use a file:
|
||||
```yaml
|
||||
blocked_words_file: "/path/to/sensitive_terms.yaml"
|
||||
```
|
||||
|
||||
### 4. Compliance
|
||||
Ensure regulatory compliance by filtering sensitive data types:
|
||||
|
||||
```yaml
|
||||
patterns:
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "visa"
|
||||
action: "BLOCK"
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "us_ssn"
|
||||
action: "BLOCK"
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Pattern Not Matching
|
||||
|
||||
**Issue:** Regex pattern isn't detecting expected content
|
||||
|
||||
**Solution:** Test your regex pattern:
|
||||
```python
|
||||
import re
|
||||
pattern = r'\b[A-Z]{3}-\d{4}\b'
|
||||
test_text = "Employee ID: ABC-1234"
|
||||
print(re.search(pattern, test_text)) # Should match
|
||||
```
|
||||
|
||||
### Multiple Pattern Matches
|
||||
|
||||
**Issue:** Text contains multiple sensitive patterns
|
||||
|
||||
**Solution:** First matching pattern/keyword is processed. Order patterns by priority:
|
||||
```yaml
|
||||
patterns:
|
||||
# Most critical first
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "us_ssn"
|
||||
action: "BLOCK"
|
||||
# Less critical
|
||||
- pattern_type: "prebuilt"
|
||||
pattern_name: "email"
|
||||
action: "MASK"
|
||||
```
|
||||
|
||||
@@ -41,6 +41,7 @@ const sidebars = {
|
||||
"proxy/guardrails/ibm_guardrails",
|
||||
"proxy/guardrails/grayswan",
|
||||
"proxy/guardrails/lasso_security",
|
||||
"proxy/guardrails/litellm_content_filter",
|
||||
"proxy/guardrails/guardrails_ai",
|
||||
"proxy/guardrails/lakera_ai",
|
||||
"proxy/guardrails/model_armor",
|
||||
|
||||
@@ -4369,7 +4369,7 @@ class StandardLoggingPayloadSetup:
|
||||
|
||||
s3_object_key = get_s3_object_key(
|
||||
s3_path=s3_path, # Use actual s3_path from logger configuration
|
||||
team_alias_prefix="", # Don't split by team alias for cold storage
|
||||
prefix="", # Don't split by team alias for cold storage
|
||||
start_time=start_time,
|
||||
s3_file_name=s3_file_name,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
|
||||
ContentFilterGuardrail,
|
||||
)
|
||||
from litellm.types.guardrails import SupportedGuardrailIntegrations
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.types.guardrails import Guardrail, LitellmParams
|
||||
|
||||
|
||||
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
|
||||
"""
|
||||
Initialize the Content Filter Guardrail.
|
||||
|
||||
Args:
|
||||
litellm_params: Guardrail configuration parameters
|
||||
guardrail: Guardrail metadata
|
||||
|
||||
Returns:
|
||||
Initialized ContentFilterGuardrail instance
|
||||
"""
|
||||
guardrail_name = guardrail.get("guardrail_name")
|
||||
if not guardrail_name:
|
||||
raise ValueError("Content Filter: guardrail_name is required")
|
||||
|
||||
content_filter_guardrail = ContentFilterGuardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
patterns=litellm_params.patterns,
|
||||
blocked_words=litellm_params.blocked_words,
|
||||
blocked_words_file=litellm_params.blocked_words_file,
|
||||
event_hook=litellm_params.mode, # type: ignore
|
||||
default_on=litellm_params.default_on or False,
|
||||
)
|
||||
|
||||
litellm.logging_callback_manager.add_litellm_callback(
|
||||
content_filter_guardrail
|
||||
)
|
||||
|
||||
return content_filter_guardrail
|
||||
|
||||
|
||||
guardrail_initializer_registry = {
|
||||
SupportedGuardrailIntegrations.LITELLM_CONTENT_FILTER.value: initialize_guardrail,
|
||||
}
|
||||
|
||||
|
||||
guardrail_class_registry = {
|
||||
SupportedGuardrailIntegrations.LITELLM_CONTENT_FILTER.value: ContentFilterGuardrail,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
"""
|
||||
Content Filter Guardrail for LiteLLM.
|
||||
|
||||
This guardrail provides regex pattern matching and keyword filtering
|
||||
to detect and block/mask sensitive content.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Any, AsyncGenerator, Dict, List, Optional, Pattern, Tuple, Union
|
||||
|
||||
import yaml
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.types.guardrails import (
|
||||
BlockedWord,
|
||||
ContentFilterAction,
|
||||
ContentFilterPattern,
|
||||
GuardrailEventHooks,
|
||||
Mode,
|
||||
PiiEntityType,
|
||||
)
|
||||
from litellm.types.utils import ModelResponseStream
|
||||
|
||||
from .patterns import get_compiled_pattern
|
||||
|
||||
|
||||
class ContentFilterGuardrail(CustomGuardrail):
|
||||
"""
|
||||
Content filter guardrail that detects sensitive information using:
|
||||
- Prebuilt regex patterns (SSN, credit cards, API keys, etc.)
|
||||
- Custom user-defined regex patterns
|
||||
- Dictionary-based keyword matching
|
||||
|
||||
Actions:
|
||||
- BLOCK: Reject the request with an error
|
||||
- MASK: Replace the sensitive content with a redacted placeholder
|
||||
"""
|
||||
|
||||
# Redaction format constants
|
||||
PATTERN_REDACTION_FORMAT = "[{pattern_name}_REDACTED]"
|
||||
KEYWORD_REDACTION_STR = "[KEYWORD_REDACTED]"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
guardrail_name: Optional[str] = None,
|
||||
patterns: Optional[List[ContentFilterPattern]] = None,
|
||||
blocked_words: Optional[List[BlockedWord]] = None,
|
||||
blocked_words_file: Optional[str] = None,
|
||||
event_hook: Optional[Union[GuardrailEventHooks, List[GuardrailEventHooks], Mode]] = None,
|
||||
default_on: bool = False,
|
||||
pattern_redaction_format: Optional[str] = None,
|
||||
keyword_redaction_tag: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Initialize the Content Filter Guardrail.
|
||||
|
||||
Args:
|
||||
guardrail_name: Name of this guardrail instance
|
||||
patterns: List of ContentFilterPattern objects to detect
|
||||
blocked_words: List of BlockedWord objects with keywords and actions
|
||||
blocked_words_file: Path to YAML file containing blocked_words list
|
||||
event_hook: When to run this guardrail (pre_call, post_call, etc.)
|
||||
default_on: If True, runs on all requests by default
|
||||
pattern_redaction_format: Format string for pattern redaction (use {pattern_name} placeholder)
|
||||
keyword_redaction_tag: Tag to use for keyword redaction
|
||||
"""
|
||||
super().__init__(
|
||||
guardrail_name=guardrail_name,
|
||||
supported_event_hooks=[
|
||||
GuardrailEventHooks.pre_call,
|
||||
GuardrailEventHooks.post_call,
|
||||
GuardrailEventHooks.during_call,
|
||||
],
|
||||
event_hook=event_hook or GuardrailEventHooks.pre_call,
|
||||
default_on=default_on,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.guardrail_provider = "litellm_content_filter"
|
||||
self.pattern_redaction_format = pattern_redaction_format or self.PATTERN_REDACTION_FORMAT
|
||||
self.keyword_redaction_tag = keyword_redaction_tag or self.KEYWORD_REDACTION_STR
|
||||
|
||||
# Compile regex patterns
|
||||
self.compiled_patterns: List[Tuple[Pattern, str, ContentFilterAction]] = []
|
||||
if patterns:
|
||||
for pattern_config in patterns:
|
||||
self._add_pattern(pattern_config)
|
||||
|
||||
# Load blocked words
|
||||
self.blocked_words: Dict[str, Tuple[ContentFilterAction, Optional[str]]] = {}
|
||||
if blocked_words:
|
||||
for word in blocked_words:
|
||||
self.blocked_words[word.keyword.lower()] = (
|
||||
word.action,
|
||||
word.description,
|
||||
)
|
||||
|
||||
# Load blocked words from file if provided
|
||||
if blocked_words_file:
|
||||
self._load_blocked_words_file(blocked_words_file)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"ContentFilterGuardrail initialized with {len(self.compiled_patterns)} patterns "
|
||||
f"and {len(self.blocked_words)} blocked words"
|
||||
)
|
||||
|
||||
def _add_pattern(self, pattern_config: ContentFilterPattern) -> None:
|
||||
"""
|
||||
Add a pattern to the compiled patterns list.
|
||||
|
||||
Args:
|
||||
pattern_config: ContentFilterPattern configuration
|
||||
"""
|
||||
try:
|
||||
if pattern_config.pattern_type == "prebuilt":
|
||||
if not pattern_config.pattern_name:
|
||||
raise ValueError("pattern_name is required for prebuilt patterns")
|
||||
compiled = get_compiled_pattern(pattern_config.pattern_name)
|
||||
pattern_name = pattern_config.pattern_name
|
||||
elif pattern_config.pattern_type == "regex":
|
||||
if not pattern_config.pattern:
|
||||
raise ValueError("pattern is required for regex patterns")
|
||||
compiled = re.compile(pattern_config.pattern, re.IGNORECASE)
|
||||
pattern_name = pattern_config.name or "custom_regex"
|
||||
else:
|
||||
raise ValueError(f"Unknown pattern_type: {pattern_config.pattern_type}")
|
||||
|
||||
self.compiled_patterns.append((compiled, pattern_name, pattern_config.action))
|
||||
verbose_proxy_logger.debug(f"Added pattern: {pattern_name} with action {pattern_config.action}")
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error adding pattern {pattern_config}: {e}")
|
||||
raise
|
||||
|
||||
def _load_blocked_words_file(self, file_path: str) -> None:
|
||||
"""
|
||||
Load blocked words from a YAML file.
|
||||
|
||||
Args:
|
||||
file_path: Path to YAML file containing blocked_words list
|
||||
|
||||
Expected format:
|
||||
```yaml
|
||||
blocked_words:
|
||||
- keyword: "sensitive_term"
|
||||
action: "BLOCK"
|
||||
description: "Optional description"
|
||||
```
|
||||
"""
|
||||
try:
|
||||
with open(file_path, "r") as f:
|
||||
data = yaml.safe_load(f)
|
||||
|
||||
if not isinstance(data, dict) or "blocked_words" not in data:
|
||||
raise ValueError(
|
||||
"Invalid format: file must contain 'blocked_words' key with list of words"
|
||||
)
|
||||
|
||||
for word_data in data["blocked_words"]:
|
||||
if not isinstance(word_data, dict) or "keyword" not in word_data or "action" not in word_data:
|
||||
verbose_proxy_logger.warning(f"Skipping invalid word entry: {word_data}")
|
||||
continue
|
||||
|
||||
keyword = word_data["keyword"].lower()
|
||||
action = ContentFilterAction(word_data["action"])
|
||||
description = word_data.get("description")
|
||||
|
||||
self.blocked_words[keyword] = (action, description)
|
||||
|
||||
verbose_proxy_logger.info(
|
||||
f"Loaded {len(data['blocked_words'])} blocked words from {file_path}"
|
||||
)
|
||||
except FileNotFoundError:
|
||||
raise FileNotFoundError(f"Blocked words file not found: {file_path}")
|
||||
except Exception as e:
|
||||
raise Exception(f"Error loading blocked words file {file_path}: {str(e)}")
|
||||
|
||||
def _check_patterns(self, text: str) -> Optional[Tuple[str, str, ContentFilterAction]]:
|
||||
"""
|
||||
Check text against all compiled regex patterns.
|
||||
|
||||
Args:
|
||||
text: Text to check
|
||||
|
||||
Returns:
|
||||
Tuple of (matched_text, pattern_name, action) if match found, None otherwise
|
||||
"""
|
||||
for compiled_pattern, pattern_name, action in self.compiled_patterns:
|
||||
match = compiled_pattern.search(text)
|
||||
if match:
|
||||
matched_text = match.group(0)
|
||||
verbose_proxy_logger.debug(
|
||||
f"Pattern '{pattern_name}' matched."
|
||||
)
|
||||
return (matched_text, pattern_name, action)
|
||||
return None
|
||||
|
||||
def _check_blocked_words(self, text: str) -> Optional[Tuple[str, ContentFilterAction, Optional[str]]]:
|
||||
"""
|
||||
Check text for blocked keywords.
|
||||
|
||||
Args:
|
||||
text: Text to check
|
||||
|
||||
Returns:
|
||||
Tuple of (keyword, action, description) if match found, None otherwise
|
||||
"""
|
||||
text_lower = text.lower()
|
||||
for keyword, (action, description) in self.blocked_words.items():
|
||||
if keyword in text_lower:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Blocked word '{keyword}' found with action {action}"
|
||||
)
|
||||
return (keyword, action, description)
|
||||
return None
|
||||
|
||||
def _mask_content(self, text: str, pattern_name: str) -> str:
|
||||
"""
|
||||
Mask sensitive content in text.
|
||||
|
||||
Args:
|
||||
text: Text containing sensitive content
|
||||
pattern_name: Name of the pattern that matched
|
||||
|
||||
Returns:
|
||||
Text with sensitive content masked
|
||||
"""
|
||||
redaction_tag = self.pattern_redaction_format.format(
|
||||
pattern_name=pattern_name.upper()
|
||||
)
|
||||
return redaction_tag
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
text: str,
|
||||
language: Optional[str] = None,
|
||||
entities: Optional[List[PiiEntityType]] = None,
|
||||
request_data: Optional[dict] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Apply content filtering guardrail to the given text.
|
||||
|
||||
This method checks for sensitive patterns and blocked keywords,
|
||||
either blocking the request or masking the sensitive content.
|
||||
|
||||
Args:
|
||||
text: The text to apply the guardrail to
|
||||
language: Optional language parameter (not used)
|
||||
entities: Optional entities parameter (not used)
|
||||
request_data: Optional request data dictionary for logging metadata
|
||||
|
||||
Returns:
|
||||
Text with sensitive content masked (if action is MASK)
|
||||
|
||||
Raises:
|
||||
HTTPException: If sensitive content is detected and action is BLOCK
|
||||
"""
|
||||
verbose_proxy_logger.debug("ContentFilterGuardrail: Applying guardrail to text")
|
||||
|
||||
# Check regex patterns
|
||||
pattern_match = self._check_patterns(text)
|
||||
if pattern_match:
|
||||
matched_text, pattern_name, action = pattern_match
|
||||
|
||||
if action == ContentFilterAction.BLOCK:
|
||||
error_msg = f"Content blocked: {pattern_name} pattern detected"
|
||||
verbose_proxy_logger.warning(error_msg)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": error_msg, "pattern": pattern_name},
|
||||
)
|
||||
elif action == ContentFilterAction.MASK:
|
||||
# Replace the matched text with redaction tag
|
||||
redaction_tag = self._mask_content(matched_text, pattern_name)
|
||||
text = text.replace(matched_text, redaction_tag)
|
||||
verbose_proxy_logger.info(f"Masked {pattern_name} in content")
|
||||
|
||||
# Check blocked words
|
||||
word_match = self._check_blocked_words(text)
|
||||
if word_match:
|
||||
keyword, action, description = word_match
|
||||
|
||||
if action == ContentFilterAction.BLOCK:
|
||||
error_msg = f"Content blocked: keyword '{keyword}' detected"
|
||||
if description:
|
||||
error_msg += f" ({description})"
|
||||
verbose_proxy_logger.warning(error_msg)
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": error_msg,
|
||||
"keyword": keyword,
|
||||
"description": description,
|
||||
},
|
||||
)
|
||||
elif action == ContentFilterAction.MASK:
|
||||
# Replace keyword with redaction tag (case-insensitive)
|
||||
text = re.sub(
|
||||
re.escape(keyword),
|
||||
self.keyword_redaction_tag,
|
||||
text,
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
verbose_proxy_logger.info(f"Masked keyword '{keyword}' in content")
|
||||
|
||||
verbose_proxy_logger.debug("ContentFilterGuardrail: Guardrail applied successfully")
|
||||
return text
|
||||
|
||||
async def async_post_call_streaming_iterator_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: Any,
|
||||
request_data: dict,
|
||||
) -> AsyncGenerator[ModelResponseStream, None]:
|
||||
"""
|
||||
Streaming hook to check each chunk as it's yielded.
|
||||
|
||||
This implementation checks each chunk individually and yields it immediately,
|
||||
allowing for low-latency streaming with content filtering.
|
||||
|
||||
Args:
|
||||
user_api_key_dict: User API key authentication
|
||||
response: Async generator of response chunks
|
||||
request_data: Original request data
|
||||
|
||||
Yields:
|
||||
Checked and potentially masked chunks
|
||||
|
||||
Raises:
|
||||
HTTPException: If chunk content should be blocked
|
||||
"""
|
||||
verbose_proxy_logger.debug(
|
||||
"ContentFilterGuardrail: Running streaming check (per-chunk mode)"
|
||||
)
|
||||
|
||||
# Process each chunk individually
|
||||
async for chunk in response:
|
||||
if isinstance(chunk, ModelResponseStream):
|
||||
for choice in chunk.choices:
|
||||
if hasattr(choice, "delta") and choice.delta.content:
|
||||
if isinstance(choice.delta.content, str):
|
||||
# Check the chunk content using apply_guardrail
|
||||
try:
|
||||
processed_content = await self.apply_guardrail(
|
||||
text=choice.delta.content,
|
||||
request_data=request_data,
|
||||
)
|
||||
if processed_content != choice.delta.content:
|
||||
choice.delta.content = processed_content
|
||||
verbose_proxy_logger.debug(
|
||||
"ContentFilterGuardrail: Modified streaming chunk"
|
||||
)
|
||||
except HTTPException as e:
|
||||
# If content should be blocked, raise immediately
|
||||
verbose_proxy_logger.warning(
|
||||
f"ContentFilterGuardrail: Blocked streaming chunk: {e.detail}"
|
||||
)
|
||||
raise
|
||||
|
||||
yield chunk
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"ContentFilterGuardrail: Streaming check completed"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
"""
|
||||
Prebuilt regex patterns for content filtering.
|
||||
|
||||
This module contains predefined regex patterns for detecting sensitive information
|
||||
like SSNs, credit cards, API keys, etc.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Dict, Pattern
|
||||
|
||||
# US Social Security Number patterns
|
||||
US_SSN_PATTERN = r"\b\d{3}-\d{2}-\d{4}\b" # Format: 123-45-6789
|
||||
US_SSN_NO_DASH_PATTERN = r"\b(?!000|666|9\d{2})\d{3}(?!00)\d{2}(?!0000)\d{4}\b" # Format: 123456789 (with validation)
|
||||
|
||||
# Credit Card patterns
|
||||
VISA_PATTERN = r"\b4\d{3}[\s\-]?\d{4}[\s\-]?\d{4}[\s\-]?\d{4}\b" # Visa starts with 4
|
||||
MASTERCARD_PATTERN = r"\b5[1-5]\d{2}[\s\-]?\d{4}[\s\-]?\d{4}[\s\-]?\d{4}\b" # Mastercard starts with 51-55
|
||||
AMEX_PATTERN = r"\b3[47]\d{2}[\s\-]?\d{6}[\s\-]?\d{5}\b" # Amex starts with 34 or 37
|
||||
DISCOVER_PATTERN = r"\b6(?:011|5\d{2})[\s\-]?\d{4}[\s\-]?\d{4}[\s\-]?\d{4}\b" # Discover starts with 6011 or 65
|
||||
|
||||
# Email pattern
|
||||
EMAIL_PATTERN = r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b"
|
||||
|
||||
# Phone number patterns (US)
|
||||
US_PHONE_PATTERN = r"\b(?:\+?1[\s.-]?)?\(?([0-9]{3})\)?[\s.-]?([0-9]{3})[\s.-]?([0-9]{4})\b"
|
||||
|
||||
# API Key patterns (common formats)
|
||||
AWS_ACCESS_KEY_PATTERN = r"\b(AKIA[0-9A-Z]{16})\b" # AWS Access Key ID
|
||||
AWS_SECRET_KEY_PATTERN = r"\b([A-Za-z0-9/+=]{40})\b" # AWS Secret Access Key (generic 40 char)
|
||||
GITHUB_TOKEN_PATTERN = r"\b(gh[ps]_[a-zA-Z0-9]{36})\b" # GitHub Personal Access Token
|
||||
SLACK_TOKEN_PATTERN = r"\b(xox[pboa]-[0-9]{10,13}-[0-9]{10,13}-[a-zA-Z0-9]{24,32})\b" # Slack tokens
|
||||
GENERIC_API_KEY_PATTERN = r"\b([Aa][Pp][Ii][-_]?[Kk][Ee][Yy][\s:=]+['\"]?[A-Za-z0-9_\-]{20,}['\"]?)\b"
|
||||
|
||||
# IP Address patterns
|
||||
IPV4_PATTERN = r"\b(?:(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\.){3}(?:25[0-5]|2[0-4][0-9]|[01]?[0-9][0-9]?)\b"
|
||||
IPV6_PATTERN = r"\b(?:[0-9a-fA-F]{1,4}:){7}[0-9a-fA-F]{1,4}\b"
|
||||
|
||||
# URL patterns
|
||||
URL_PATTERN = r"\b(?:https?://|www\.)[^\s/$.?#].[^\s]*\b"
|
||||
|
||||
|
||||
PREBUILT_PATTERNS: Dict[str, str] = {
|
||||
# SSN patterns
|
||||
"us_ssn": US_SSN_PATTERN,
|
||||
"us_ssn_no_dash": US_SSN_NO_DASH_PATTERN,
|
||||
|
||||
# Credit card patterns
|
||||
"visa": VISA_PATTERN,
|
||||
"mastercard": MASTERCARD_PATTERN,
|
||||
"amex": AMEX_PATTERN,
|
||||
"discover": DISCOVER_PATTERN,
|
||||
"credit_card": rf"(?:{VISA_PATTERN}|{MASTERCARD_PATTERN}|{AMEX_PATTERN}|{DISCOVER_PATTERN})",
|
||||
|
||||
# Contact information
|
||||
"email": EMAIL_PATTERN,
|
||||
"us_phone": US_PHONE_PATTERN,
|
||||
|
||||
# API keys and secrets
|
||||
"aws_access_key": AWS_ACCESS_KEY_PATTERN,
|
||||
"aws_secret_key": AWS_SECRET_KEY_PATTERN,
|
||||
"github_token": GITHUB_TOKEN_PATTERN,
|
||||
"slack_token": SLACK_TOKEN_PATTERN,
|
||||
"generic_api_key": GENERIC_API_KEY_PATTERN,
|
||||
|
||||
# Network identifiers
|
||||
"ipv4": IPV4_PATTERN,
|
||||
"ipv6": IPV6_PATTERN,
|
||||
"url": URL_PATTERN,
|
||||
}
|
||||
|
||||
|
||||
def get_compiled_pattern(pattern_name: str) -> Pattern:
|
||||
"""
|
||||
Get a compiled regex pattern by name.
|
||||
|
||||
Args:
|
||||
pattern_name: Name of the prebuilt pattern
|
||||
|
||||
Returns:
|
||||
Compiled regex pattern
|
||||
|
||||
Raises:
|
||||
ValueError: If pattern_name is not found in PREBUILT_PATTERNS
|
||||
"""
|
||||
if pattern_name not in PREBUILT_PATTERNS:
|
||||
available_patterns = ", ".join(PREBUILT_PATTERNS.keys())
|
||||
raise ValueError(
|
||||
f"Unknown pattern name: '{pattern_name}'. "
|
||||
f"Available patterns: {available_patterns}"
|
||||
)
|
||||
|
||||
return re.compile(PREBUILT_PATTERNS[pattern_name], re.IGNORECASE)
|
||||
|
||||
|
||||
def get_all_pattern_names():
|
||||
"""
|
||||
Get a list of all available prebuilt pattern names.
|
||||
|
||||
Returns:
|
||||
List of pattern names
|
||||
"""
|
||||
return list(PREBUILT_PATTERNS.keys())
|
||||
|
||||
@@ -52,6 +52,7 @@ class SupportedGuardrailIntegrations(Enum):
|
||||
JAVELIN = "javelin"
|
||||
ENKRYPTAI = "enkryptai"
|
||||
IBM_GUARDRAILS = "ibm_guardrails"
|
||||
LITELLM_CONTENT_FILTER = "litellm_content_filter"
|
||||
|
||||
|
||||
class Role(Enum):
|
||||
@@ -443,6 +444,65 @@ class JavelinGuardrailConfigModel(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class ContentFilterAction(str, Enum):
|
||||
"""Action to take when content filter detects a match"""
|
||||
|
||||
BLOCK = "BLOCK"
|
||||
MASK = "MASK"
|
||||
|
||||
|
||||
class BlockedWord(BaseModel):
|
||||
"""Represents a blocked word with its action and optional description"""
|
||||
|
||||
keyword: str = Field(description="The keyword to block or mask")
|
||||
action: ContentFilterAction = Field(
|
||||
description="Action to take when keyword is detected (BLOCK or MASK)"
|
||||
)
|
||||
description: Optional[str] = Field(
|
||||
default=None, description="Optional description explaining why this keyword is sensitive"
|
||||
)
|
||||
|
||||
|
||||
class ContentFilterPattern(BaseModel):
|
||||
"""Represents a content filter pattern (prebuilt or custom regex)"""
|
||||
|
||||
pattern_type: Literal["prebuilt", "regex"] = Field(
|
||||
description="Type of pattern: 'prebuilt' for predefined patterns or 'regex' for custom"
|
||||
)
|
||||
pattern_name: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Name of prebuilt pattern (e.g., 'us_ssn', 'credit_card'). Required if pattern_type is 'prebuilt'"
|
||||
)
|
||||
pattern: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Custom regex pattern. Required if pattern_type is 'regex'"
|
||||
)
|
||||
name: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Name for this pattern (used in logging and error messages)"
|
||||
)
|
||||
action: ContentFilterAction = Field(
|
||||
description="Action to take when pattern matches (BLOCK or MASK)"
|
||||
)
|
||||
|
||||
|
||||
class ContentFilterConfigModel(BaseModel):
|
||||
"""Configuration parameters for the content filter guardrail"""
|
||||
|
||||
patterns: Optional[List[ContentFilterPattern]] = Field(
|
||||
default=None,
|
||||
description="List of patterns (prebuilt or custom regex) to detect"
|
||||
)
|
||||
blocked_words: Optional[List[BlockedWord]] = Field(
|
||||
default=None,
|
||||
description="List of blocked words with individual actions"
|
||||
)
|
||||
blocked_words_file: Optional[str] = Field(
|
||||
default=None,
|
||||
description="Path to YAML file containing blocked_words list"
|
||||
)
|
||||
|
||||
|
||||
class BaseLitellmParams(BaseModel): # works for new and patch update guardrails
|
||||
api_key: Optional[str] = Field(
|
||||
default=None, description="API key for the guardrail service"
|
||||
@@ -534,6 +594,7 @@ class LitellmParams(
|
||||
NomaGuardrailConfigModel,
|
||||
ToolPermissionGuardrailConfigModel,
|
||||
JavelinGuardrailConfigModel,
|
||||
ContentFilterConfigModel,
|
||||
BaseLitellmParams,
|
||||
EnkryptAIGuardrailConfigs,
|
||||
IBMGuardrailsBaseConfigModel,
|
||||
|
||||
+469
@@ -0,0 +1,469 @@
|
||||
"""
|
||||
Tests for the Content Filter Guardrail
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import (
|
||||
ContentFilterGuardrail,
|
||||
)
|
||||
from litellm.types.guardrails import (
|
||||
BlockedWord,
|
||||
ContentFilterAction,
|
||||
ContentFilterPattern,
|
||||
GuardrailEventHooks,
|
||||
)
|
||||
|
||||
|
||||
class TestContentFilterGuardrail:
|
||||
"""Test the ContentFilterGuardrail class"""
|
||||
|
||||
def test_init_with_patterns(self):
|
||||
"""
|
||||
Test initialization with prebuilt patterns
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="us_ssn",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-content-filter",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
assert guardrail.guardrail_name == "test-content-filter"
|
||||
assert len(guardrail.compiled_patterns) == 1
|
||||
|
||||
def test_init_with_blocked_words(self):
|
||||
"""
|
||||
Test initialization with blocked words
|
||||
"""
|
||||
blocked_words = [
|
||||
BlockedWord(
|
||||
keyword="secret_project",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
description="Top secret project"
|
||||
),
|
||||
BlockedWord(
|
||||
keyword="internal_api",
|
||||
action=ContentFilterAction.MASK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-content-filter",
|
||||
blocked_words=blocked_words,
|
||||
)
|
||||
|
||||
assert len(guardrail.blocked_words) == 2
|
||||
assert "secret_project" in guardrail.blocked_words
|
||||
assert guardrail.blocked_words["secret_project"][0] == ContentFilterAction.BLOCK
|
||||
|
||||
def test_check_patterns_ssn(self):
|
||||
"""
|
||||
Test SSN pattern detection
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="us_ssn",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-ssn",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
# Test with SSN
|
||||
result = guardrail._check_patterns("My SSN is 123-45-6789")
|
||||
assert result is not None
|
||||
assert result[1] == "us_ssn"
|
||||
assert result[2] == ContentFilterAction.BLOCK
|
||||
|
||||
# Test without SSN
|
||||
result = guardrail._check_patterns("This is a normal message")
|
||||
assert result is None
|
||||
|
||||
def test_check_patterns_email(self):
|
||||
"""
|
||||
Test email pattern detection
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="email",
|
||||
action=ContentFilterAction.MASK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-email",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
result = guardrail._check_patterns("Contact me at test@example.com")
|
||||
assert result is not None
|
||||
assert result[1] == "email"
|
||||
assert result[2] == ContentFilterAction.MASK
|
||||
|
||||
def test_check_patterns_custom_regex(self):
|
||||
"""
|
||||
Test custom regex pattern detection
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="regex",
|
||||
pattern=r"\b[A-Z]{3}-\d{4}\b",
|
||||
name="custom_id",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-custom",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
result = guardrail._check_patterns("My ID is ABC-1234")
|
||||
assert result is not None
|
||||
assert result[1] == "custom_id"
|
||||
|
||||
def test_check_blocked_words(self):
|
||||
"""
|
||||
Test blocked word detection
|
||||
"""
|
||||
blocked_words = [
|
||||
BlockedWord(
|
||||
keyword="confidential",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-words",
|
||||
blocked_words=blocked_words,
|
||||
)
|
||||
|
||||
# Test with blocked word
|
||||
result = guardrail._check_blocked_words("This is CONFIDENTIAL information")
|
||||
assert result is not None
|
||||
assert result[0] == "confidential"
|
||||
assert result[1] == ContentFilterAction.BLOCK
|
||||
|
||||
# Test without blocked word
|
||||
result = guardrail._check_blocked_words("This is normal information")
|
||||
assert result is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_block(self):
|
||||
"""
|
||||
Test apply_guardrail with BLOCK action
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="us_ssn",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-block",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.apply_guardrail(text="My SSN is 123-45-6789")
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "us_ssn" in str(exc_info.value.detail)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_mask(self):
|
||||
"""
|
||||
Test apply_guardrail with MASK action
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="email",
|
||||
action=ContentFilterAction.MASK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-mask",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
result = await guardrail.apply_guardrail(text="Contact me at test@example.com")
|
||||
|
||||
assert result is not None
|
||||
assert "[EMAIL_REDACTED]" in result
|
||||
assert "test@example.com" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_blocked_word_mask(self):
|
||||
"""
|
||||
Test apply_guardrail with blocked word MASK action
|
||||
"""
|
||||
blocked_words = [
|
||||
BlockedWord(
|
||||
keyword="proprietary",
|
||||
action=ContentFilterAction.MASK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-word-mask",
|
||||
blocked_words=blocked_words,
|
||||
)
|
||||
|
||||
result = await guardrail.apply_guardrail(text="This is PROPRIETARY information")
|
||||
|
||||
assert result is not None
|
||||
assert "[KEYWORD_REDACTED]" in result
|
||||
assert "PROPRIETARY" not in result
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_multiple_patterns(self):
|
||||
"""
|
||||
Test apply_guardrail with multiple patterns in the same text
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="email",
|
||||
action=ContentFilterAction.MASK,
|
||||
),
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="us_ssn",
|
||||
action=ContentFilterAction.MASK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-multiple",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
result = await guardrail.apply_guardrail(
|
||||
text="Contact user@test.com or SSN: 123-45-6789"
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
# At least one pattern should be redacted (first match wins)
|
||||
assert "[EMAIL_REDACTED]" in result or "[US_SSN_REDACTED]" in result
|
||||
|
||||
def test_mask_content(self):
|
||||
"""
|
||||
Test content masking
|
||||
"""
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-mask",
|
||||
)
|
||||
|
||||
masked = guardrail._mask_content("sensitive text", "us_ssn")
|
||||
assert masked == "[US_SSN_REDACTED]"
|
||||
|
||||
def test_load_blocked_words_file(self):
|
||||
"""
|
||||
Test loading blocked words from a YAML file
|
||||
"""
|
||||
import tempfile
|
||||
|
||||
# Create a temporary blocked words file
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
|
||||
f.write("""blocked_words:
|
||||
- keyword: "test_keyword"
|
||||
action: "BLOCK"
|
||||
description: "Test keyword"
|
||||
- keyword: "another_word"
|
||||
action: "MASK"
|
||||
""")
|
||||
temp_file = f.name
|
||||
|
||||
try:
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-file-load",
|
||||
blocked_words_file=temp_file,
|
||||
)
|
||||
|
||||
assert len(guardrail.blocked_words) == 2
|
||||
assert "test_keyword" in guardrail.blocked_words
|
||||
assert guardrail.blocked_words["test_keyword"][0] == ContentFilterAction.BLOCK
|
||||
assert guardrail.blocked_words["test_keyword"][1] == "Test keyword"
|
||||
assert "another_word" in guardrail.blocked_words
|
||||
assert guardrail.blocked_words["another_word"][0] == ContentFilterAction.MASK
|
||||
finally:
|
||||
os.unlink(temp_file)
|
||||
|
||||
def test_credit_card_patterns(self):
|
||||
"""
|
||||
Test credit card pattern detection
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="visa",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-cc",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
# Test Visa card
|
||||
result = guardrail._check_patterns("My card is 4532-1234-5678-9010")
|
||||
assert result is not None
|
||||
assert result[1] == "visa"
|
||||
|
||||
def test_api_key_patterns(self):
|
||||
"""
|
||||
Test API key pattern detection
|
||||
"""
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="aws_access_key",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-api-key",
|
||||
patterns=patterns,
|
||||
)
|
||||
|
||||
# Test AWS Access Key
|
||||
result = guardrail._check_patterns("My key is AKIAIOSFODNN7EXAMPLE")
|
||||
assert result is not None
|
||||
assert result[1] == "aws_access_key"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_mask(self):
|
||||
"""
|
||||
Test streaming hook with MASK action
|
||||
"""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="email",
|
||||
action=ContentFilterAction.MASK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-streaming-mask",
|
||||
patterns=patterns,
|
||||
event_hook=GuardrailEventHooks.during_call,
|
||||
)
|
||||
|
||||
# Create mock streaming chunks
|
||||
async def mock_stream():
|
||||
# Chunk 1: contains email
|
||||
chunk1 = ModelResponseStream(
|
||||
id="chunk1",
|
||||
choices=[StreamingChoices(delta=Delta(content="Contact me at test@example.com"), index=0)],
|
||||
model="gpt-4",
|
||||
)
|
||||
yield chunk1
|
||||
|
||||
# Chunk 2: normal content
|
||||
chunk2 = ModelResponseStream(
|
||||
id="chunk2",
|
||||
choices=[StreamingChoices(delta=Delta(content=" for more info"), index=0)],
|
||||
model="gpt-4",
|
||||
)
|
||||
yield chunk2
|
||||
|
||||
user_api_key_dict = MagicMock()
|
||||
request_data = {}
|
||||
|
||||
# Process streaming response
|
||||
result_chunks = []
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
result_chunks.append(chunk)
|
||||
|
||||
assert len(result_chunks) == 2
|
||||
# First chunk should have email masked
|
||||
assert "[EMAIL_REDACTED]" in result_chunks[0].choices[0].delta.content
|
||||
assert "test@example.com" not in result_chunks[0].choices[0].delta.content
|
||||
# Second chunk should be unchanged
|
||||
assert result_chunks[1].choices[0].delta.content == " for more info"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_hook_block(self):
|
||||
"""
|
||||
Test streaming hook with BLOCK action
|
||||
"""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
|
||||
|
||||
patterns = [
|
||||
ContentFilterPattern(
|
||||
pattern_type="prebuilt",
|
||||
pattern_name="us_ssn",
|
||||
action=ContentFilterAction.BLOCK,
|
||||
),
|
||||
]
|
||||
|
||||
guardrail = ContentFilterGuardrail(
|
||||
guardrail_name="test-streaming-block",
|
||||
patterns=patterns,
|
||||
event_hook=GuardrailEventHooks.during_call,
|
||||
)
|
||||
|
||||
# Create mock streaming chunks with SSN
|
||||
async def mock_stream():
|
||||
chunk = ModelResponseStream(
|
||||
id="chunk1",
|
||||
choices=[StreamingChoices(delta=Delta(content="SSN: 123-45-6789"), index=0)],
|
||||
model="gpt-4",
|
||||
)
|
||||
yield chunk
|
||||
|
||||
user_api_key_dict = MagicMock()
|
||||
request_data = {}
|
||||
|
||||
# Should raise HTTPException when SSN is detected
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=mock_stream(),
|
||||
request_data=request_data,
|
||||
):
|
||||
pass
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert "us_ssn" in str(exc_info.value.detail)
|
||||
|
||||
Reference in New Issue
Block a user