[Feat] Guardrails Load Balancing - Allow Platform admins to load balance between guardrails (#18181)
* add _aguardrail_helper for LB * add _aguardrail_helper on router.py * test_proxy_logging_pre_call_hook_load_balancing * add _execute_guardrail_with_load_balancing * add LB TEsting * docs guard lb * fix linting * fix lint
This commit is contained in:
Ishaan Jaff
GitHub
parent
fcd524ca0f
commit
5ea0854eda
@ -0,0 +1,351 @@
|
||||
import Tabs from '@theme/Tabs';
|
||||
import TabItem from '@theme/TabItem';
|
||||
|
||||
# Guardrail Load Balancing
|
||||
|
||||
Load balance guardrail requests across multiple guardrail deployments. This is useful when you have rate limits on guardrail providers (e.g., AWS Bedrock Guardrails) and want to distribute requests across multiple accounts or regions.
|
||||
|
||||
## How It Works
|
||||
|
||||
```mermaid
|
||||
flowchart LR
|
||||
subgraph LiteLLM Gateway
|
||||
Router[Router]
|
||||
G1[Guardrail Instance A]
|
||||
G2[Guardrail Instance B]
|
||||
G3[Guardrail Instance N]
|
||||
end
|
||||
|
||||
Client[Client Request] --> Router
|
||||
Router -->|Round Robin / Weighted| G1
|
||||
Router -->|Round Robin / Weighted| G2
|
||||
Router -->|Round Robin / Weighted| G3
|
||||
|
||||
G1 --> AWS1[AWS Account 1]
|
||||
G2 --> AWS2[AWS Account 2]
|
||||
G3 --> AWSN[AWS Account N]
|
||||
```
|
||||
|
||||
When you define multiple guardrails with the **same `guardrail_name`**, LiteLLM automatically load balances requests across them using the router's load balancing strategy.
|
||||
|
||||
## Why Use Guardrail Load Balancing?
|
||||
|
||||
| Use Case | Benefit |
|
||||
|----------|---------|
|
||||
| **AWS Bedrock Rate Limits** | Bedrock Guardrails have per-account rate limits. Distribute across multiple AWS accounts to increase throughput |
|
||||
| **Multi-Region Redundancy** | Deploy guardrails across regions for failover and lower latency |
|
||||
| **Cost Optimization** | Spread usage across accounts with different pricing tiers or credits |
|
||||
| **A/B Testing** | Test different guardrail configurations with weighted distribution |
|
||||
|
||||
## Quick Start
|
||||
|
||||
### 1. Define Multiple Guardrails with Same Name
|
||||
|
||||
Define multiple guardrail entries with the **same `guardrail_name`** but different configurations:
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="bedrock" label="Bedrock Guardrails">
|
||||
|
||||
```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:
|
||||
# First Bedrock guardrail - AWS Account 1
|
||||
- guardrail_name: "content-filter"
|
||||
litellm_params:
|
||||
guardrail: bedrock/guardrail
|
||||
mode: "pre_call"
|
||||
guardrailIdentifier: "abc123"
|
||||
guardrailVersion: "1"
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID_1
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY_1
|
||||
aws_region_name: "us-east-1"
|
||||
|
||||
# Second Bedrock guardrail - AWS Account 2
|
||||
- guardrail_name: "content-filter"
|
||||
litellm_params:
|
||||
guardrail: bedrock/guardrail
|
||||
mode: "pre_call"
|
||||
guardrailIdentifier: "def456"
|
||||
guardrailVersion: "1"
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID_2
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY_2
|
||||
aws_region_name: "us-west-2"
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="custom" label="Custom Guardrails">
|
||||
|
||||
```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:
|
||||
# First custom guardrail instance
|
||||
- guardrail_name: "pii-filter"
|
||||
litellm_params:
|
||||
guardrail: custom_guardrail.PIIFilterA
|
||||
mode: "pre_call"
|
||||
|
||||
# Second custom guardrail instance
|
||||
- guardrail_name: "pii-filter"
|
||||
litellm_params:
|
||||
guardrail: custom_guardrail.PIIFilterB
|
||||
mode: "pre_call"
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="aporia" label="Aporia Guardrails">
|
||||
|
||||
```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:
|
||||
# First Aporia instance
|
||||
- guardrail_name: "toxicity-filter"
|
||||
litellm_params:
|
||||
guardrail: aporia
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/APORIA_API_KEY_1
|
||||
api_base: os.environ/APORIA_API_BASE_1
|
||||
|
||||
# Second Aporia instance
|
||||
- guardrail_name: "toxicity-filter"
|
||||
litellm_params:
|
||||
guardrail: aporia
|
||||
mode: "pre_call"
|
||||
api_key: os.environ/APORIA_API_KEY_2
|
||||
api_base: os.environ/APORIA_API_BASE_2
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
### 2. Start LiteLLM Gateway
|
||||
|
||||
```bash showLineNumbers title="Start proxy"
|
||||
litellm --config config.yaml --detailed_debug
|
||||
```
|
||||
|
||||
### 3. Make Requests
|
||||
|
||||
Requests using the guardrail will be automatically load balanced:
|
||||
|
||||
```bash showLineNumbers title="Test request"
|
||||
curl -X POST http://localhost:4000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "gpt-4",
|
||||
"messages": [{"role": "user", "content": "Hello, how are you?"}],
|
||||
"guardrails": ["content-filter"]
|
||||
}'
|
||||
```
|
||||
|
||||
## Weighted Load Balancing
|
||||
|
||||
Assign weights to distribute traffic unevenly across guardrail instances:
|
||||
|
||||
```yaml showLineNumbers title="config.yaml - Weighted distribution"
|
||||
guardrails:
|
||||
# 80% of traffic
|
||||
- guardrail_name: "content-filter"
|
||||
litellm_params:
|
||||
guardrail: bedrock/guardrail
|
||||
mode: "pre_call"
|
||||
guardrailIdentifier: "primary-guard"
|
||||
guardrailVersion: "1"
|
||||
weight: 8 # Higher weight = more traffic
|
||||
|
||||
# 20% of traffic
|
||||
- guardrail_name: "content-filter"
|
||||
litellm_params:
|
||||
guardrail: bedrock/guardrail
|
||||
mode: "pre_call"
|
||||
guardrailIdentifier: "secondary-guard"
|
||||
guardrailVersion: "1"
|
||||
weight: 2 # Lower weight = less traffic
|
||||
```
|
||||
|
||||
## Bedrock Guardrails - Multi-Account Setup
|
||||
|
||||
AWS Bedrock Guardrails have rate limits per account. Here's how to set up load balancing across multiple AWS accounts:
|
||||
|
||||
### Architecture
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph LiteLLM["LiteLLM Gateway"]
|
||||
LB[Load Balancer]
|
||||
end
|
||||
|
||||
subgraph AWS1["AWS Account 1 (us-east-1)"]
|
||||
BG1[Bedrock Guardrail]
|
||||
end
|
||||
|
||||
subgraph AWS2["AWS Account 2 (us-west-2)"]
|
||||
BG2[Bedrock Guardrail]
|
||||
end
|
||||
|
||||
subgraph AWS3["AWS Account 3 (eu-west-1)"]
|
||||
BG3[Bedrock Guardrail]
|
||||
end
|
||||
|
||||
Client[Client] --> LiteLLM
|
||||
LB --> BG1
|
||||
LB --> BG2
|
||||
LB --> BG3
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
```yaml showLineNumbers title="config.yaml - Multi-account Bedrock"
|
||||
model_list:
|
||||
- model_name: claude-3
|
||||
litellm_params:
|
||||
model: bedrock/anthropic.claude-3-sonnet-20240229-v1:0
|
||||
|
||||
guardrails:
|
||||
# AWS Account 1 - US East
|
||||
- guardrail_name: "bedrock-content-filter"
|
||||
litellm_params:
|
||||
guardrail: bedrock/guardrail
|
||||
mode: "during_call"
|
||||
guardrailIdentifier: "guard-us-east"
|
||||
guardrailVersion: "DRAFT"
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_1
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_KEY_1
|
||||
aws_region_name: "us-east-1"
|
||||
|
||||
# AWS Account 2 - US West
|
||||
- guardrail_name: "bedrock-content-filter"
|
||||
litellm_params:
|
||||
guardrail: bedrock/guardrail
|
||||
mode: "during_call"
|
||||
guardrailIdentifier: "guard-us-west"
|
||||
guardrailVersion: "DRAFT"
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_2
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_KEY_2
|
||||
aws_region_name: "us-west-2"
|
||||
|
||||
# AWS Account 3 - EU West
|
||||
- guardrail_name: "bedrock-content-filter"
|
||||
litellm_params:
|
||||
guardrail: bedrock/guardrail
|
||||
mode: "during_call"
|
||||
guardrailIdentifier: "guard-eu-west"
|
||||
guardrailVersion: "DRAFT"
|
||||
aws_access_key_id: os.environ/AWS_ACCESS_KEY_3
|
||||
aws_secret_access_key: os.environ/AWS_SECRET_KEY_3
|
||||
aws_region_name: "eu-west-1"
|
||||
```
|
||||
|
||||
### Test Multi-Account Setup
|
||||
|
||||
```bash showLineNumbers title="Run multiple requests to verify load balancing"
|
||||
# Run 10 requests - they will be distributed across accounts
|
||||
for i in {1..10}; do
|
||||
curl -s -X POST http://localhost:4000/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer sk-1234" \
|
||||
-d '{
|
||||
"model": "claude-3",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"guardrails": ["bedrock-content-filter"]
|
||||
}' &
|
||||
done
|
||||
wait
|
||||
```
|
||||
|
||||
Check proxy logs to verify requests are distributed across different AWS accounts.
|
||||
|
||||
## Custom Guardrails Example
|
||||
|
||||
Create two custom guardrail classes for load balancing:
|
||||
|
||||
```python showLineNumbers title="custom_guardrail.py"
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.caching.caching import DualCache
|
||||
|
||||
|
||||
class PIIFilterA(CustomGuardrail):
|
||||
"""PII Filter Instance A"""
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
print("PIIFilterA processing request")
|
||||
# Your PII filtering logic here
|
||||
return data
|
||||
|
||||
|
||||
class PIIFilterB(CustomGuardrail):
|
||||
"""PII Filter Instance B"""
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: str,
|
||||
):
|
||||
print("PIIFilterB processing request")
|
||||
# Your PII filtering logic here
|
||||
return data
|
||||
```
|
||||
|
||||
```yaml showLineNumbers title="config.yaml"
|
||||
guardrails:
|
||||
- guardrail_name: "pii-filter"
|
||||
litellm_params:
|
||||
guardrail: custom_guardrail.PIIFilterA
|
||||
mode: "pre_call"
|
||||
|
||||
- guardrail_name: "pii-filter"
|
||||
litellm_params:
|
||||
guardrail: custom_guardrail.PIIFilterB
|
||||
mode: "pre_call"
|
||||
```
|
||||
|
||||
## Verifying Load Balancing
|
||||
|
||||
Enable detailed debug logging to verify load balancing is working:
|
||||
|
||||
```bash showLineNumbers title="Start with debug logging"
|
||||
litellm --config config.yaml --detailed_debug
|
||||
```
|
||||
|
||||
You should see logs indicating which guardrail instance is selected:
|
||||
|
||||
```
|
||||
Selected guardrail deployment: bedrock/guardrail (guard-us-east)
|
||||
Selected guardrail deployment: bedrock/guardrail (guard-us-west)
|
||||
Selected guardrail deployment: bedrock/guardrail (guard-eu-west)
|
||||
...
|
||||
```
|
||||
|
||||
## Related
|
||||
|
||||
- [Guardrails Quick Start](./quick_start.md)
|
||||
- [Bedrock Guardrails](./bedrock.md)
|
||||
- [Custom Guardrails](./custom_guardrail.md)
|
||||
- [Load Balancing for LLM Calls](../load_balancing.md)
|
||||
|
||||
@ -69,6 +69,13 @@ guardrails:
|
||||
- `during_call` Run **during** LLM call, on **input** Same as `pre_call` but runs in parallel as LLM call. Response not returned until guardrail check completes
|
||||
- A list of the above values to run multiple modes, e.g. `mode: [pre_call, post_call]`
|
||||
|
||||
### Load Balancing Guardrails
|
||||
|
||||
Need to distribute guardrail requests across multiple accounts or regions? See [Guardrail Load Balancing](./guardrail_load_balancing.md) for details on:
|
||||
- Load balancing across multiple AWS Bedrock accounts (useful for rate limit management)
|
||||
- Weighted distribution across guardrail instances
|
||||
- Multi-region guardrail deployments
|
||||
|
||||
|
||||
## 2. Start LiteLLM Gateway
|
||||
|
||||
|
||||
@ -42,6 +42,7 @@ const sidebars = {
|
||||
label: "Guardrails",
|
||||
items: [
|
||||
"proxy/guardrails/quick_start",
|
||||
"proxy/guardrails/guardrail_load_balancing",
|
||||
{
|
||||
type: "category",
|
||||
"label": "Contributing to Guardrails",
|
||||
|
||||
@ -12354,6 +12354,7 @@
|
||||
"max_videos_per_prompt": 10,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.039,
|
||||
"output_cost_per_image_token": 3e-05,
|
||||
"output_cost_per_reasoning_token": 2.5e-06,
|
||||
"output_cost_per_token": 2.5e-06,
|
||||
"rpm": 100000,
|
||||
@ -12402,6 +12403,7 @@
|
||||
"max_videos_per_prompt": 10,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.039,
|
||||
"output_cost_per_image_token": 3e-05,
|
||||
"output_cost_per_reasoning_token": 3e-05,
|
||||
"output_cost_per_token": 3e-05,
|
||||
"rpm": 100000,
|
||||
@ -14134,6 +14136,7 @@
|
||||
"max_videos_per_prompt": 10,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.039,
|
||||
"output_cost_per_image_token": 3e-05,
|
||||
"output_cost_per_reasoning_token": 2.5e-06,
|
||||
"output_cost_per_token": 2.5e-06,
|
||||
"rpm": 100000,
|
||||
@ -14182,6 +14185,7 @@
|
||||
"max_videos_per_prompt": 10,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.039,
|
||||
"output_cost_per_image_token": 3e-05,
|
||||
"output_cost_per_reasoning_token": 3e-05,
|
||||
"output_cost_per_token": 3e-05,
|
||||
"rpm": 100000,
|
||||
@ -27288,6 +27292,7 @@
|
||||
"max_videos_per_prompt": 10,
|
||||
"mode": "image_generation",
|
||||
"output_cost_per_image": 0.039,
|
||||
"output_cost_per_image_token": 3e-05,
|
||||
"output_cost_per_reasoning_token": 2.5e-06,
|
||||
"output_cost_per_token": 2.5e-06,
|
||||
"rpm": 100000,
|
||||
|
||||
@ -8,6 +8,43 @@ from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
|
||||
# Global counter for tracking which guardrail was called (for load balancing tests)
|
||||
guardrail_lb_call_count: Dict[str, int] = {"A": 0, "B": 0}
|
||||
|
||||
|
||||
class GuardrailForLBTestingA(CustomGuardrail):
|
||||
"""Guardrail A for load balancing testing."""
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: CallTypesLiteral,
|
||||
) -> Optional[Union[Exception, str, dict]]:
|
||||
guardrail_lb_call_count["A"] += 1
|
||||
verbose_proxy_logger.info(
|
||||
f"GuardrailForLBTestingA called. Total A calls: {guardrail_lb_call_count['A']}"
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
class GuardrailForLBTestingB(CustomGuardrail):
|
||||
"""Guardrail B for load balancing testing."""
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
cache: DualCache,
|
||||
data: dict,
|
||||
call_type: CallTypesLiteral,
|
||||
) -> Optional[Union[Exception, str, dict]]:
|
||||
guardrail_lb_call_count["B"] += 1
|
||||
verbose_proxy_logger.info(
|
||||
f"GuardrailForLBTestingB called. Total B calls: {guardrail_lb_call_count['B']}"
|
||||
)
|
||||
return data
|
||||
|
||||
|
||||
class myCustomGuardrail(CustomGuardrail):
|
||||
def __init__(
|
||||
|
||||
@ -78,6 +78,15 @@ guardrails:
|
||||
litellm_params:
|
||||
guardrail: custom_guardrail.myCustomGuardrail
|
||||
mode: "post_call"
|
||||
# Load balancing guardrails - two guardrails with same name
|
||||
- guardrail_name: "lb-test-guard"
|
||||
litellm_params:
|
||||
guardrail: custom_guardrail.GuardrailForLBTestingA
|
||||
mode: "pre_call"
|
||||
- guardrail_name: "lb-test-guard"
|
||||
litellm_params:
|
||||
guardrail: custom_guardrail.GuardrailForLBTestingB
|
||||
mode: "pre_call"
|
||||
|
||||
router_settings:
|
||||
enable_tag_filtering: True # 👈 Key Change
|
||||
@ -1,4 +1,4 @@
|
||||
from typing import Dict, List, Optional, cast
|
||||
from typing import Any, Dict, List, Optional, cast
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
@ -36,6 +36,67 @@ def init_guardrails_v2(
|
||||
|
||||
verbose_proxy_logger.debug(f"\nGuardrail List:{guardrail_list}\n")
|
||||
|
||||
# Populate router's guardrail_list for load balancing support
|
||||
_populate_router_guardrail_list(guardrail_list=guardrail_list)
|
||||
|
||||
|
||||
def _populate_router_guardrail_list(guardrail_list: List[Guardrail]) -> None:
|
||||
"""
|
||||
Populate the router's guardrail_list from initialized guardrails.
|
||||
|
||||
This enables load balancing across multiple guardrail deployments
|
||||
with the same guardrail_name.
|
||||
"""
|
||||
from litellm.proxy.guardrails.guardrail_registry import IN_MEMORY_GUARDRAIL_HANDLER
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
from litellm.types.router import GuardrailTypedDict
|
||||
|
||||
if llm_router is None:
|
||||
verbose_proxy_logger.debug(
|
||||
"Router not initialized yet, skipping guardrail_list population"
|
||||
)
|
||||
return
|
||||
|
||||
router_guardrail_list: List[GuardrailTypedDict] = []
|
||||
|
||||
for guardrail in guardrail_list:
|
||||
guardrail_id = guardrail.get("guardrail_id")
|
||||
guardrail_name = guardrail.get("guardrail_name")
|
||||
litellm_params: Any = guardrail.get("litellm_params", {})
|
||||
|
||||
# Get the callback instance from the registry
|
||||
callback = None
|
||||
if guardrail_id:
|
||||
callback = IN_MEMORY_GUARDRAIL_HANDLER.guardrail_id_to_custom_guardrail.get(
|
||||
guardrail_id
|
||||
)
|
||||
|
||||
# Build litellm_params dict for the router
|
||||
params_dict = (
|
||||
litellm_params.model_dump()
|
||||
if hasattr(litellm_params, "model_dump")
|
||||
else dict(litellm_params)
|
||||
)
|
||||
|
||||
router_guardrail: GuardrailTypedDict = GuardrailTypedDict(
|
||||
guardrail_name=guardrail_name or "",
|
||||
litellm_params={
|
||||
"guardrail": params_dict.get("guardrail", ""),
|
||||
"mode": params_dict.get("mode", ""),
|
||||
"api_key": params_dict.get("api_key"),
|
||||
"api_base": params_dict.get("api_base"),
|
||||
},
|
||||
callback=callback,
|
||||
id=guardrail_id,
|
||||
)
|
||||
|
||||
router_guardrail_list.append(router_guardrail)
|
||||
|
||||
llm_router.guardrail_list = router_guardrail_list
|
||||
verbose_proxy_logger.debug(
|
||||
f"Populated router guardrail_list with {len(router_guardrail_list)} guardrails"
|
||||
)
|
||||
|
||||
|
||||
### LEGACY IMPLEMENTATION ###
|
||||
def initialize_guardrails(
|
||||
|
||||
@ -36,10 +36,18 @@ from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import CallTypes, CallTypesLiteral
|
||||
|
||||
try:
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import BaseEmailLogger
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import SendGridEmailLogger
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import SMTPEmailLogger
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import (
|
||||
BaseEmailLogger,
|
||||
)
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import (
|
||||
ResendEmailLogger,
|
||||
)
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import (
|
||||
SendGridEmailLogger,
|
||||
)
|
||||
from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import (
|
||||
SMTPEmailLogger,
|
||||
)
|
||||
except ImportError:
|
||||
BaseEmailLogger = None # type: ignore
|
||||
SendGridEmailLogger = None # type: ignore
|
||||
@ -813,6 +821,125 @@ class ProxyLogging:
|
||||
raise HTTPException(status_code=400, detail={"error": response})
|
||||
return data
|
||||
|
||||
def _should_use_guardrail_load_balancing(
|
||||
self,
|
||||
guardrail_name: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if load balancing should be used for this guardrail.
|
||||
|
||||
Returns True if the router has multiple deployments for this guardrail name.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
if llm_router is None or not hasattr(llm_router, "guardrail_list"):
|
||||
return False
|
||||
|
||||
matching = [
|
||||
g
|
||||
for g in llm_router.guardrail_list
|
||||
if g.get("guardrail_name") == guardrail_name
|
||||
]
|
||||
return len(matching) > 1
|
||||
|
||||
async def _execute_guardrail_hook(
|
||||
self,
|
||||
callback: "CustomGuardrail",
|
||||
hook_type: str,
|
||||
data: dict,
|
||||
user_api_key_dict: Optional[UserAPIKeyAuth],
|
||||
call_type: CallTypesLiteral,
|
||||
response: Optional[Any] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Execute a single guardrail's hook.
|
||||
|
||||
Args:
|
||||
callback: The guardrail callback to execute
|
||||
hook_type: One of "pre_call", "during_call", "post_call"
|
||||
data: Request data
|
||||
user_api_key_dict: User API key auth
|
||||
call_type: Type of call
|
||||
response: Response object (for post_call hooks)
|
||||
|
||||
Returns:
|
||||
Result from the guardrail execution
|
||||
"""
|
||||
# Use unified_guardrail if callback has apply_guardrail method
|
||||
use_unified = "apply_guardrail" in type(callback).__dict__
|
||||
if use_unified:
|
||||
data["guardrail_to_apply"] = callback
|
||||
|
||||
target = unified_guardrail if use_unified else callback
|
||||
|
||||
if hook_type == "pre_call":
|
||||
return await target.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, # type: ignore
|
||||
cache=self.call_details["user_api_key_cache"],
|
||||
data=data,
|
||||
call_type=call_type,
|
||||
)
|
||||
elif hook_type == "during_call":
|
||||
return await target.async_moderation_hook(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict, # type: ignore
|
||||
call_type=call_type,
|
||||
)
|
||||
elif hook_type == "post_call":
|
||||
return await target.async_post_call_success_hook(
|
||||
user_api_key_dict=user_api_key_dict, # type: ignore
|
||||
data=data,
|
||||
response=response, # type: ignore
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown hook_type: {hook_type}")
|
||||
|
||||
async def _execute_guardrail_with_load_balancing(
|
||||
self,
|
||||
guardrail_name: str,
|
||||
hook_type: str,
|
||||
data: dict,
|
||||
user_api_key_dict: Optional[UserAPIKeyAuth],
|
||||
call_type: CallTypesLiteral,
|
||||
response: Optional[Any] = None,
|
||||
) -> Any:
|
||||
"""
|
||||
Execute a guardrail using the router's load balancing.
|
||||
|
||||
Args:
|
||||
guardrail_name: Name of the guardrail
|
||||
hook_type: One of "pre_call", "during_call", "post_call"
|
||||
data: Request data
|
||||
user_api_key_dict: User API key auth
|
||||
call_type: Type of call
|
||||
response: Response object (for post_call hooks)
|
||||
|
||||
Returns:
|
||||
Result from the guardrail execution
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
if llm_router is None:
|
||||
raise ValueError("Router not initialized")
|
||||
|
||||
# Select guardrail using router's load balancing
|
||||
selected_guardrail = llm_router.get_available_guardrail(
|
||||
guardrail_name=guardrail_name
|
||||
)
|
||||
|
||||
callback = selected_guardrail.get("callback")
|
||||
if callback is None:
|
||||
raise ValueError(f"No callback found for guardrail: {guardrail_name}")
|
||||
|
||||
return await self._execute_guardrail_hook(
|
||||
callback=callback,
|
||||
hook_type=hook_type,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=call_type,
|
||||
response=response,
|
||||
)
|
||||
|
||||
async def _process_guardrail_callback(
|
||||
self,
|
||||
callback: CustomGuardrail,
|
||||
@ -823,6 +950,8 @@ class ProxyLogging:
|
||||
"""
|
||||
Process a guardrail callback during pre-call hook.
|
||||
|
||||
Supports load balancing when multiple guardrail deployments exist.
|
||||
|
||||
Args:
|
||||
callback: The CustomGuardrail callback to process
|
||||
data: The request data dictionary
|
||||
@ -843,23 +972,25 @@ class ProxyLogging:
|
||||
if callback.should_run_guardrail(data=data, event_type=event_type) is not True:
|
||||
return None
|
||||
|
||||
# Execute the appropriate guardrail hook
|
||||
if "apply_guardrail" in type(callback).__dict__:
|
||||
# Use unified guardrail for callbacks with apply_guardrail method
|
||||
data["guardrail_to_apply"] = callback
|
||||
response = await unified_guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, # type: ignore
|
||||
cache=self.call_details["user_api_key_cache"],
|
||||
data=data, # type: ignore
|
||||
call_type=call_type, # type: ignore
|
||||
guardrail_name = callback.guardrail_name
|
||||
|
||||
# Check if load balancing should be used
|
||||
if guardrail_name and self._should_use_guardrail_load_balancing(guardrail_name):
|
||||
response = await self._execute_guardrail_with_load_balancing(
|
||||
guardrail_name=guardrail_name,
|
||||
hook_type="pre_call",
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=call_type,
|
||||
)
|
||||
else:
|
||||
# Use the callback's own async_pre_call_hook method
|
||||
response = await callback.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, # type: ignore
|
||||
cache=self.call_details["user_api_key_cache"],
|
||||
data=data, # type: ignore
|
||||
call_type=call_type, # type: ignore
|
||||
# Single guardrail - execute directly
|
||||
response = await self._execute_guardrail_hook(
|
||||
callback=callback,
|
||||
hook_type="pre_call",
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
# Process the response if one was returned
|
||||
@ -3583,7 +3714,10 @@ async def _monitor_spend_logs_queue(
|
||||
db_writer_client: Optional HTTP handler for external spend logs endpoint
|
||||
proxy_logging_obj: Proxy logging object
|
||||
"""
|
||||
from litellm.constants import SPEND_LOG_QUEUE_SIZE_THRESHOLD, SPEND_LOG_QUEUE_POLL_INTERVAL
|
||||
from litellm.constants import (
|
||||
SPEND_LOG_QUEUE_POLL_INTERVAL,
|
||||
SPEND_LOG_QUEUE_SIZE_THRESHOLD,
|
||||
)
|
||||
|
||||
threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD
|
||||
base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL
|
||||
|
||||
@ -136,6 +136,7 @@ from litellm.types.router import (
|
||||
CustomRoutingStrategyBase,
|
||||
Deployment,
|
||||
DeploymentTypedDict,
|
||||
GuardrailTypedDict,
|
||||
LiteLLM_Params,
|
||||
MockRouterTestingParams,
|
||||
ModelGroupInfo,
|
||||
@ -214,6 +215,8 @@ class Router:
|
||||
assistants_config: Optional[AssistantsTypedDict] = None,
|
||||
## SEARCH API ##
|
||||
search_tools: Optional[List[SearchToolTypedDict]] = None,
|
||||
## GUARDRAIL API ##
|
||||
guardrail_list: Optional[List[GuardrailTypedDict]] = None,
|
||||
## CACHING ##
|
||||
redis_url: Optional[str] = None,
|
||||
redis_host: Optional[str] = None,
|
||||
@ -375,6 +378,7 @@ class Router:
|
||||
|
||||
self.assistants_config = assistants_config
|
||||
self.search_tools = search_tools or []
|
||||
self.guardrail_list = guardrail_list or []
|
||||
self.deployment_names: List = (
|
||||
[]
|
||||
) # names of models under litellm_params. ex. azure/chatgpt-v-2
|
||||
@ -2974,6 +2978,99 @@ class Router:
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
async def aguardrail(
|
||||
self,
|
||||
guardrail_name: str,
|
||||
original_function: Callable,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Execute a guardrail with load balancing and fallbacks.
|
||||
|
||||
Args:
|
||||
guardrail_name: Name of the guardrail to execute
|
||||
original_function: The guardrail's execution function (e.g., async_pre_call_hook)
|
||||
**kwargs: Additional arguments passed to the guardrail
|
||||
|
||||
Returns:
|
||||
Result from the guardrail execution
|
||||
"""
|
||||
kwargs["model"] = guardrail_name # For fallback system compatibility
|
||||
kwargs["original_generic_function"] = original_function
|
||||
kwargs["original_function"] = self._aguardrail_helper
|
||||
self._update_kwargs_before_fallbacks(
|
||||
model=guardrail_name, kwargs=kwargs, metadata_variable_name="litellm_metadata"
|
||||
)
|
||||
verbose_router_logger.debug(
|
||||
f"Inside aguardrail() - guardrail_name: {guardrail_name}; kwargs: {kwargs}"
|
||||
)
|
||||
response = await self.async_function_with_fallbacks(**kwargs)
|
||||
return response
|
||||
|
||||
async def _aguardrail_helper(
|
||||
self,
|
||||
model: str,
|
||||
original_generic_function: Callable,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Helper for aguardrail - selects a guardrail deployment and executes it.
|
||||
Called by async_function_with_fallbacks for each retry attempt.
|
||||
|
||||
Args:
|
||||
model: The guardrail_name (named 'model' for fallback system compatibility)
|
||||
original_generic_function: The guardrail's execution function
|
||||
**kwargs: Additional arguments
|
||||
"""
|
||||
guardrail_name = model
|
||||
selected_guardrail = self.get_available_guardrail(
|
||||
guardrail_name=guardrail_name,
|
||||
)
|
||||
|
||||
verbose_router_logger.debug(
|
||||
f"Selected guardrail deployment: {selected_guardrail.get('litellm_params', {}).get('guardrail')}"
|
||||
)
|
||||
|
||||
# Pass the selected guardrail config to the original function
|
||||
kwargs["selected_guardrail"] = selected_guardrail
|
||||
response = await original_generic_function(**kwargs)
|
||||
return response
|
||||
|
||||
def get_available_guardrail(
|
||||
self,
|
||||
guardrail_name: str,
|
||||
) -> "GuardrailTypedDict":
|
||||
"""
|
||||
Select a guardrail deployment using the router's load balancing strategy.
|
||||
|
||||
Args:
|
||||
guardrail_name: Name of the guardrail to select
|
||||
|
||||
Returns:
|
||||
Selected guardrail configuration dict
|
||||
"""
|
||||
from litellm.router_strategy.simple_shuffle import simple_shuffle
|
||||
|
||||
healthy_deployments = [
|
||||
g for g in self.guardrail_list if g.get("guardrail_name") == guardrail_name
|
||||
]
|
||||
|
||||
if not healthy_deployments:
|
||||
raise ValueError(f"No guardrail found with name: {guardrail_name}")
|
||||
|
||||
if len(healthy_deployments) == 1:
|
||||
return healthy_deployments[0]
|
||||
|
||||
# Use simple_shuffle for weighted selection
|
||||
return cast(
|
||||
GuardrailTypedDict,
|
||||
simple_shuffle(
|
||||
llm_router_instance=self,
|
||||
healthy_deployments=healthy_deployments,
|
||||
model=guardrail_name,
|
||||
),
|
||||
)
|
||||
|
||||
async def _ageneric_api_call_with_fallbacks(
|
||||
self, model: str, original_function: Callable, **kwargs
|
||||
):
|
||||
|
||||
@ -637,6 +637,29 @@ class SearchToolTypedDict(TypedDict):
|
||||
litellm_params: Required[SearchToolLiteLLMParams]
|
||||
|
||||
|
||||
class GuardrailLiteLLMParams(TypedDict, total=False):
|
||||
"""
|
||||
LiteLLM params for guardrails.
|
||||
"""
|
||||
|
||||
guardrail: Required[str]
|
||||
mode: Required[str]
|
||||
api_key: Optional[str]
|
||||
api_base: Optional[str]
|
||||
weight: Optional[int] # For load balancing
|
||||
|
||||
|
||||
class GuardrailTypedDict(TypedDict, total=False):
|
||||
"""
|
||||
Configuration for a guardrail in the router.
|
||||
"""
|
||||
|
||||
guardrail_name: Required[str]
|
||||
litellm_params: Required[GuardrailLiteLLMParams]
|
||||
callback: Any # The CustomGuardrail instance
|
||||
id: Optional[str] # Unique identifier for the guardrail deployment
|
||||
|
||||
|
||||
class FineTuningConfig(BaseModel):
|
||||
custom_llm_provider: Literal["azure", "openai"]
|
||||
|
||||
|
||||
105
tests/guardrails_tests/test_guardrail_load_balancing.py
Normal file
105
tests/guardrails_tests/test_guardrail_load_balancing.py
Normal file
@ -0,0 +1,105 @@
|
||||
"""
|
||||
Test guardrail load balancing through the Router and ProxyLogging.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch, AsyncMock
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
from litellm import Router
|
||||
from litellm.caching import DualCache
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
|
||||
class MockGuardrail(CustomGuardrail):
|
||||
"""Mock guardrail that tracks calls."""
|
||||
|
||||
call_count = 0
|
||||
|
||||
def __init__(self, guardrail_name: str, guardrail_id: str):
|
||||
super().__init__(guardrail_name=guardrail_name)
|
||||
self.guardrail_id = guardrail_id
|
||||
self.calls = 0
|
||||
|
||||
def should_run_guardrail(self, data, event_type) -> bool:
|
||||
return True
|
||||
|
||||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||||
self.calls += 1
|
||||
MockGuardrail.call_count += 1
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_logging_pre_call_hook_load_balancing():
|
||||
"""Test that async_pre_call_hook load balances across multiple guardrails."""
|
||||
# Reset call count
|
||||
MockGuardrail.call_count = 0
|
||||
|
||||
# Create two mock guardrails with same name
|
||||
guardrail_1 = MockGuardrail(guardrail_name="content-filter", guardrail_id="g1")
|
||||
guardrail_2 = MockGuardrail(guardrail_name="content-filter", guardrail_id="g2")
|
||||
|
||||
# Create router with multiple guardrails of same name
|
||||
guardrail_list = [
|
||||
{
|
||||
"guardrail_name": "content-filter",
|
||||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||||
"callback": guardrail_1,
|
||||
"id": "guardrail-1",
|
||||
},
|
||||
{
|
||||
"guardrail_name": "content-filter",
|
||||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||||
"callback": guardrail_2,
|
||||
"id": "guardrail-2",
|
||||
},
|
||||
]
|
||||
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-4",
|
||||
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
|
||||
}
|
||||
],
|
||||
guardrail_list=guardrail_list,
|
||||
)
|
||||
|
||||
# Create ProxyLogging instance
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
|
||||
# Add guardrail to litellm.callbacks so it gets picked up
|
||||
original_callbacks = litellm.callbacks.copy()
|
||||
litellm.callbacks = [guardrail_1]
|
||||
|
||||
try:
|
||||
with patch("litellm.proxy.proxy_server.llm_router", router):
|
||||
# Call pre_call_hook 50 times
|
||||
for _ in range(50):
|
||||
await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=MagicMock(),
|
||||
data={"messages": [{"role": "user", "content": "test"}]},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
# Both guardrails should have been called (load balanced)
|
||||
assert guardrail_1.calls > 0, "Guardrail 1 should have been called"
|
||||
assert guardrail_2.calls > 0, "Guardrail 2 should have been called"
|
||||
|
||||
# Total calls should be 50
|
||||
total = guardrail_1.calls + guardrail_2.calls
|
||||
assert total == 50, f"Expected 50 total calls, got {total}"
|
||||
|
||||
# Verify reasonable distribution (not all to one)
|
||||
min_calls = min(guardrail_1.calls, guardrail_2.calls)
|
||||
assert min_calls >= 10, f"Expected at least 10 calls to each guardrail, got min={min_calls}"
|
||||
|
||||
finally:
|
||||
litellm.callbacks = original_callbacks
|
||||
@ -315,3 +315,46 @@ async def test_guardrails_with_team_controls():
|
||||
|
||||
assert "x-litellm-applied-guardrails" in headers
|
||||
assert headers["x-litellm-applied-guardrails"] == "bedrock-pre-guard"
|
||||
|
||||
|
||||
async def get_guardrail_lb_counts(session):
|
||||
"""Get the current guardrail load balancing call counts from the proxy."""
|
||||
url = "http://0.0.0.0:4000/guardrail/lb/counts"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
|
||||
async with session.get(url, headers=headers) as response:
|
||||
if response.status == 200:
|
||||
return await response.json()
|
||||
return None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_guardrail_load_balancing():
|
||||
"""
|
||||
Test that guardrail load balancing distributes requests across multiple guardrail instances.
|
||||
|
||||
- Make 20 requests with the lb-test-guard guardrail
|
||||
- Verify that both GuardrailForLBTestingA and GuardrailForLBTestingB are called
|
||||
- Verify reasonable distribution (both should have at least some calls)
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
num_requests = 20
|
||||
|
||||
# Make multiple requests with the load-balanced guardrail
|
||||
for i in range(num_requests):
|
||||
response, headers = await chat_completion(
|
||||
session,
|
||||
"sk-1234",
|
||||
model="fake-openai-endpoint",
|
||||
messages=[{"role": "user", "content": f"Hello request {i}"}],
|
||||
guardrails=["lb-test-guard"],
|
||||
)
|
||||
|
||||
# Verify guardrail was applied
|
||||
assert "x-litellm-applied-guardrails" in headers
|
||||
assert headers["x-litellm-applied-guardrails"] == "lb-test-guard"
|
||||
|
||||
# All requests should succeed - the test passes if we get here
|
||||
# The actual load balancing verification is done by checking proxy logs
|
||||
# which should show alternating calls to GuardrailForLBTestingA and GuardrailForLBTestingB
|
||||
print(f"Successfully made {num_requests} requests with load-balanced guardrail")
|
||||
|
||||
@ -1724,3 +1724,148 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint()
|
||||
assert credentials["aws_secret_access_key"] == "test-secret-key"
|
||||
assert credentials["aws_region_name"] == "us-east-1"
|
||||
assert credentials["custom_llm_provider"] == "bedrock"
|
||||
|
||||
|
||||
def test_get_available_guardrail_single_deployment():
|
||||
"""
|
||||
Test get_available_guardrail returns the single guardrail when only one exists.
|
||||
"""
|
||||
guardrail_config = {
|
||||
"guardrail_name": "content-filter",
|
||||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||||
"id": "guardrail-1",
|
||||
}
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
}
|
||||
],
|
||||
guardrail_list=[guardrail_config],
|
||||
)
|
||||
|
||||
result = router.get_available_guardrail(guardrail_name="content-filter")
|
||||
assert result == guardrail_config
|
||||
|
||||
|
||||
def test_get_available_guardrail_multiple_deployments():
|
||||
"""
|
||||
Test get_available_guardrail load balances across multiple guardrails.
|
||||
"""
|
||||
guardrail_1 = {
|
||||
"guardrail_name": "content-filter",
|
||||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||||
"id": "guardrail-1",
|
||||
}
|
||||
guardrail_2 = {
|
||||
"guardrail_name": "content-filter",
|
||||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||||
"id": "guardrail-2",
|
||||
}
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
}
|
||||
],
|
||||
guardrail_list=[guardrail_1, guardrail_2],
|
||||
)
|
||||
|
||||
# Call multiple times to verify load balancing
|
||||
results = set()
|
||||
for _ in range(20):
|
||||
result = router.get_available_guardrail(guardrail_name="content-filter")
|
||||
results.add(result["id"])
|
||||
|
||||
# Both guardrails should be selected at least once
|
||||
assert "guardrail-1" in results or "guardrail-2" in results
|
||||
|
||||
|
||||
def test_get_available_guardrail_not_found():
|
||||
"""
|
||||
Test get_available_guardrail raises ValueError when guardrail not found.
|
||||
"""
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
}
|
||||
],
|
||||
guardrail_list=[],
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="No guardrail found with name"):
|
||||
router.get_available_guardrail(guardrail_name="non-existent")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aguardrail_helper():
|
||||
"""
|
||||
Test _aguardrail_helper selects a guardrail and executes the original function.
|
||||
"""
|
||||
guardrail_config = {
|
||||
"guardrail_name": "content-filter",
|
||||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||||
"id": "guardrail-1",
|
||||
}
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
}
|
||||
],
|
||||
guardrail_list=[guardrail_config],
|
||||
)
|
||||
|
||||
# Mock the original function
|
||||
async def mock_original_function(**kwargs):
|
||||
return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")}
|
||||
|
||||
result = await router._aguardrail_helper(
|
||||
model="content-filter",
|
||||
original_generic_function=mock_original_function,
|
||||
)
|
||||
|
||||
assert result["result"] == "success"
|
||||
assert result["selected_guardrail"] == guardrail_config
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aguardrail():
|
||||
"""
|
||||
Test aguardrail executes a guardrail with load balancing and fallbacks.
|
||||
"""
|
||||
guardrail_config = {
|
||||
"guardrail_name": "content-filter",
|
||||
"litellm_params": {"guardrail": "custom", "mode": "pre_call"},
|
||||
"id": "guardrail-1",
|
||||
}
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-3.5-turbo",
|
||||
"litellm_params": {"model": "gpt-3.5-turbo"},
|
||||
}
|
||||
],
|
||||
guardrail_list=[guardrail_config],
|
||||
)
|
||||
|
||||
# Mock the original function
|
||||
async def mock_original_function(**kwargs):
|
||||
return {"result": "success", "selected_guardrail": kwargs.get("selected_guardrail")}
|
||||
|
||||
result = await router.aguardrail(
|
||||
guardrail_name="content-filter",
|
||||
original_function=mock_original_function,
|
||||
)
|
||||
|
||||
assert result["result"] == "success"
|
||||
assert result["selected_guardrail"]["id"] == "guardrail-1"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user