diff --git a/litellm/proxy/management_endpoints/policy_endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints.py new file mode 100644 index 00000000000..84048a2105d --- /dev/null +++ b/litellm/proxy/management_endpoints/policy_endpoints.py @@ -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, + } diff --git a/litellm/proxy/policy_engine/policy_validator.py b/litellm/proxy/policy_engine/policy_validator.py new file mode 100644 index 00000000000..01c6e0735aa --- /dev/null +++ b/litellm/proxy/policy_engine/policy_validator.py @@ -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) diff --git a/litellm/types/proxy/policy_engine/__init__.py b/litellm/types/proxy/policy_engine/__init__.py index c7eb2f42e34..bb622e846d1 100644 --- a/litellm/types/proxy/policy_engine/__init__.py +++ b/litellm/types/proxy/policy_engine/__init__.py @@ -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", ] diff --git a/litellm/types/proxy/policy_engine/resolver_types.py b/litellm/types/proxy/policy_engine/resolver_types.py index fa513e123f2..81ae248d436 100644 --- a/litellm/types/proxy/policy_engine/resolver_types.py +++ b/litellm/types/proxy/policy_engine/resolver_types.py @@ -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 diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_validator.py b/tests/test_litellm/proxy/policy_engine/test_policy_validator.py new file mode 100644 index 00000000000..51f54e33c87 --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_policy_validator.py @@ -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 + )