Custom Auth
You can now override the default api key auth.
Overview
Custom Auth
You can now override the default api key auth.
Usage
1. Create a custom auth file.
Make sure the response type follows the UserAPIKeyAuth pydantic object. This is used by for logging usage specific to that user key.
from fastapi import Request
from litellm.proxy._types import UserAPIKeyAuth
async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth:
try:
modified_master_key = "sk-my-master-key"
if api_key == modified_master_key:
return UserAPIKeyAuth(api_key=api_key)
raise Exception
except:
raise Exception
UserAPIKeyAuth Fields Reference
The UserAPIKeyAuth object supports the following fields for comprehensive auth configuration:
Core Authentication Fields
UserAPIKeyAuth(
# Basic auth fields
api_key: Optional[str] = None, # The API key (will be hashed automatically)
token: Optional[str] = None, # Hashed token for internal use
key_name: Optional[str] = None, # Human-readable key name
key_alias: Optional[str] = None, # Key alias for identification
# User identification
user_id: Optional[str] = None, # Unique user identifier
user_email: Optional[str] = None, # User email address
user_role: Optional[LitellmUserRoles] = None, # User role (PROXY_ADMIN, INTERNAL_USER, etc.)
# Team/Organization
team_id: Optional[str] = None, # Team identifier
team_alias: Optional[str] = None, # Team display name
org_id: Optional[str] = None, # Organization identifier
)
Budget and Spend Tracking
UserAPIKeyAuth(
# User budgets
max_budget: Optional[float] = None, # Maximum budget for the key
spend: float = 0.0, # Current spend amount
soft_budget: Optional[float] = None, # Soft budget limit (warnings)
model_max_budget: Dict = {}, # Per-model budget limits
model_spend: Dict = {}, # Per-model spend tracking
# Team budgets
team_max_budget: Optional[float] = None, # Team's maximum budget
team_spend: Optional[float] = None, # Team's current spend
team_member_spend: Optional[float] = None, # This user's spend within the team
# Budget timing
budget_duration: Optional[str] = None, # Budget reset period
budget_reset_at: Optional[datetime] = None, # When budget resets
)
Rate Limiting
UserAPIKeyAuth(
# User limits
tpm_limit: Optional[int] = None, # Tokens per minute limit
rpm_limit: Optional[int] = None, # Requests per minute limit
user_tpm_limit: Optional[int] = None, # User-specific TPM limit
user_rpm_limit: Optional[int] = None, # User-specific RPM limit
# Team limits
team_tpm_limit: Optional[int] = None, # Team TPM limit
team_rpm_limit: Optional[int] = None, # Team RPM limit
team_member_tpm_limit: Optional[int] = None, # Per-member TPM limit
team_member_rpm_limit: Optional[int] = None, # Per-member RPM limit
# Per-model limits
rpm_limit_per_model: Optional[Dict[str, int]] = None, # RPM limits by model
tpm_limit_per_model: Optional[Dict[str, int]] = None, # TPM limits by model
)
End User Tracking
UserAPIKeyAuth(
# End user identification and limits
end_user_id: Optional[str] = None, # End user identifier
end_user_tpm_limit: Optional[int] = None, # End user TPM limit
end_user_rpm_limit: Optional[int] = None, # End user RPM limit
end_user_max_budget: Optional[float] = None, # End user budget limit
)
Model and Route Access
UserAPIKeyAuth(
# Model access control
models: List = [], # Allowed models list
team_models: List = [], # Team's allowed models
aliases: Dict = {}, # Model aliases
# Route permissions
allowed_routes: Optional[list] = [], # Allowed API routes
allowed_cache_controls: Optional[list] = [], # Cache control permissions
permissions: Dict = {}, # General permissions
)
Object Permission Example (MCP, agents, etc.)
from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
global_mcp_server_manager,
)
def _server_id(name: str) -> str:
server = global_mcp_server_manager.get_mcp_server_by_name(name)
if not server:
raise ValueError(f"Unknown MCP server '{name}'")
return server.server_id
object_permission = LiteLLM_ObjectPermissionTable(
mcp_servers=[_server_id("deepwiki"), _server_id("everything")], # MCP servers this key is allowed to use
mcp_tool_permissions={"deepwiki": ["search", "read_doc"]}, # optional per-server tool allow-list
)
UserAPIKeyAuth(
object_permission=object_permission,
)
Advanced Configuration
UserAPIKeyAuth(
# Request handling
max_parallel_requests: Optional[int] = None, # Concurrent request limit
allowed_model_region: Optional[AllowedModelRegion] = None, # Geographic restrictions
# Expiration and status
expires: Optional[Union[str, datetime]] = None, # Key expiration
blocked: Optional[bool] = None, # Whether key is blocked
# Metadata and configuration
metadata: Dict = {}, # Custom metadata
config: Dict = {}, # Configuration settings
team_metadata: Optional[Dict] = None, # Team metadata
# Internal tracking
request_route: Optional[str] = None, # Current request route
last_refreshed_at: Optional[float] = None, # Cache refresh timestamp
)
Complete Example
from fastapi import Request
from datetime import datetime, timedelta
from litellm.proxy._types import UserAPIKeyAuth, LitellmUserRoles
async def user_api_key_auth(request: Request, api_key: str) -> UserAPIKeyAuth:
try:
# Example: Comprehensive auth configuration
if api_key.startswith("sk-admin-"):
return UserAPIKeyAuth(
api_key=api_key,
user_id="admin_user_123",
user_email="admin@company.com",
user_role=LitellmUserRoles.PROXY_ADMIN,
team_id="admin_team",
team_alias="Administrative Team",
max_budget=1000.0,
soft_budget=800.0,
tpm_limit=10000,
rpm_limit=100,
models=["gpt-4", "claude-3-sonnet", "gpt-3.5-turbo"],
allowed_routes=["/chat/completions", "/embeddings"],
expires=datetime.now() + timedelta(days=30),
metadata={"department": "engineering", "cost_center": "ai_ops"}
)
elif api_key.startswith("sk-team-"):
return UserAPIKeyAuth(
api_key=api_key,
user_id="team_user_456",
user_email="user@company.com",
user_role=LitellmUserRoles.INTERNAL_USER,
team_id="dev_team",
team_alias="Development Team",
max_budget=100.0,
tpm_limit=1000,
rpm_limit=20,
models=["gpt-3.5-turbo", "claude-3-haiku"],
team_member_tpm_limit=500, # Limit within team
end_user_tpm_limit=100, # Per end-user limit
metadata={"project": "chatbot_v2"}
)
else:
raise Exception("Invalid API key")
except Exception:
raise Exception("Authentication failed")
2. Pass the filepath (relative to the config.yaml)
Pass the filepath to the config.yaml
e.g. if they're both in the same dir - ./config.yaml and ./custom_auth.py, this is what it looks like:
model_list:
- model_name: "openai-model"
litellm_params:
model: "gpt-3.5-turbo"
litellm_settings:
drop_params: True
set_verbose: True
general_settings:
custom_auth: custom_auth.user_api_key_auth
3. Start the proxy
$ litellm --config /path/to/config.yaml
✨ Support LiteLLM Virtual Keys + Custom Auth
Supported from v1.72.2+
:::info
✨ Supporting Custom Auth + LiteLLM Virtual Keys is on LiteLLM Enterprise
Usage
- Setup custom auth file
"""
Example custom auth function.
This will allow all keys starting with "my-custom-key" to pass through.
"""
from typing import Union
from fastapi import Request
from litellm.proxy._types import UserAPIKeyAuth
async def user_api_key_auth(
request: Request, api_key: str
) -> Union[UserAPIKeyAuth, str]:
try:
if api_key.startswith("my-custom-key"):
return "sk-P1zJMdsqCPNN54alZd_ETw"
else:
raise Exception("Invalid API key")
except Exception:
raise Exception("Invalid API key")
- Setup config.yaml
Key change set mode: auto. This will check both litellm api key auth + custom auth.
model_list:
- model_name: "openai-model"
litellm_params:
model: "gpt-3.5-turbo"
api_key: os.environ/OPENAI_API_KEY
general_settings:
custom_auth: custom_auth_auto.user_api_key_auth
custom_auth_settings:
mode: "auto" # can be 'on', 'off', 'auto' - 'auto' checks both litellm api key auth + custom auth
Flow:
-
Checks custom auth first
-
If custom auth fails, checks litellm api key auth
-
If both fail, returns 401
-
Test it!
curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \\
-H 'Content-Type: application/json' \\
-H 'Authorization: Bearer sk-P1zJMdsqCPNN54alZd_ETw' \\
-d '{
"model": "openai-model",
"messages": [
{
"role": "user",
"content": "Hey! My name is John"
}
]
}'
Bubble up custom exceptions
If you want to bubble up custom exceptions, you can do so by raising a ProxyException.
"""
Example custom auth function.
This will allow all keys starting with "my-custom-key" to pass through.
"""
from typing import Union
from fastapi import Request
from litellm.proxy._types import UserAPIKeyAuth, ProxyException
async def user_api_key_auth(
request: Request, api_key: str
) -> Union[UserAPIKeyAuth, str]:
try:
if api_key.startswith("my-custom-key"):
return "sk-P1zJMdsqCPNN54alZd_ETw"
if api_key == "invalid-api-key":
# raise a custom exception back to the client
raise ProxyException(
message="Invalid API key",
type="invalid_request_error",
param="api_key",
code=401,
)
else:
raise Exception("Invalid API key")
except Exception:
raise Exception("Invalid API key")