[Feat] Guardrail Policy Management - Allow using UI to manage guardrail policies

.gitignore

@@ -1,5 +1,6 @@
 .python-version
 .venv
+.venv_policy_test
 .env
 .newenv
 newenv/*

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -124,7 +124,7 @@ model LiteLLM_TeamTable {
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
     model_spend      Json @default("{}")
     model_max_budget Json @default("{}")
-    router_settings Json? @default("{}")
+    router_settings  Json? @default("{}")
     team_member_permissions String[] @default([])
     model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases
     litellm_organization_table LiteLLM_OrganizationTable?   @relation(fields: [organization_id], references: [organization_id])

litellm/proxy/_types.py

@@ -73,6 +73,7 @@ class SupportedDBObjectType(str, enum.Enum):
     MODELS = "models"
     MCP = "mcp"
     GUARDRAILS = "guardrails"
+    POLICIES = "policies"
     VECTOR_STORES = "vector_stores"
     PASS_THROUGH_ENDPOINTS = "pass_through_endpoints"
     PROMPTS = "prompts"
@@ -844,6 +845,7 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase):
     model_rpm_limit: Optional[dict] = None
     model_tpm_limit: Optional[dict] = None
     guardrails: Optional[List[str]] = None
+    policies: Optional[List[str]] = None
     prompts: Optional[List[str]] = None
     blocked: Optional[bool] = None
     aliases: Optional[dict] = {}
@@ -1477,6 +1479,7 @@ class NewTeamRequest(TeamBase):
     model_aliases: Optional[dict] = None
     tags: Optional[list] = None
     guardrails: Optional[List[str]] = None
+    policies: Optional[List[str]] = None
     prompts: Optional[List[str]] = None
     object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
     allowed_passthrough_routes: Optional[list] = None
@@ -1526,6 +1529,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
     blocked: Optional[bool] = None
     budget_duration: Optional[str] = None
     guardrails: Optional[List[str]] = None
+    policies: Optional[List[str]] = None
     """
 
     team_id: str  # required
@@ -1541,6 +1545,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
     tags: Optional[list] = None
     model_aliases: Optional[dict] = None
     guardrails: Optional[List[str]] = None
+    policies: Optional[List[str]] = None
     object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
     team_member_budget: Optional[float] = None
     team_member_budget_duration: Optional[str] = None
@@ -3499,6 +3504,7 @@ class PassThroughEndpointLoggingTypedDict(TypedDict):
 
 LiteLLM_ManagementEndpoint_MetadataFields_Premium = [
     "guardrails",
+    "policies",
     "tags",
     "team_member_key_duration",
     "prompts",

litellm/proxy/litellm_pre_call_utils.py

@@ -1311,6 +1311,118 @@ def _add_guardrails_from_key_or_team_metadata(
         data[metadata_variable_name]["guardrails"] = list(combined_guardrails)
 
 
+def _add_guardrails_from_policies_in_metadata(
+    key_metadata: Optional[dict],
+    team_metadata: Optional[dict],
+    data: dict,
+    metadata_variable_name: str,
+) -> None:
+    """
+    Helper to resolve guardrails from policies attached to key/team metadata.
+
+    This function:
+    1. Gets policy names from key and team metadata
+    2. Resolves guardrails from those policies (including inheritance)
+    3. Adds resolved guardrails to request metadata
+
+    Args:
+        key_metadata: The key metadata dictionary to check for policies
+        team_metadata: The team metadata dictionary to check for policies
+        data: The request data to update
+        metadata_variable_name: The name of the metadata field in data
+    """
+    from litellm._logging import verbose_proxy_logger
+    from litellm.proxy.policy_engine.policy_registry import get_policy_registry
+    from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
+    from litellm.proxy.utils import _premium_user_check
+    from litellm.types.proxy.policy_engine import PolicyMatchContext
+
+    # Collect policy names from key and team metadata
+    policy_names: set = set()
+
+    # Add key-level policies first
+    if key_metadata and "policies" in key_metadata:
+        if (
+            isinstance(key_metadata["policies"], list)
+            and len(key_metadata["policies"]) > 0
+        ):
+            _premium_user_check()
+            policy_names.update(key_metadata["policies"])
+
+    # Add team-level policies
+    if team_metadata and "policies" in team_metadata:
+        if (
+            isinstance(team_metadata["policies"], list)
+            and len(team_metadata["policies"]) > 0
+        ):
+            _premium_user_check()
+            policy_names.update(team_metadata["policies"])
+
+    if not policy_names:
+        return
+
+    verbose_proxy_logger.debug(
+        f"Policy engine: resolving guardrails from key/team policies: {policy_names}"
+    )
+
+    # Check if policy registry is initialized
+    registry = get_policy_registry()
+    if not registry.is_initialized():
+        verbose_proxy_logger.debug(
+            "Policy engine not initialized, skipping policy resolution from metadata"
+        )
+        return
+
+    # Build context for policy resolution (model from request data)
+    context = PolicyMatchContext(model=data.get("model"))
+
+    # Get all policies from registry
+    all_policies = registry.get_all_policies()
+
+    # Resolve guardrails from the specified policies
+    resolved_guardrails: set = set()
+    for policy_name in policy_names:
+        if registry.has_policy(policy_name):
+            resolved_policy = PolicyResolver.resolve_policy_guardrails(
+                policy_name=policy_name,
+                policies=all_policies,
+                context=context,
+            )
+            resolved_guardrails.update(resolved_policy.guardrails)
+            verbose_proxy_logger.debug(
+                f"Policy engine: resolved guardrails from policy '{policy_name}': {resolved_policy.guardrails}"
+            )
+        else:
+            verbose_proxy_logger.warning(
+                f"Policy engine: policy '{policy_name}' not found in registry"
+            )
+
+    if not resolved_guardrails:
+        return
+
+    # Add resolved guardrails to request metadata
+    if metadata_variable_name not in data:
+        data[metadata_variable_name] = {}
+
+    existing_guardrails = data[metadata_variable_name].get("guardrails", [])
+    if not isinstance(existing_guardrails, list):
+        existing_guardrails = []
+
+    # Combine existing guardrails with policy-resolved guardrails (no duplicates)
+    combined = set(existing_guardrails)
+    combined.update(resolved_guardrails)
+    data[metadata_variable_name]["guardrails"] = list(combined)
+
+    # Store applied policies in metadata for tracking
+    if "applied_policies" not in data[metadata_variable_name]:
+        data[metadata_variable_name]["applied_policies"] = []
+    data[metadata_variable_name]["applied_policies"].extend(list(policy_names))
+
+    verbose_proxy_logger.debug(
+        f"Policy engine: added guardrails from key/team policies to request metadata: {list(resolved_guardrails)}"
+    )
+
+
 def move_guardrails_to_metadata(
     data: dict,
     _metadata_variable_name: str,
@@ -1321,6 +1433,7 @@ def move_guardrails_to_metadata(
 
     - If guardrails set on API Key metadata then sets guardrails on request metadata
     - If guardrails not set on API key, then checks request metadata
+    - Adds guardrails from policies attached to key/team metadata
     - Adds guardrails from policy engine based on team/key/model context
     """
     # Check key-level guardrails
@@ -1331,6 +1444,16 @@ def move_guardrails_to_metadata(
         metadata_variable_name=_metadata_variable_name,
     )
 
+    #########################################################################################
+    # Add guardrails from policies attached to key/team metadata
+    #########################################################################################
+    _add_guardrails_from_policies_in_metadata(
+        key_metadata=user_api_key_dict.metadata,
+        team_metadata=user_api_key_dict.team_metadata,
+        data=data,
+        metadata_variable_name=_metadata_variable_name,
+    )
+
     #########################################################################################
     # Add guardrails from policy engine based on team/key/model context
     #########################################################################################

litellm/proxy/management_endpoints/internal_user_endpoints.py

@@ -14,7 +14,6 @@
 import asyncio
 import json
 import traceback
-from litellm._uuid import uuid
 from datetime import datetime, timezone
 from typing import Any, Dict, List, Optional, Union, cast
 
@@ -23,6 +22,7 @@
 
 import litellm
 from litellm._logging import verbose_proxy_logger
+from litellm._uuid import uuid
 from litellm.proxy._types import *
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.hooks.user_management_event_hooks import UserManagementEventHooks
@@ -355,6 +355,7 @@ async def new_user(
     - allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request-
     - blocked: Optional[bool] - [Not Implemented Yet] Whether the user is blocked.
     - guardrails: Optional[List[str]] - [Not Implemented Yet] List of active guardrails for the user
+    - policies: Optional[List[str]] - List of policy names to apply to the user. Policies define guardrails, conditions, and inheritance rules.
     - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
     - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
     - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
@@ -1060,6 +1061,7 @@ async def user_update(
         - allowed_cache_controls: Optional[list] - List of allowed cache control values. Example - ["no-cache", "no-store"]. See all values - https://docs.litellm.ai/docs/proxy/caching#turn-on--off-caching-per-request-
         - blocked: Optional[bool] - [Not Implemented Yet] Whether the user is blocked.
         - guardrails: Optional[List[str]] - [Not Implemented Yet] List of active guardrails for the user
+        - policies: Optional[List[str]] - List of policy names to apply to the user. Policies define guardrails, conditions, and inheritance rules.
         - permissions: Optional[dict] - [Not Implemented Yet] User-specific permissions, eg. turning off pii masking.
         - metadata: Optional[dict] - Metadata for user, store information for user. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
         - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -14,11 +14,11 @@
 import json
 import secrets
 import traceback
-import yaml
 from datetime import datetime, timedelta, timezone
 from typing import Any, Dict, List, Literal, Optional, Tuple, cast
-from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+
 import fastapi
+import yaml
 from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status
 
 import litellm
@@ -31,6 +31,7 @@
     UI_SESSION_TOKEN_TEAM_ID,
 )
 from litellm.litellm_core_utils.duration_parser import duration_in_seconds
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._experimental.mcp_server.db import (
     rotate_mcp_server_credentials_master_key,
 )
@@ -1010,6 +1011,7 @@ async def generate_key_fn(
     - max_parallel_requests: Optional[int] - Rate limit a user based on the number of parallel requests. Raises 429 error, if user's parallel requests > x.
     - metadata: Optional[dict] - Metadata for key, store information for key. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
     - guardrails: Optional[List[str]] - List of active guardrails for the key
+    - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
     - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
     - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false}
     - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget.
@@ -1480,6 +1482,7 @@ async def update_key_fn(
     - permissions: Optional[dict] - Key-specific permissions
     - send_invite_email: Optional[bool] - Send invite email to user_id
     - guardrails: Optional[List[str]] - List of active guardrails for the key
+    - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules.
     - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
     - prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
     - blocked: Optional[bool] - Whether the key is blocked
@@ -2077,6 +2080,7 @@ async def generate_key_helper_fn(  # noqa: PLR0915
     model_rpm_limit: Optional[dict] = None,
     model_tpm_limit: Optional[dict] = None,
     guardrails: Optional[list] = None,
+    policies: Optional[list] = None,
     prompts: Optional[list] = None,
     teams: Optional[list] = None,
     organization_id: Optional[str] = None,
@@ -2139,6 +2143,9 @@ async def generate_key_helper_fn(  # noqa: PLR0915
     if guardrails is not None:
         metadata = metadata or {}
         metadata["guardrails"] = guardrails
+    if policies is not None:
+        metadata = metadata or {}
+        metadata["policies"] = policies
     if prompts is not None:
         metadata = metadata or {}
         metadata["prompts"] = prompts

litellm/proxy/management_endpoints/team_endpoints.py

@@ -22,19 +22,20 @@
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import (
     BlockTeamRequest,
     CommonProxyErrors,
     DeleteTeamRequest,
     LiteLLM_AuditLogs,
+    LiteLLM_DeletedTeamTable,
     LiteLLM_ManagementEndpoint_MetadataFields,
     LiteLLM_ManagementEndpoint_MetadataFields_Premium,
     LiteLLM_ModelTable,
     LiteLLM_OrganizationTable,
     LiteLLM_OrganizationTableWithMembers,
     LiteLLM_TeamMembership,
     LiteLLM_TeamTable,
-    LiteLLM_DeletedTeamTable,
     LiteLLM_TeamTableCachedObj,
     LiteLLM_UserTable,
     LiteLLM_VerificationToken,
@@ -102,7 +103,7 @@
     TeamMemberAddResult,
     UpdateTeamMemberPermissionsRequest,
 )
-from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+
 router = APIRouter()
 
 
@@ -689,6 +690,7 @@ async def new_team(  # noqa: PLR0915
     - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
     - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
     - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
+    - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
     - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
     - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
     - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.
@@ -1228,6 +1230,7 @@ async def update_team(   # noqa: PLR0915
     - organization_id: Optional[str] - The organization id of the team. Default is None. Create via `/organization/new`.
     - model_aliases: Optional[dict] - Model aliases for the team. [Docs](https://docs.litellm.ai/docs/proxy/team_based_routing#create-team-with-model-alias)
     - guardrails: Optional[List[str]] - Guardrails for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails)
+    - policies: Optional[List[str]] - Policies for the team. [Docs](https://docs.litellm.ai/docs/proxy/guardrails/guardrail_policies)
     - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key.
     - object_permission: Optional[LiteLLM_ObjectPermissionBase] - team-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "agents": ["agent_1", "agent_2"], "agent_access_groups": ["dev_group"]}. IF null or {} then no object permission.
     - team_member_budget: Optional[float] - The maximum budget allocated to an individual team member.