mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
inint PolicyValidator
This commit is contained in:
parent
6809260a7d
commit
0bd59c5246
5 changed files with 785 additions and 1 deletions
243
litellm/proxy/management_endpoints/policy_endpoints.py
Normal file
243
litellm/proxy/management_endpoints/policy_endpoints.py
Normal 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,
|
||||
}
|
||||
377
litellm/proxy/policy_engine/policy_validator.py
Normal file
377
litellm/proxy/policy_engine/policy_validator.py
Normal 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)
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
Loading…
Add table
Reference in a new issue