From 6809260a7d1122135ece60ec229ce0fe59b7b5a7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 22 Jan 2026 14:43:40 -0800 Subject: [PATCH] init policy from config --- litellm/proxy/policy_engine/__init__.py | 23 +++ litellm/proxy/policy_engine/init_policies.py | 163 +++++++++++++++ .../proxy/policy_engine/policy_registry.py | 190 ++++++++++++++++++ litellm/types/policy_engine.py | 36 ++++ litellm/types/proxy/policy_engine/__init__.py | 40 ++++ .../types/proxy/policy_engine/policy_types.py | 154 ++++++++++++++ .../proxy/policy_engine/validation_types.py | 80 ++++++++ .../proxy/policy_engine/__init__.py | 0 8 files changed, 686 insertions(+) create mode 100644 litellm/proxy/policy_engine/__init__.py create mode 100644 litellm/proxy/policy_engine/init_policies.py create mode 100644 litellm/proxy/policy_engine/policy_registry.py create mode 100644 litellm/types/policy_engine.py create mode 100644 litellm/types/proxy/policy_engine/__init__.py create mode 100644 litellm/types/proxy/policy_engine/policy_types.py create mode 100644 litellm/types/proxy/policy_engine/validation_types.py create mode 100644 tests/test_litellm/proxy/policy_engine/__init__.py diff --git a/litellm/proxy/policy_engine/__init__.py b/litellm/proxy/policy_engine/__init__.py new file mode 100644 index 00000000000..2ee7bdf1b17 --- /dev/null +++ b/litellm/proxy/policy_engine/__init__.py @@ -0,0 +1,23 @@ +""" +LiteLLM Policy Engine + +The Policy Engine allows administrators to define policies that combine guardrails +with scoping rules. Policies can target specific teams, API keys, and models using +wildcard patterns, and support inheritance from base policies. +""" + +from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher +from litellm.proxy.policy_engine.policy_registry import ( + PolicyRegistry, + get_policy_registry, +) +from litellm.proxy.policy_engine.policy_resolver import PolicyResolver +from litellm.proxy.policy_engine.policy_validator import PolicyValidator + +__all__ = [ + "PolicyRegistry", + "get_policy_registry", + "PolicyMatcher", + "PolicyResolver", + "PolicyValidator", +] diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py new file mode 100644 index 00000000000..520bfe40747 --- /dev/null +++ b/litellm/proxy/policy_engine/init_policies.py @@ -0,0 +1,163 @@ +""" +Policy Initialization - Loads policies from config and validates on startup. +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.proxy.policy_engine.policy_registry import get_policy_registry +from litellm.proxy.policy_engine.policy_validator import PolicyValidator +from litellm.types.proxy.policy_engine import PolicyValidationResponse + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +async def init_policies( + policies_config: Dict[str, Any], + prisma_client: Optional["PrismaClient"] = None, + validate_db: bool = True, + fail_on_error: bool = True, +) -> PolicyValidationResponse: + """ + Initialize policies from configuration. + + This function: + 1. Parses the policy configuration + 2. Validates policies (guardrails exist, teams/keys exist in DB) + 3. Loads policies into the global registry + + Args: + policies_config: Dictionary mapping policy names to policy definitions + prisma_client: Optional Prisma client for database validation + validate_db: Whether to validate team/key aliases against database + fail_on_error: If True, raise exception on validation errors + + Returns: + PolicyValidationResponse with validation results + + Raises: + ValueError: If fail_on_error is True and validation errors are found + """ + verbose_proxy_logger.info(f"Initializing {len(policies_config)} policies...") + + # Get the global registry + registry = get_policy_registry() + + # Create validator + validator = PolicyValidator(prisma_client=prisma_client) + + # Validate the configuration + validation_result = await validator.validate_policy_config( + policies_config, + validate_db=validate_db, + ) + + # Log validation results + if validation_result.errors: + for error in validation_result.errors: + verbose_proxy_logger.error( + f"Policy validation error in '{error.policy_name}': " + f"[{error.error_type}] {error.message}" + ) + + if validation_result.warnings: + for warning in validation_result.warnings: + verbose_proxy_logger.warning( + f"Policy validation warning in '{warning.policy_name}': " + f"[{warning.error_type}] {warning.message}" + ) + + # Fail if there are errors and fail_on_error is True + if not validation_result.valid and fail_on_error: + error_messages = [ + f"[{e.policy_name}] {e.message}" for e in validation_result.errors + ] + raise ValueError( + f"Policy validation failed with {len(validation_result.errors)} error(s):\n" + + "\n".join(error_messages) + ) + + # Load policies into registry (even with warnings) + try: + registry.load_policies(policies_config) + verbose_proxy_logger.info( + f"Successfully loaded {len(policies_config)} policies" + ) + except Exception as e: + verbose_proxy_logger.error(f"Failed to load policies: {str(e)}") + raise + + return validation_result + + +def init_policies_sync( + policies_config: Dict[str, Any], + fail_on_error: bool = True, +) -> None: + """ + Synchronous version of init_policies (without DB validation). + + Use this when async is not available or DB validation is not needed. + + Args: + policies_config: Dictionary mapping policy names to policy definitions + fail_on_error: If True, raise exception on validation errors + """ + import asyncio + + # Run the async function without DB validation + try: + loop = asyncio.get_event_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + loop.run_until_complete( + init_policies( + policies_config=policies_config, + prisma_client=None, + validate_db=False, + fail_on_error=fail_on_error, + ) + ) + + +def get_policies_summary() -> Dict[str, Any]: + """ + Get a summary of loaded policies for debugging/display. + + Returns: + Dictionary with policy information + """ + from litellm.proxy.policy_engine.policy_resolver import PolicyResolver + + registry = get_policy_registry() + + if not registry.is_initialized(): + return {"initialized": False, "policies": {}} + + resolved = PolicyResolver.get_all_resolved_policies() + + summary = { + "initialized": True, + "policy_count": len(resolved), + "policies": {}, + } + + for policy_name, resolved_policy in resolved.items(): + policy = registry.get_policy(policy_name) + summary["policies"][policy_name] = { + "inherit": policy.inherit if policy else None, + "scope": { + "teams": policy.scope.get_teams() if policy else [], + "keys": policy.scope.get_keys() if policy else [], + "models": policy.scope.get_models() if policy else [], + }, + "guardrails_add": policy.guardrails.get_add() if policy else [], + "guardrails_remove": policy.guardrails.get_remove() if policy else [], + "resolved_guardrails": resolved_policy.guardrails, + "inheritance_chain": resolved_policy.inheritance_chain, + } + + return summary diff --git a/litellm/proxy/policy_engine/policy_registry.py b/litellm/proxy/policy_engine/policy_registry.py new file mode 100644 index 00000000000..b090e3579b2 --- /dev/null +++ b/litellm/proxy/policy_engine/policy_registry.py @@ -0,0 +1,190 @@ +""" +Policy Registry - In-memory storage for policies. + +Handles storing, retrieving, and managing policies. +""" + +from typing import Any, Dict, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.types.proxy.policy_engine import ( + Policy, + PolicyConfig, + PolicyGuardrails, + PolicyScope, +) + + +class PolicyRegistry: + """ + In-memory registry for storing and managing policies. + + This is a singleton that holds all loaded policies and provides + methods to access them. + """ + + def __init__(self): + self._policies: Dict[str, Policy] = {} + self._initialized: bool = False + + def load_policies(self, policies_config: Dict[str, Any]) -> None: + """ + Load policies from a configuration dictionary. + + Args: + policies_config: Dictionary mapping policy names to policy definitions. + This is the raw config from the YAML file. + """ + self._policies = {} + + for policy_name, policy_data in policies_config.items(): + try: + policy = self._parse_policy(policy_name, policy_data) + self._policies[policy_name] = policy + verbose_proxy_logger.debug(f"Loaded policy: {policy_name}") + except Exception as e: + verbose_proxy_logger.error( + f"Error loading policy '{policy_name}': {str(e)}" + ) + raise ValueError(f"Invalid policy '{policy_name}': {str(e)}") from e + + self._initialized = True + verbose_proxy_logger.info(f"Loaded {len(self._policies)} policies") + + def _parse_policy(self, policy_name: str, policy_data: Dict[str, Any]) -> Policy: + """ + Parse a policy from raw configuration data. + + Args: + policy_name: Name of the policy + policy_data: Raw policy configuration + + Returns: + Parsed Policy object + """ + # Parse guardrails + guardrails_data = policy_data.get("guardrails", {}) + if isinstance(guardrails_data, dict): + guardrails = PolicyGuardrails( + add=guardrails_data.get("add"), + remove=guardrails_data.get("remove"), + ) + else: + # Handle legacy format where guardrails might be a list + guardrails = PolicyGuardrails(add=guardrails_data if guardrails_data else None) + + # Parse scope + scope_data = policy_data.get("scope", {}) + scope = PolicyScope( + teams=scope_data.get("teams"), + keys=scope_data.get("keys"), + models=scope_data.get("models"), + ) + + return Policy( + inherit=policy_data.get("inherit"), + guardrails=guardrails, + scope=scope, + ) + + def get_policy(self, policy_name: str) -> Optional[Policy]: + """ + Get a policy by name. + + Args: + policy_name: Name of the policy to retrieve + + Returns: + Policy object if found, None otherwise + """ + return self._policies.get(policy_name) + + def get_all_policies(self) -> Dict[str, Policy]: + """ + Get all loaded policies. + + Returns: + Dictionary mapping policy names to Policy objects + """ + return self._policies.copy() + + def get_policy_names(self) -> List[str]: + """ + Get list of all policy names. + + Returns: + List of policy names + """ + return list(self._policies.keys()) + + def has_policy(self, policy_name: str) -> bool: + """ + Check if a policy exists. + + Args: + policy_name: Name of the policy to check + + Returns: + True if policy exists, False otherwise + """ + return policy_name in self._policies + + def is_initialized(self) -> bool: + """ + Check if the registry has been initialized with policies. + + Returns: + True if policies have been loaded, False otherwise + """ + return self._initialized + + def clear(self) -> None: + """ + Clear all policies from the registry. + """ + self._policies = {} + self._initialized = False + + def add_policy(self, policy_name: str, policy: Policy) -> None: + """ + Add or update a single policy. + + Args: + policy_name: Name of the policy + policy: Policy object to add + """ + self._policies[policy_name] = policy + verbose_proxy_logger.debug(f"Added/updated policy: {policy_name}") + + def remove_policy(self, policy_name: str) -> bool: + """ + Remove a policy by name. + + Args: + policy_name: Name of the policy to remove + + Returns: + True if policy was removed, False if it didn't exist + """ + if policy_name in self._policies: + del self._policies[policy_name] + verbose_proxy_logger.debug(f"Removed policy: {policy_name}") + return True + return False + + +# Global singleton instance +_policy_registry: Optional[PolicyRegistry] = None + + +def get_policy_registry() -> PolicyRegistry: + """ + Get the global PolicyRegistry singleton. + + Returns: + The global PolicyRegistry instance + """ + global _policy_registry + if _policy_registry is None: + _policy_registry = PolicyRegistry() + return _policy_registry diff --git a/litellm/types/policy_engine.py b/litellm/types/policy_engine.py new file mode 100644 index 00000000000..d5eb7e2b140 --- /dev/null +++ b/litellm/types/policy_engine.py @@ -0,0 +1,36 @@ +""" +Type definitions for the LiteLLM Policy Engine. + +This module re-exports types from litellm.types.proxy.policy_engine for backward compatibility. +The canonical location for these types is litellm/types/proxy/policy_engine/. +""" + +# Re-export all types from the new location +from litellm.types.proxy.policy_engine import ( # Policy types; Validation types; Resolver types + Policy, + PolicyConfig, + PolicyGuardrails, + PolicyMatchContext, + PolicyScope, + PolicyValidateRequest, + PolicyValidationError, + PolicyValidationErrorType, + PolicyValidationResponse, + ResolvedPolicy, +) + +__all__ = [ + # Policy types + "Policy", + "PolicyConfig", + "PolicyGuardrails", + "PolicyScope", + # Validation types + "PolicyValidateRequest", + "PolicyValidationError", + "PolicyValidationErrorType", + "PolicyValidationResponse", + # Resolver types + "PolicyMatchContext", + "ResolvedPolicy", +] diff --git a/litellm/types/proxy/policy_engine/__init__.py b/litellm/types/proxy/policy_engine/__init__.py new file mode 100644 index 00000000000..c7eb2f42e34 --- /dev/null +++ b/litellm/types/proxy/policy_engine/__init__.py @@ -0,0 +1,40 @@ +""" +Type definitions for the LiteLLM Policy Engine. + +The Policy Engine allows administrators to define policies that combine guardrails +with scoping rules. Policies can target specific teams, API keys, and models using +wildcard patterns, and support inheritance from base policies. +""" + +from litellm.types.proxy.policy_engine.policy_types import ( + Policy, + PolicyConfig, + PolicyGuardrails, + PolicyScope, +) +from litellm.types.proxy.policy_engine.resolver_types import ( + PolicyMatchContext, + ResolvedPolicy, +) +from litellm.types.proxy.policy_engine.validation_types import ( + PolicyValidateRequest, + PolicyValidationError, + PolicyValidationErrorType, + PolicyValidationResponse, +) + +__all__ = [ + # Policy types + "Policy", + "PolicyConfig", + "PolicyGuardrails", + "PolicyScope", + # Validation types + "PolicyValidateRequest", + "PolicyValidationError", + "PolicyValidationErrorType", + "PolicyValidationResponse", + # Resolver types + "PolicyMatchContext", + "ResolvedPolicy", +] diff --git a/litellm/types/proxy/policy_engine/policy_types.py b/litellm/types/proxy/policy_engine/policy_types.py new file mode 100644 index 00000000000..7b29565fde5 --- /dev/null +++ b/litellm/types/proxy/policy_engine/policy_types.py @@ -0,0 +1,154 @@ +""" +Core policy type definitions. + +These types define the structure of policies in the configuration. +""" + +from typing import Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field + + +class PolicyScope(BaseModel): + """ + Defines the scope for a policy - which requests it applies to. + + Scope Fields: + | Field | What it matches | Wildcard support | + |--------|-----------------|----------------------| + | teams | Team aliases | *, healthcare-* | + | keys | Key aliases | *, dev-key-* | + | models | Model names | *, bedrock/*, gpt-* | + + If a field is None or empty, it defaults to matching everything (["*"]). + A request must match ALL specified scope fields for the policy to apply. + """ + + teams: Optional[List[str]] = Field( + default=None, + description="Team aliases or wildcard patterns. Use '*' for all teams.", + ) + keys: Optional[List[str]] = Field( + default=None, + description="Key aliases or wildcard patterns. Use '*' for all keys.", + ) + models: Optional[List[str]] = Field( + default=None, + description="Model names or wildcard patterns. Use '*' for all models.", + ) + + model_config = ConfigDict(extra="forbid") + + def get_teams(self) -> List[str]: + """Returns teams list, defaulting to ['*'] if not specified.""" + return self.teams if self.teams else ["*"] + + def get_keys(self) -> List[str]: + """Returns keys list, defaulting to ['*'] if not specified.""" + return self.keys if self.keys else ["*"] + + def get_models(self) -> List[str]: + """Returns models list, defaulting to ['*'] if not specified.""" + return self.models if self.models else ["*"] + + +class PolicyGuardrails(BaseModel): + """ + Defines guardrails to add or remove in a policy. + + - `add`: List of guardrail names to add (on top of inherited guardrails) + - `remove`: List of guardrail names to remove (from inherited guardrails) + """ + + add: Optional[List[str]] = Field( + default=None, + description="Guardrail names to add to this policy.", + ) + remove: Optional[List[str]] = Field( + default=None, + description="Guardrail names to remove (typically from inherited policy).", + ) + + model_config = ConfigDict(extra="forbid") + + def get_add(self) -> List[str]: + """Returns add list, defaulting to empty list if not specified.""" + return self.add if self.add else [] + + def get_remove(self) -> List[str]: + """Returns remove list, defaulting to empty list if not specified.""" + return self.remove if self.remove else [] + + +class Policy(BaseModel): + """ + A policy that defines which guardrails apply to requests matching its scope. + + Policies can inherit from other policies using the `inherit` field. + When inheriting: + - Guardrails from `guardrails.add` are added to the inherited guardrails + - Guardrails from `guardrails.remove` are removed from the inherited guardrails + + Example configuration: + ```yaml + 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, medical-research] + models: [gpt-4, bedrock/claude-*] + + internal-dev: + inherit: global-baseline + guardrails: + add: + - toxicity_filter + remove: + - phi_blocker + scope: + keys: [dev-key-*, test-key-*] + ``` + """ + + inherit: Optional[str] = Field( + default=None, + description="Name of the parent policy to inherit from.", + ) + guardrails: PolicyGuardrails = Field( + default_factory=PolicyGuardrails, + description="Guardrails configuration with add/remove lists.", + ) + scope: PolicyScope = Field( + default_factory=PolicyScope, + description="Scope defining which requests this policy applies to.", + ) + + model_config = ConfigDict(extra="forbid") + + +class PolicyConfig(BaseModel): + """ + Root configuration for all policies. + + Maps policy names to their Policy definitions. + """ + + policies: Dict[str, Policy] = Field( + default_factory=dict, + description="Map of policy names to Policy objects.", + ) + + model_config = ConfigDict(extra="forbid") diff --git a/litellm/types/proxy/policy_engine/validation_types.py b/litellm/types/proxy/policy_engine/validation_types.py new file mode 100644 index 00000000000..e079febcc9e --- /dev/null +++ b/litellm/types/proxy/policy_engine/validation_types.py @@ -0,0 +1,80 @@ +""" +Policy validation type definitions. + +These types are used for validating policy configurations and returning +validation results. +""" + +from enum import Enum +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field + + +class PolicyValidationErrorType(str, Enum): + """Types of validation errors that can occur.""" + + INVALID_GUARDRAIL = "invalid_guardrail" + INVALID_TEAM = "invalid_team" + INVALID_KEY = "invalid_key" + INVALID_MODEL = "invalid_model" + INVALID_INHERITANCE = "invalid_inheritance" + CIRCULAR_INHERITANCE = "circular_inheritance" + INVALID_SCOPE = "invalid_scope" + INVALID_SYNTAX = "invalid_syntax" + + +class PolicyValidationError(BaseModel): + """ + Represents a validation error or warning for a policy. + """ + + policy_name: str = Field(description="Name of the policy with the issue.") + error_type: PolicyValidationErrorType = Field( + description="Type of validation error." + ) + message: str = Field(description="Human-readable error message.") + field: Optional[str] = Field( + default=None, + description="Specific field that caused the error (e.g., 'guardrails.add', 'scope.teams').", + ) + value: Optional[str] = Field( + default=None, + description="The invalid value that caused the error.", + ) + + model_config = ConfigDict(extra="forbid") + + +class PolicyValidationResponse(BaseModel): + """ + Response from policy validation. + + - `valid`: True if no blocking errors were found + - `errors`: List of blocking errors (prevent policy from being applied) + - `warnings`: List of non-blocking warnings (policy can still be applied) + """ + + valid: bool = Field(description="True if the policy configuration is valid.") + errors: List[PolicyValidationError] = Field( + default_factory=list, + description="List of blocking validation errors.", + ) + warnings: List[PolicyValidationError] = Field( + default_factory=list, + description="List of non-blocking validation warnings.", + ) + + model_config = ConfigDict(extra="forbid") + + +class PolicyValidateRequest(BaseModel): + """ + Request body for the /policy/validate endpoint. + """ + + policies: Dict[str, Any] = Field( + description="Policy configuration to validate. Map of policy names to policy definitions." + ) + + model_config = ConfigDict(extra="forbid") diff --git a/tests/test_litellm/proxy/policy_engine/__init__.py b/tests/test_litellm/proxy/policy_engine/__init__.py new file mode 100644 index 00000000000..e69de29bb2d