From 55b33efb48d3bb70446e5a017757cf34e2cc4bb1 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Thu, 22 Jan 2026 18:09:24 -0800 Subject: [PATCH] fix policies --- litellm/proxy/litellm_pre_call_utils.py | 34 ++++- litellm/proxy/policy_engine/init_policies.py | 12 +- litellm/proxy/policy_engine/policy_matcher.py | 43 ++++++ litellm/proxy/proxy_config.yaml | 129 +++++++++++++----- litellm/proxy/proxy_server.py | 9 +- 5 files changed, 179 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 38c07a687f7..21be88c2635 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1321,6 +1321,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 policy engine based on team/key/model context """ # Check key-level guardrails _add_guardrails_from_key_or_team_metadata( @@ -1330,6 +1331,15 @@ def move_guardrails_to_metadata( metadata_variable_name=_metadata_variable_name, ) + ######################################################################################### + # Add guardrails from policy engine based on team/key/model context + ######################################################################################### + add_guardrails_from_policy_engine( + data=data, + metadata_variable_name=_metadata_variable_name, + user_api_key_dict=user_api_key_dict, + ) + ######################################################################################### # User's might send "guardrails" in the request body, we need to add them to the request metadata. # Since downstream logic requires "guardrails" to be in the request metadata @@ -1407,19 +1417,33 @@ def add_guardrails_from_policy_engine( f"key_alias={context.key_alias}, model={context.model}" ) + from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator + # Get matching policies via attachments matching_policy_names = PolicyMatcher.get_matching_policies(context=context) - verbose_proxy_logger.debug(f"Policy engine: matched policies: {matching_policy_names}") + verbose_proxy_logger.debug(f"Policy engine: matched policies via attachments: {matching_policy_names}") if not matching_policy_names: return - # Track applied policies in metadata + # Get all policies to check conditions + all_policies = registry.get_all_policies() + + # Track applied policies - only include policies whose conditions actually match + applied_policy_names = [] for policy_name in matching_policy_names: - add_policy_to_applied_policies_header( - request_data=data, policy_name=policy_name - ) + policy = all_policies.get(policy_name) + if policy is None: + continue + # Check if policy condition matches (or has no condition) + if policy.condition is None or ConditionEvaluator.evaluate(policy.condition, context): + applied_policy_names.append(policy_name) + add_policy_to_applied_policies_header( + request_data=data, policy_name=policy_name + ) + + verbose_proxy_logger.debug(f"Policy engine: applied policies (conditions matched): {applied_policy_names}") # Resolve guardrails from matching policies resolved_guardrails = PolicyResolver.resolve_guardrails_for_context(context=context) diff --git a/litellm/proxy/policy_engine/init_policies.py b/litellm/proxy/policy_engine/init_policies.py index c2dc0725f6a..b734c0cb5cc 100644 --- a/litellm/proxy/policy_engine/init_policies.py +++ b/litellm/proxy/policy_engine/init_policies.py @@ -2,7 +2,7 @@ Policy Initialization - Loads policies from config and validates on startup. Configuration structure: -- policies: Define WHAT guardrails to apply (with inheritance and statements) +- policies: Define WHAT guardrails to apply (with inheritance and conditions) - policy_attachments: Define WHERE policies apply (teams, keys, models) """ @@ -41,7 +41,7 @@ def _print_policies_on_startup( for policy_name, policy_data in policies_config.items(): guardrails = policy_data.get("guardrails", {}) inherit = policy_data.get("inherit") - statements = policy_data.get("statements", []) + condition = policy_data.get("condition") description = policy_data.get("description") guardrails_add = guardrails.get("add", []) if isinstance(guardrails, dict) else [] @@ -57,8 +57,10 @@ def _print_policies_on_startup( print(f" guardrails.add: {guardrails_add}") # noqa: T201 if guardrails_remove: print(f" guardrails.remove: {guardrails_remove}") # noqa: T201 - if statements: - print(f" statements: {len(statements)} conditional statement(s)") # noqa: T201 + if condition: + model_condition = condition.get("model") if isinstance(condition, dict) else None + if model_condition: + print(f" condition.model: {model_condition}") # noqa: T201 # Print attachments if policy_attachments_config: @@ -256,7 +258,7 @@ def get_policies_summary() -> Dict[str, Any]: "description": policy.description if policy else None, "guardrails_add": policy.guardrails.get_add() if policy else [], "guardrails_remove": policy.guardrails.get_remove() if policy else [], - "statements_count": len(policy.statements) if policy and policy.statements else 0, + "condition": policy.condition.model_dump() if policy and policy.condition else None, "resolved_guardrails": resolved_policy.guardrails, "inheritance_chain": resolved_policy.inheritance_chain, } diff --git a/litellm/proxy/policy_engine/policy_matcher.py b/litellm/proxy/policy_engine/policy_matcher.py index 0b8c5e9544b..ab73970bfab 100644 --- a/litellm/proxy/policy_engine/policy_matcher.py +++ b/litellm/proxy/policy_engine/policy_matcher.py @@ -123,3 +123,46 @@ class PolicyMatcher: List of policy names that match the context """ return PolicyMatcher.get_matching_policies(context=context) + + @staticmethod + def get_policies_with_matching_conditions( + policy_names: List[str], + context: PolicyMatchContext, + policies: Optional[Dict[str, Policy]] = None, + ) -> List[str]: + """ + Filter policies to only those whose conditions match the context. + + A policy's condition matches if: + - The policy has no condition (condition is None), OR + - The policy's condition evaluates to True for the given context + + Args: + policy_names: List of policy names to filter + context: The request context to evaluate conditions against + policies: Dictionary of all policies (if None, uses global registry) + + Returns: + List of policy names whose conditions match the context + """ + from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + + if policies is None: + registry = get_policy_registry() + if not registry.is_initialized(): + return [] + policies = registry.get_all_policies() + + matching_policies = [] + for policy_name in policy_names: + policy = policies.get(policy_name) + if policy is None: + continue + # Policy matches if it has no condition OR condition evaluates to True + if policy.condition is None or ConditionEvaluator.evaluate( + policy.condition, context + ): + matching_policies.append(policy_name) + + return matching_policies diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 958ddbf613c..ea405c1dea4 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,42 +1,101 @@ model_list: - # Anthropic direct - - model_name: anthropic-claude + - model_name: "*" litellm_params: - model: anthropic/claude-sonnet-4-20250514 - api_key: os.environ/ANTHROPIC_API_KEY - - # Azure AI Anthropic - - model_name: azure-ai-claude + model: "*" + - model_name: "gpt-4" litellm_params: - model: azure_ai/claude-3-5-sonnet - api_base: https://krish-mh44t553-eastus2.services.ai.azure.com/ - api_key: os.environ/AZURE_ANTHROPIC_API_KEY - - # Azure AI Anthropic (alternate endpoint format) - - model_name: claude-4.5-haiku + model: "gpt-4" + api_key: os.environ/OPENAI_API_KEY + - model_name: "gpt-3.5-turbo" litellm_params: - model: anthropic/claude-haiku-4-5 - api_base: https://krish-mh44t553-eastus2.services.ai.azure.com/anthropic/v1/messages - api_version: "2023-06-01" - api_key: os.environ/AZURE_ANTHROPIC_API_KEY - - - -# Search Tools Configuration - Define search providers for WebSearch interception -# search_tools: -# - search_tool_name: "my-perplexity-search" -# litellm_params: -# search_provider: "perplexity" # Can be: perplexity, brave, etc. - -litellm_settings: - callbacks: ["websearch_interception"] - # WebSearch Interception - Automatically intercepts and executes WebSearch tool calls - # for models that don't natively support web search (e.g., Bedrock/Claude) - websearch_interception_params: - enabled_providers: ["bedrock"] # List of providers to enable interception for - search_tool_name: "my-perplexity-search" # Optional: Name of search tool from search_tools config + model: "gpt-3.5-turbo" + api_key: os.environ/OPENAI_API_KEY general_settings: - store_prompts_in_spend_logs: true - forward_client_headers_to_llm_api: true + master_key: sk-1234 +# ─────────────────────────────────────────────── +# POLICIES - Define WHAT guardrails to apply +# ─────────────────────────────────────────────── +# +# Policies define guardrails with: +# - inherit: Inherit guardrails from another policy +# - description: Human-readable description +# - guardrails.add: Add guardrails (on top of inherited) +# - guardrails.remove: Remove guardrails (from inherited) +# - condition.model: Model pattern (exact or regex) for when guardrails apply +# +policies: + # Global baseline policy + global-baseline: + description: "Base guardrails for all requests" + guardrails: + add: + - pii_blocker + + # Healthcare policy - inherits from global-baseline + healthcare-compliance: + inherit: global-baseline + description: "HIPAA compliance for healthcare teams" + guardrails: + add: + - hipaa_audit + + # Dev policy - inherits but removes PII blocker for testing + internal-dev: + inherit: global-baseline + description: "Relaxed policy for internal development" + guardrails: + add: + - toxicity_filter + remove: + - pii_blocker + + # Policy with model condition (regex pattern) + gpt4-safety: + description: "Extra safety for GPT-4 models" + guardrails: + add: + - toxicity_filter + condition: + model: "gpt-4.*" # regex: matches gpt-4, gpt-4-turbo, gpt-4o, etc. + + # Policy with model condition (exact match list) + bedrock-compliance: + description: "Compliance for Bedrock models" + guardrails: + add: + - strict_pii_blocker + condition: + model: ["bedrock/claude-3", "bedrock/claude-2"] # exact matches + +# ─────────────────────────────────────────────── +# POLICY ATTACHMENTS - Define WHERE policies apply +# ─────────────────────────────────────────────── +# +# Attachments are REQUIRED to make policies active. +# A policy without an attachment will not be applied. +# +policy_attachments: + # Global attachment - applies to all requests + - policy: global-baseline + scope: "*" + + # Team-specific attachment + - policy: healthcare-compliance + teams: + - healthcare-team + - medical-research + + # Key pattern attachment + - policy: internal-dev + keys: + - "dev-key-*" + - "test-key-*" + + # Model-specific policies (attached globally, condition filters by model) + - policy: gpt4-safety + scope: "*" + + - policy: bedrock-compliance + scope: "*" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ca2ace47b74..eb4d6bd1501 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2853,6 +2853,9 @@ class ProxyConfig: prisma_client: Optional Prisma client for DB validation llm_router: Optional LLM router for model validation """ + + from litellm.proxy.policy_engine.init_policies import init_policies + from litellm.proxy.policy_engine.policy_validator import PolicyValidator if config is None: verbose_proxy_logger.debug("Policy engine: config is None, skipping") return @@ -2862,10 +2865,9 @@ class ProxyConfig: verbose_proxy_logger.debug("Policy engine: no policies in config, skipping") return - verbose_proxy_logger.info(f"Policy engine: found {len(policies_config)} policies in config") + policy_attachments_config = config.get("policy_attachments", None) - from litellm.proxy.policy_engine.init_policies import init_policies - from litellm.proxy.policy_engine.policy_validator import PolicyValidator + verbose_proxy_logger.info(f"Policy engine: found {len(policies_config)} policies in config") # Create validator with router for model validation validator = PolicyValidator( @@ -2876,6 +2878,7 @@ class ProxyConfig: # Initialize policies await init_policies( policies_config=policies_config, + policy_attachments_config=policy_attachments_config, prisma_client=prisma_client, validate_db=prisma_client is not None, fail_on_error=True,