inint PolicyValidator

This commit is contained in:
Ishaan Jaffer 2026-01-22 14:52:53 -08:00
parent 6809260a7d
commit 0bd59c5246
5 changed files with 785 additions and 1 deletions

View file

@ -0,0 +1,243 @@
"""
POLICY MANAGEMENT
All /policy management endpoints
/policy/validate - Validate a policy configuration
/policy/list - List all loaded policies
/policy/info - Get information about a specific policy
"""
from typing import Any, Dict, List, Optional
from fastapi import APIRouter, Depends, HTTPException, Request
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
from litellm.types.proxy.policy_engine import (
PolicyMatchContext,
PolicyValidateRequest,
PolicyValidationResponse,
ResolvedPolicy,
)
router = APIRouter()
@router.post(
"/policy/validate",
tags=["policy management"],
dependencies=[Depends(user_api_key_auth)],
response_model=PolicyValidationResponse,
)
@management_endpoint_wrapper
async def validate_policy(
request: Request,
data: PolicyValidateRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> PolicyValidationResponse:
"""
Validate a policy configuration before applying it.
Checks:
- All referenced guardrails exist in the guardrail registry
- All non-wildcard team aliases exist in the database
- All non-wildcard key aliases exist in the database
- Inheritance chains are valid (no cycles, parents exist)
- Scope patterns are syntactically valid
Returns:
- valid: True if the policy configuration is valid (no blocking errors)
- errors: List of blocking validation errors
- warnings: List of non-blocking validation warnings
Example request:
```json
{
"policies": {
"global-baseline": {
"guardrails": {
"add": ["pii_blocker", "phi_blocker"]
},
"scope": {
"teams": ["*"],
"keys": ["*"],
"models": ["*"]
}
},
"healthcare-compliance": {
"inherit": "global-baseline",
"guardrails": {
"add": ["hipaa_audit"]
},
"scope": {
"teams": ["healthcare-team"]
}
}
}
}
```
"""
from litellm.proxy.policy_engine.policy_validator import PolicyValidator
from litellm.proxy.proxy_server import prisma_client
verbose_proxy_logger.debug(
f"Validating policy configuration with {len(data.policies)} policies"
)
validator = PolicyValidator(prisma_client=prisma_client)
result = await validator.validate_policy_config(
data.policies,
validate_db=prisma_client is not None,
)
return result
@router.get(
"/policy/list",
tags=["policy management"],
dependencies=[Depends(user_api_key_auth)],
)
@management_endpoint_wrapper
async def list_policies(
request: Request,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> Dict[str, Any]:
"""
List all loaded policies with their resolved guardrails.
Returns information about each policy including:
- Inheritance configuration
- Scope (teams, keys, models)
- Guardrails to add/remove
- Resolved guardrails (after inheritance)
- Inheritance chain
"""
from litellm.proxy.policy_engine.init_policies import get_policies_summary
return get_policies_summary()
@router.get(
"/policy/info/{policy_name}",
tags=["policy management"],
dependencies=[Depends(user_api_key_auth)],
)
@management_endpoint_wrapper
async def get_policy_info(
request: Request,
policy_name: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> Dict[str, Any]:
"""
Get detailed information about a specific policy.
Returns:
- Policy configuration
- Resolved guardrails (after inheritance)
- Inheritance chain
"""
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
registry = get_policy_registry()
if not registry.is_initialized():
raise HTTPException(
status_code=404,
detail="Policy engine not initialized. No policies loaded.",
)
policy = registry.get_policy(policy_name)
if policy is None:
raise HTTPException(
status_code=404,
detail=f"Policy '{policy_name}' not found",
)
resolved = PolicyResolver.resolve_policy_guardrails(
policy_name, registry.get_all_policies()
)
return {
"policy_name": policy_name,
"inherit": policy.inherit,
"scope": {
"teams": policy.scope.get_teams(),
"keys": policy.scope.get_keys(),
"models": policy.scope.get_models(),
},
"guardrails": {
"add": policy.guardrails.get_add(),
"remove": policy.guardrails.get_remove(),
},
"resolved_guardrails": resolved.guardrails,
"inheritance_chain": resolved.inheritance_chain,
}
@router.post(
"/policy/test",
tags=["policy management"],
dependencies=[Depends(user_api_key_auth)],
)
@management_endpoint_wrapper
async def test_policy_matching(
request: Request,
context: PolicyMatchContext,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> Dict[str, Any]:
"""
Test which policies would match a given request context.
This is useful for debugging and understanding policy behavior.
Request body:
```json
{
"team_alias": "healthcare-team",
"key_alias": "my-api-key",
"model": "gpt-4"
}
```
Returns:
- matching_policies: List of policy names that match
- resolved_guardrails: Final list of guardrails that would be applied
"""
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
registry = get_policy_registry()
if not registry.is_initialized():
return {
"matching_policies": [],
"resolved_guardrails": [],
"message": "Policy engine not initialized. No policies loaded.",
}
policies = registry.get_all_policies()
# Get matching policies
matching_policy_names = PolicyMatcher.get_matching_policies(policies, context)
# Resolve guardrails
resolved_guardrails = PolicyResolver.resolve_guardrails_for_context(
context, policies
)
return {
"context": {
"team_alias": context.team_alias,
"key_alias": context.key_alias,
"model": context.model,
},
"matching_policies": matching_policy_names,
"resolved_guardrails": resolved_guardrails,
}

View file

@ -0,0 +1,377 @@
"""
Policy Validator - Validates policy configurations.
Validates:
- Guardrail names exist in the guardrail registry
- Non-wildcard team aliases exist in the database
- Non-wildcard key aliases exist in the database
- Non-wildcard model names exist in the router or match a wildcard route
- Inheritance chains are valid (no cycles, parents exist)
"""
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set
from litellm._logging import verbose_proxy_logger
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.types.proxy.policy_engine import (
Policy,
PolicyValidationError,
PolicyValidationErrorType,
PolicyValidationResponse,
)
if TYPE_CHECKING:
from litellm.proxy.utils import PrismaClient
from litellm.router import Router
class PolicyValidator:
"""
Validates policy configurations against actual data.
"""
def __init__(
self,
prisma_client: Optional["PrismaClient"] = None,
llm_router: Optional["Router"] = None,
):
"""
Initialize the validator.
Args:
prisma_client: Optional Prisma client for database validation
llm_router: Optional LLM router for model validation
"""
self.prisma_client = prisma_client
self.llm_router = llm_router
@staticmethod
def is_wildcard_pattern(pattern: str) -> bool:
"""
Check if a pattern contains wildcards.
Args:
pattern: The pattern to check
Returns:
True if the pattern contains wildcard characters
"""
return "*" in pattern or "?" in pattern
def get_available_guardrails(self) -> Set[str]:
"""
Get set of available guardrail names from the guardrail registry.
Returns:
Set of guardrail names
"""
try:
from litellm.proxy.guardrails.guardrail_registry import (
IN_MEMORY_GUARDRAIL_HANDLER,
)
guardrails = IN_MEMORY_GUARDRAIL_HANDLER.list_in_memory_guardrails()
return {g.get("guardrail_name", "") for g in guardrails if g.get("guardrail_name")}
except Exception as e:
verbose_proxy_logger.warning(
f"Could not get guardrails from registry: {str(e)}"
)
return set()
async def check_team_alias_exists(self, team_alias: str) -> bool:
"""
Check if a specific team alias exists in the database.
Args:
team_alias: The team alias to check
Returns:
True if the team alias exists
"""
if self.prisma_client is None:
return True # Can't validate without DB, assume valid
try:
team = await self.prisma_client.db.litellm_teamtable.find_first(
where={"team_alias": team_alias},
)
return team is not None
except Exception as e:
verbose_proxy_logger.warning(
f"Could not check team alias '{team_alias}': {str(e)}"
)
return True # Assume valid on error
async def check_key_alias_exists(self, key_alias: str) -> bool:
"""
Check if a specific key alias exists in the database.
Args:
key_alias: The key alias to check
Returns:
True if the key alias exists
"""
if self.prisma_client is None:
return True # Can't validate without DB, assume valid
try:
key = await self.prisma_client.db.litellm_verificationtoken.find_first(
where={"key_alias": key_alias},
)
return key is not None
except Exception as e:
verbose_proxy_logger.warning(
f"Could not check key alias '{key_alias}': {str(e)}"
)
return True # Assume valid on error
def check_model_exists(self, model: str) -> bool:
"""
Check if a model exists in the router or matches a wildcard pattern.
Args:
model: The model name to check
Returns:
True if the model exists or matches a pattern in the router
"""
if self.llm_router is None:
return True # Can't validate without router, assume valid
try:
# Check if model is in router's model names
if model in self.llm_router.model_names:
return True
# Check if model matches any pattern via pattern router
if hasattr(self.llm_router, "pattern_router"):
pattern_deployments = self.llm_router.pattern_router.get_deployments_by_pattern(
model=model
)
if pattern_deployments:
return True
return False
except Exception as e:
verbose_proxy_logger.warning(
f"Could not check model '{model}': {str(e)}"
)
return True # Assume valid on error
def _validate_inheritance_chain(
self,
policy_name: str,
policies: Dict[str, Policy],
visited: Optional[Set[str]] = None,
) -> List[PolicyValidationError]:
"""
Validate the inheritance chain for a policy.
Checks for:
- Parent policy exists
- No circular inheritance
Args:
policy_name: Name of the policy to validate
policies: All policies
visited: Set of already visited policy names (for cycle detection)
Returns:
List of validation errors
"""
errors: List[PolicyValidationError] = []
if visited is None:
visited = set()
if policy_name in visited:
errors.append(
PolicyValidationError(
policy_name=policy_name,
error_type=PolicyValidationErrorType.CIRCULAR_INHERITANCE,
message=f"Circular inheritance detected: {' -> '.join(visited)} -> {policy_name}",
field="inherit",
)
)
return errors
policy = policies.get(policy_name)
if policy is None:
return errors
if policy.inherit:
if policy.inherit not in policies:
errors.append(
PolicyValidationError(
policy_name=policy_name,
error_type=PolicyValidationErrorType.INVALID_INHERITANCE,
message=f"Parent policy '{policy.inherit}' not found",
field="inherit",
value=policy.inherit,
)
)
else:
# Recursively check parent
visited.add(policy_name)
errors.extend(
self._validate_inheritance_chain(policy.inherit, policies, visited)
)
return errors
async def validate_policies(
self,
policies: Dict[str, Policy],
validate_db: bool = True,
) -> PolicyValidationResponse:
"""
Validate a set of policies.
Args:
policies: Dictionary mapping policy names to Policy objects
validate_db: Whether to validate against database (teams, keys)
Returns:
PolicyValidationResponse with errors and warnings
"""
errors: List[PolicyValidationError] = []
warnings: List[PolicyValidationError] = []
# Get available guardrails
available_guardrails = self.get_available_guardrails()
for policy_name, policy in policies.items():
# Validate guardrails
for guardrail in policy.guardrails.get_add():
if available_guardrails and guardrail not in available_guardrails:
errors.append(
PolicyValidationError(
policy_name=policy_name,
error_type=PolicyValidationErrorType.INVALID_GUARDRAIL,
message=f"Guardrail '{guardrail}' not found in guardrail registry",
field="guardrails.add",
value=guardrail,
)
)
for guardrail in policy.guardrails.get_remove():
if available_guardrails and guardrail not in available_guardrails:
warnings.append(
PolicyValidationError(
policy_name=policy_name,
error_type=PolicyValidationErrorType.INVALID_GUARDRAIL,
message=f"Guardrail '{guardrail}' in remove list not found in guardrail registry",
field="guardrails.remove",
value=guardrail,
)
)
# Validate team aliases (non-wildcard only, query per alias)
if validate_db and self.prisma_client:
for team_pattern in policy.scope.get_teams():
if not self.is_wildcard_pattern(team_pattern):
exists = await self.check_team_alias_exists(team_alias=team_pattern)
if not exists:
warnings.append(
PolicyValidationError(
policy_name=policy_name,
error_type=PolicyValidationErrorType.INVALID_TEAM,
message=f"Team alias '{team_pattern}' not found in database",
field="scope.teams",
value=team_pattern,
)
)
# Validate key aliases (non-wildcard only, query per alias)
if validate_db and self.prisma_client:
for key_pattern in policy.scope.get_keys():
if not self.is_wildcard_pattern(key_pattern):
exists = await self.check_key_alias_exists(key_alias=key_pattern)
if not exists:
warnings.append(
PolicyValidationError(
policy_name=policy_name,
error_type=PolicyValidationErrorType.INVALID_KEY,
message=f"Key alias '{key_pattern}' not found in database",
field="scope.keys",
value=key_pattern,
)
)
# Validate models (non-wildcard only, check against router)
if self.llm_router:
for model_pattern in policy.scope.get_models():
if not self.is_wildcard_pattern(model_pattern):
exists = self.check_model_exists(model=model_pattern)
if not exists:
warnings.append(
PolicyValidationError(
policy_name=policy_name,
error_type=PolicyValidationErrorType.INVALID_MODEL,
message=f"Model '{model_pattern}' not found in router",
field="scope.models",
value=model_pattern,
)
)
# Validate inheritance
inheritance_errors = self._validate_inheritance_chain(
policy_name=policy_name, policies=policies
)
errors.extend(inheritance_errors)
return PolicyValidationResponse(
valid=len(errors) == 0,
errors=errors,
warnings=warnings,
)
async def validate_policy_config(
self,
policy_config: Dict[str, Any],
validate_db: bool = True,
) -> PolicyValidationResponse:
"""
Validate a raw policy configuration dictionary.
This parses the config and then validates it.
Args:
policy_config: Raw policy configuration from YAML
validate_db: Whether to validate against database
Returns:
PolicyValidationResponse with errors and warnings
"""
from litellm.proxy.policy_engine.policy_registry import PolicyRegistry
# First, try to parse the policies
errors: List[PolicyValidationError] = []
policies: Dict[str, Policy] = {}
temp_registry = PolicyRegistry()
for policy_name, policy_data in policy_config.items():
try:
policy = temp_registry._parse_policy(policy_name, policy_data)
policies[policy_name] = policy
except Exception as e:
errors.append(
PolicyValidationError(
policy_name=policy_name,
error_type=PolicyValidationErrorType.INVALID_SYNTAX,
message=f"Failed to parse policy: {str(e)}",
)
)
# If there were parsing errors, return early
if errors:
return PolicyValidationResponse(
valid=False,
errors=errors,
warnings=[],
)
# Validate the parsed policies
return await self.validate_policies(policies, validate_db=validate_db)

View file

@ -13,7 +13,13 @@ from litellm.types.proxy.policy_engine.policy_types import (
PolicyScope,
)
from litellm.types.proxy.policy_engine.resolver_types import (
PolicyGuardrailsResponse,
PolicyInfoResponse,
PolicyListResponse,
PolicyMatchContext,
PolicyScopeResponse,
PolicySummaryItem,
PolicyTestResponse,
ResolvedPolicy,
)
from litellm.types.proxy.policy_engine.validation_types import (
@ -37,4 +43,11 @@ __all__ = [
# Resolver types
"PolicyMatchContext",
"ResolvedPolicy",
# API Response types
"PolicyGuardrailsResponse",
"PolicyInfoResponse",
"PolicyListResponse",
"PolicyScopeResponse",
"PolicySummaryItem",
"PolicyTestResponse",
]

View file

@ -5,7 +5,7 @@ These types are used for matching requests to policies and resolving
the final guardrails list.
"""
from typing import List, Optional
from typing import Dict, List, Optional
from pydantic import BaseModel, ConfigDict, Field
@ -51,3 +51,60 @@ class ResolvedPolicy(BaseModel):
)
model_config = ConfigDict(extra="forbid")
# ─────────────────────────────────────────────────────────────────────────────
# API Response Types
# ─────────────────────────────────────────────────────────────────────────────
class PolicyScopeResponse(BaseModel):
"""Scope configuration for a policy."""
teams: List[str] = Field(default_factory=list)
keys: List[str] = Field(default_factory=list)
models: List[str] = Field(default_factory=list)
class PolicyGuardrailsResponse(BaseModel):
"""Guardrails configuration for a policy."""
add: List[str] = Field(default_factory=list)
remove: List[str] = Field(default_factory=list)
class PolicyInfoResponse(BaseModel):
"""Response for /policy/info/{policy_name} endpoint."""
policy_name: str
inherit: Optional[str] = None
scope: PolicyScopeResponse
guardrails: PolicyGuardrailsResponse
resolved_guardrails: List[str]
inheritance_chain: List[str]
class PolicySummaryItem(BaseModel):
"""Summary of a single policy for list endpoint."""
inherit: Optional[str] = None
scope: PolicyScopeResponse
guardrails: PolicyGuardrailsResponse
resolved_guardrails: List[str]
inheritance_chain: List[str]
class PolicyListResponse(BaseModel):
"""Response for /policy/list endpoint."""
policies: Dict[str, PolicySummaryItem]
total_count: int
class PolicyTestResponse(BaseModel):
"""Response for /policy/test endpoint."""
context: PolicyMatchContext
matching_policies: List[str]
resolved_guardrails: List[str]
message: Optional[str] = None

View file

@ -0,0 +1,94 @@
"""
Unit tests for PolicyValidator - tests policy configuration validation.
Tests validation of:
- Inheritance chains (parent exists, no circular deps)
- Guardrail names exist in registry
- Model names exist in router
"""
from unittest.mock import MagicMock, patch
import pytest
from litellm.proxy.policy_engine.policy_validator import PolicyValidator
from litellm.types.proxy.policy_engine import (
Policy,
PolicyGuardrails,
PolicyScope,
PolicyValidationErrorType,
)
class TestPolicyValidator:
"""Test policy validation logic."""
@pytest.mark.asyncio
async def test_validate_missing_parent_policy(self):
"""Test that referencing non-existent parent policy fails."""
policies = {
"child": Policy(
inherit="nonexistent-parent",
guardrails=PolicyGuardrails(add=["hipaa_audit"]),
scope=PolicyScope(teams=["healthcare-team"]),
),
}
validator = PolicyValidator(prisma_client=None)
result = await validator.validate_policies(policies=policies, validate_db=False)
assert result.valid is False
assert any(
e.error_type == PolicyValidationErrorType.INVALID_INHERITANCE
for e in result.errors
)
@pytest.mark.asyncio
async def test_validate_invalid_guardrail(self):
"""Test that referencing non-existent guardrail fails."""
policies = {
"test-policy": Policy(
guardrails=PolicyGuardrails(add=["nonexistent_guardrail"]),
scope=PolicyScope(teams=["*"]),
),
}
validator = PolicyValidator(prisma_client=None)
with patch.object(
validator, "get_available_guardrails", return_value={"pii_blocker", "toxicity_filter"}
):
result = await validator.validate_policies(policies=policies, validate_db=False)
assert result.valid is False
assert any(
e.error_type == PolicyValidationErrorType.INVALID_GUARDRAIL
and e.value == "nonexistent_guardrail"
for e in result.errors
)
@pytest.mark.asyncio
async def test_validate_invalid_model(self):
"""Test that referencing non-existent model warns."""
policies = {
"test-policy": Policy(
guardrails=PolicyGuardrails(add=["pii_blocker"]),
scope=PolicyScope(models=["nonexistent-model"]),
),
}
# Mock the router with known model names
mock_router = MagicMock()
mock_router.model_names = {"gpt-4", "gpt-3.5-turbo"}
# Mock pattern_router to return empty list (no pattern matches)
mock_router.pattern_router.get_deployments_by_pattern.return_value = []
validator = PolicyValidator(prisma_client=None, llm_router=mock_router)
with patch.object(validator, "get_available_guardrails", return_value={"pii_blocker"}):
result = await validator.validate_policies(policies=policies, validate_db=False)
# Model validation is a warning, not an error
assert any(
w.error_type == PolicyValidationErrorType.INVALID_MODEL
and w.value == "nonexistent-model"
for w in result.warnings
)