init policy from config

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",
+]

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

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

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",
+]

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",
+]

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")

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")