Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/policy_engine/init_policies.py: 22%
110 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Policy Initialization - Loads policies from config and validates on startup.
4Configuration structure:
5- policies: Define WHAT guardrails to apply (with inheritance and conditions)
6- policy_attachments: Define WHERE policies apply (teams, keys, models)
7"""
9from collections.abc import Mapping, Sequence
10from typing import TYPE_CHECKING, Any, Final, Optional
12from litellm._logging import verbose_proxy_logger
13from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
14from litellm.proxy.policy_engine.policy_registry import get_policy_registry
15from litellm.proxy.policy_engine.policy_validator import PolicyValidator
16from litellm.types.proxy.policy_engine import PolicyValidationResponse
18if TYPE_CHECKING: 18 ↛ 19line 18 didn't jump to line 19 because the condition on line 18 was never true
19 from litellm.proxy.utils import PrismaClient
21# ANSI color codes for terminal output
22_green_color_code: Final = "\033[92m"
23_blue_color_code: Final = "\033[94m"
24_yellow_color_code: Final = "\033[93m"
25_reset_color_code: Final = "\033[0m"
28def _print_policies_on_startup(
29 policies_config: Mapping[str, Mapping[str, object]],
30 policy_attachments_config: Sequence[Mapping[str, object]] | None = None,
31) -> None:
32 """
33 Print loaded policies to console on startup (similar to model list).
34 """
35 import sys
37 print( # noqa: T201
38 f"{_green_color_code}\nLiteLLM Policy Engine: Loaded {len(policies_config)} policies{_reset_color_code}\n"
39 )
40 sys.stdout.flush()
42 for policy_name, policy_data in policies_config.items():
43 guardrails = policy_data.get("guardrails", {})
44 inherit = policy_data.get("inherit")
45 condition = policy_data.get("condition")
46 description = policy_data.get("description")
48 guardrails_add = guardrails.get("add", []) if isinstance(guardrails, dict) else []
49 guardrails_remove = guardrails.get("remove", []) if isinstance(guardrails, dict) else []
50 inherit_str = f" (inherits: {inherit})" if inherit else ""
52 print( # noqa: T201
53 f"{_blue_color_code} - {policy_name}{inherit_str}{_reset_color_code}"
54 )
55 if description:
56 print(f" description: {description}") # noqa: T201
57 if guardrails_add:
58 print(f" guardrails.add: {guardrails_add}") # noqa: T201
59 if guardrails_remove:
60 print(f" guardrails.remove: {guardrails_remove}") # noqa: T201
61 if condition:
62 model_condition = condition.get("model") if isinstance(condition, dict) else None
63 if model_condition:
64 print(f" condition.model: {model_condition}") # noqa: T201
66 # Print attachments
67 if policy_attachments_config:
68 print( # noqa: T201
69 f"\n{_yellow_color_code}Policy Attachments: {len(policy_attachments_config)} attachment(s){_reset_color_code}"
70 )
71 for attachment in policy_attachments_config:
72 policy = attachment.get("policy", "unknown")
73 scope = attachment.get("scope")
74 teams = attachment.get("teams")
75 keys = attachment.get("keys")
76 models = attachment.get("models")
78 scope_parts = []
79 if scope == "*":
80 scope_parts.append("scope=* (global)")
81 if teams:
82 scope_parts.append(f"teams={teams}")
83 if keys:
84 scope_parts.append(f"keys={keys}")
85 if models:
86 scope_parts.append(f"models={models}")
87 scope_str = ", ".join(scope_parts) if scope_parts else "all"
89 print(f" - {policy} -> {scope_str}") # noqa: T201
90 else:
91 print( # noqa: T201
92 f"\n{_yellow_color_code}Warning: No policy_attachments configured. Policies will not be applied to any requests.{_reset_color_code}"
93 )
95 print() # noqa: T201
96 sys.stdout.flush()
99async def init_policies(
100 policies_config: dict[str, Any],
101 policy_attachments_config: list[dict[str, Any]] | None = None,
102 prisma_client: Optional["PrismaClient"] = None,
103 validate_db: bool = True,
104 fail_on_error: bool = True,
105) -> PolicyValidationResponse:
106 """
107 Initialize policies from configuration.
109 This function:
110 1. Parses the policy configuration
111 2. Validates policies (guardrails exist, teams/keys exist in DB)
112 3. Loads policies into the global registry
113 4. Loads attachments into the attachment registry (if provided)
115 Args:
116 policies_config: Dictionary mapping policy names to policy definitions
117 policy_attachments_config: Optional list of policy attachment configurations
118 prisma_client: Optional Prisma client for database validation
119 validate_db: Whether to validate team/key aliases against database
120 fail_on_error: If True, raise exception on validation errors
122 Returns:
123 PolicyValidationResponse with validation results
125 Raises:
126 ValueError: If fail_on_error is True and validation errors are found
127 """
128 verbose_proxy_logger.info("Initializing %s policies...", len(policies_config))
130 # Print policies to console on startup
131 _print_policies_on_startup(policies_config, policy_attachments_config)
133 # Get the global registries
134 policy_registry: Final = get_policy_registry()
135 attachment_registry: Final = get_attachment_registry()
137 # Create validator
138 validator: Final = PolicyValidator(prisma_client=prisma_client)
140 # Validate the configuration
141 validation_result: Final = await validator.validate_policy_config(
142 policies_config,
143 validate_db=validate_db,
144 )
146 # Log validation results
147 if validation_result.errors:
148 for error in validation_result.errors:
149 verbose_proxy_logger.error(
150 "Policy validation error in '%s': [%s] %s", error.policy_name, error.error_type, error.message
151 )
153 if validation_result.warnings:
154 for warning in validation_result.warnings:
155 verbose_proxy_logger.warning(
156 "Policy validation warning in '%s': [%s] %s", warning.policy_name, warning.error_type, warning.message
157 )
159 # Fail if there are errors and fail_on_error is True
160 if not validation_result.valid and fail_on_error:
161 error_messages: Final = [f"[{e.policy_name}] {e.message}" for e in validation_result.errors]
162 raise ValueError(
163 f"Policy validation failed with {len(validation_result.errors)} error(s):\n" + "\n".join(error_messages)
164 )
166 # Load policies into registry (even with warnings)
167 try:
168 policy_registry.load_policies(policies_config)
169 verbose_proxy_logger.info("Successfully loaded %s policies", len(policies_config))
170 except Exception as e:
171 verbose_proxy_logger.error("Failed to load policies: %s", e)
172 raise
174 # Load attachments if provided
175 if policy_attachments_config:
176 try:
177 attachment_registry.load_attachments(policy_attachments_config)
178 verbose_proxy_logger.info("Successfully loaded %s policy attachments", len(policy_attachments_config))
179 except Exception as e:
180 verbose_proxy_logger.error("Failed to load policy attachments: %s", e)
181 raise
183 return validation_result
186def init_policies_sync(
187 policies_config: dict[str, Any],
188 policy_attachments_config: list[dict[str, Any]] | None = None,
189 fail_on_error: bool = True,
190) -> None:
191 """
192 Synchronous version of init_policies (without DB validation).
194 Use this when async is not available or DB validation is not needed.
196 Args:
197 policies_config: Dictionary mapping policy names to policy definitions
198 policy_attachments_config: Optional list of policy attachment configurations
199 fail_on_error: If True, raise exception on validation errors
200 """
201 import asyncio
203 # Run the async function without DB validation
204 try:
205 loop = asyncio.get_event_loop()
206 except RuntimeError:
207 loop = asyncio.new_event_loop()
208 asyncio.set_event_loop(loop)
210 loop.run_until_complete(
211 init_policies(
212 policies_config=policies_config,
213 policy_attachments_config=policy_attachments_config,
214 prisma_client=None,
215 validate_db=False,
216 fail_on_error=fail_on_error,
217 )
218 )
221def get_policies_summary() -> dict[str, Any]:
222 """
223 Get a summary of loaded policies for debugging/display.
225 Returns:
226 Dictionary with policy information
227 """
228 from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
230 policy_registry: Final = get_policy_registry()
231 attachment_registry: Final = get_attachment_registry()
233 if not policy_registry.is_initialized(): 233 ↛ 234line 233 didn't jump to line 234 because the condition on line 233 was never true
234 return {"initialized": False, "policies": {}, "attachments": []}
236 resolved: Final = PolicyResolver.get_all_resolved_policies()
238 summary: Final[dict[str, Any]] = {
239 "initialized": True,
240 "policy_count": len(resolved),
241 "attachment_count": len(attachment_registry.get_all_attachments()),
242 "policies": {},
243 "attachments": [],
244 }
246 for policy_name, resolved_policy in resolved.items():
247 policy = policy_registry.get_policy(policy_name)
248 summary["policies"][policy_name] = {
249 "inherit": policy.inherit if policy else None,
250 "description": policy.description if policy else None,
251 "guardrails_add": policy.guardrails.get_add() if policy else [],
252 "guardrails_remove": policy.guardrails.get_remove() if policy else [],
253 "condition": (policy.condition.model_dump() if policy and policy.condition else None),
254 "resolved_guardrails": resolved_policy.guardrails,
255 "inheritance_chain": resolved_policy.inheritance_chain,
256 }
258 # Add attachment info
259 for attachment in attachment_registry.get_all_attachments():
260 summary["attachments"].append(
261 {
262 "policy": attachment.policy,
263 "scope": attachment.scope,
264 "teams": attachment.teams,
265 "keys": attachment.keys,
266 "models": attachment.models,
267 }
268 )
270 return summary