Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/policy_engine/attachment_registry.py: 63%
154 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1"""
2Attachment Registry - Manages policy attachments from YAML config.
4Attachments define WHERE policies apply, separate from the policy definitions.
5This allows the same policy to be attached to multiple scopes.
6"""
8from collections.abc import Callable
9from datetime import datetime, timezone
10from types import MappingProxyType
11from typing import TYPE_CHECKING, Any, Final, TypedDict
13from litellm._logging import verbose_proxy_logger
14from litellm.repositories.table_repositories import PolicyAttachmentRepository
15from litellm.types.proxy.policy_engine import (
16 PolicyAttachment,
17 PolicyAttachmentCreateRequest,
18 PolicyAttachmentDBResponse,
19 PolicyMatchContext,
20)
22if TYPE_CHECKING: 22 ↛ 23line 22 didn't jump to line 23 because the condition on line 22 was never true
23 from collections.abc import Sequence
25 from prisma.models import LiteLLM_PolicyAttachmentTable
27 from litellm.proxy.utils import PrismaClient
30class PolicyAttachmentMatch(TypedDict):
31 policy_name: str
32 matched_via: str
35def _attachment_specificity(attachment: PolicyAttachment) -> tuple[int, int]:
36 if attachment.is_global(): 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true
37 return (0, 0)
39 dims: Final = tuple(
40 specificity
41 for values, specificity in (
42 (attachment.teams, 1),
43 (attachment.keys, 2),
44 (attachment.tags, 3),
45 (attachment.models, 4),
46 )
47 if values
48 )
49 return (max(dims, default=0), len(dims))
52def _attachment_sort_key(attachment: PolicyAttachment) -> tuple[int, int, int, int]:
53 specificity: Final = _attachment_specificity(attachment)
54 if attachment.priority is not None:
55 return (0, attachment.priority, *specificity)
56 return (1, 0, *specificity)
59class AttachmentRegistry:
60 """
61 In-memory registry for storing and managing policy attachments.
63 Attachments define the relationship between policies and their scopes.
64 A single policy can have multiple attachments (applied to different scopes).
66 Example YAML:
67 ```yaml
68 attachments:
69 - policy: global-baseline
70 scope: "*"
71 - policy: healthcare-compliance
72 teams: [healthcare-team]
73 - policy: dev-safety
74 keys: ["dev-key-*"]
75 ```
76 """
78 def __init__(self) -> None:
79 self._attachments: list[PolicyAttachment] = []
80 self._config_attachments: tuple[PolicyAttachment, ...] = ()
81 self._initialized: bool = False
83 def load_attachments(self, attachments_config: list[dict[str, Any]]) -> None:
84 """
85 Load attachments from a configuration list.
87 Args:
88 attachments_config: List of attachment dictionaries from YAML.
89 """
90 self._attachments = []
92 for attachment_data in attachments_config:
93 try:
94 attachment = self._parse_attachment(attachment_data)
95 self._attachments.append(attachment)
96 verbose_proxy_logger.debug("Loaded attachment for policy: %s", attachment.policy)
97 except Exception as e:
98 verbose_proxy_logger.error("Error loading attachment: %s", e)
99 raise ValueError(f"Invalid attachment: {e}") from e
101 self._config_attachments = tuple(self._attachments)
102 self._initialized = True
103 verbose_proxy_logger.info("Loaded %s policy attachments", len(self._attachments))
105 def _parse_attachment(self, attachment_data: dict[str, Any]) -> PolicyAttachment:
106 """
107 Parse an attachment from raw configuration data.
109 Args:
110 attachment_data: Raw attachment configuration
112 Returns:
113 Parsed PolicyAttachment object
114 """
115 return PolicyAttachment(
116 policy=attachment_data.get("policy", ""),
117 scope=attachment_data.get("scope"),
118 teams=attachment_data.get("teams"),
119 keys=attachment_data.get("keys"),
120 models=attachment_data.get("models"),
121 tags=attachment_data.get("tags"),
122 priority=attachment_data.get("priority"),
123 default=attachment_data.get("default", False),
124 )
126 def get_attached_policies(
127 self,
128 context: PolicyMatchContext,
129 policy_applies: Callable[[str], bool] | None = None,
130 ) -> list[str]:
131 """
132 Get list of policy names attached to the given context.
134 Args:
135 context: The request context to match against
136 policy_applies: Optional predicate; attachments whose policy does not apply are ignored
138 Returns:
139 List of policy names that are attached to matching scopes
140 """
141 return [r["policy_name"] for r in self.get_attached_policies_with_reasons(context, policy_applies)]
143 def get_attached_policies_with_reasons(
144 self,
145 context: PolicyMatchContext,
146 policy_applies: Callable[[str], bool] | None = None,
147 ) -> list[PolicyAttachmentMatch]:
148 """
149 Get list of policy names and match reasons for the given context.
151 Returns a list of dicts with 'policy_name' and 'matched_via' keys.
152 The 'matched_via' describes which dimension caused the match.
153 Attachments whose policy fails `policy_applies` are dropped before defaults are considered.
154 """
155 from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
157 in_scope: Final = tuple(
158 attachment
159 for attachment in self._attachments
160 if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context)
161 and (policy_applies is None or policy_applies(attachment.policy))
162 )
163 non_default: Final = tuple(attachment for attachment in in_scope if not attachment.default)
164 matching_attachments: Final = sorted(
165 non_default or tuple(attachment for attachment in in_scope if attachment.default),
166 key=_attachment_sort_key,
167 )
168 broadest_attachment_by_policy: Final = MappingProxyType(
169 {attachment.policy: attachment for attachment in reversed(matching_attachments)}
170 )
171 unique_attachments: Final = tuple(
172 broadest_attachment_by_policy[policy_name]
173 for policy_name in dict.fromkeys(attachment.policy for attachment in matching_attachments)
174 )
176 return [
177 {
178 "policy_name": attachment.policy,
179 "matched_via": self._describe_match_reason(attachment, context),
180 }
181 for attachment in unique_attachments
182 ]
184 @staticmethod
185 def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str:
186 """Describe why an attachment matched the context."""
187 reason: Final = AttachmentRegistry._describe_scope_match(attachment, context)
188 return f"default:{reason}" if attachment.default else reason
190 @staticmethod
191 def _describe_scope_match(attachment: PolicyAttachment, context: PolicyMatchContext) -> str:
192 from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
194 if attachment.is_global(): 194 ↛ 195line 194 didn't jump to line 195 because the condition on line 194 was never true
195 return "scope:*"
197 reasons: Final = []
198 if attachment.tags and context.tags: 198 ↛ 199line 198 didn't jump to line 199 because the condition on line 198 was never true
199 matching_tags: Final = [t for t in context.tags if PolicyMatcher.matches_pattern(t, attachment.tags)]
200 if matching_tags:
201 reasons.append(f"tag:{matching_tags[0]}")
202 if attachment.teams and context.team_alias: 202 ↛ 203line 202 didn't jump to line 203 because the condition on line 202 was never true
203 reasons.append(f"team:{context.team_alias}")
204 if attachment.keys and context.key_alias: 204 ↛ 205line 204 didn't jump to line 205 because the condition on line 204 was never true
205 reasons.append(f"key:{context.key_alias}")
206 if attachment.models and context.model: 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true
207 reasons.append(f"model:{context.model}")
209 return "+".join(reasons) if reasons else "scope:default"
211 def is_policy_attached(self, policy_name: str, context: PolicyMatchContext) -> bool:
212 """
213 Check if a specific policy is attached to the given context.
215 Args:
216 policy_name: Name of the policy to check
217 context: The request context to match against
219 Returns:
220 True if the policy is attached to a matching scope
221 """
222 attached: Final = self.get_attached_policies(context)
223 return policy_name in attached
225 def get_all_attachments(self) -> list[PolicyAttachment]:
226 """
227 Get all loaded attachments.
229 Returns:
230 List of all PolicyAttachment objects
231 """
232 return self._attachments.copy()
234 def get_config_attachments(self) -> tuple[PolicyAttachment, ...]:
235 """
236 Get the attachments loaded from config.yaml.
238 Returns:
239 Tuple of config-defined PolicyAttachment objects
240 """
241 return self._config_attachments
243 def get_attachments_for_policy(self, policy_name: str) -> list[PolicyAttachment]:
244 """
245 Get all attachments for a specific policy.
247 Args:
248 policy_name: Name of the policy
250 Returns:
251 List of attachments for the policy
252 """
253 return [a for a in self._attachments if a.policy == policy_name]
255 def is_initialized(self) -> bool:
256 """
257 Check if the registry has been initialized with attachments.
259 Returns:
260 True if attachments have been loaded, False otherwise
261 """
262 return self._initialized
264 def clear(self) -> None:
265 """
266 Clear all attachments from the registry.
267 """
268 self._attachments = []
269 self._config_attachments = ()
270 self._initialized = False
272 def add_attachment(self, attachment: PolicyAttachment) -> None:
273 """
274 Add a single attachment.
276 Args:
277 attachment: PolicyAttachment object to add
278 """
279 self._attachments.append(attachment)
280 self._initialized = True
281 verbose_proxy_logger.debug("Added attachment for policy: %s", attachment.policy)
283 def remove_attachments_for_policy(self, policy_name: str) -> int:
284 """
285 Remove all attachments for a specific policy.
287 Args:
288 policy_name: Name of the policy
290 Returns:
291 Number of attachments removed
292 """
293 original_count: Final = len(self._attachments)
294 self._attachments = [a for a in self._attachments if a.policy != policy_name]
295 removed_count: Final = original_count - len(self._attachments)
296 if removed_count > 0:
297 verbose_proxy_logger.debug("Removed %s attachment(s) for policy: %s", removed_count, policy_name)
298 return removed_count
300 def remove_attachment_by_id(self, attachment_id: str) -> bool:
301 """
302 Remove an attachment by its ID (for DB-synced attachments).
304 Args:
305 attachment_id: The ID of the attachment to remove
307 Returns:
308 True if removed, False if not found
309 """
310 # Note: In-memory attachments don't have IDs, so this is primarily
311 # for consistency after DB operations
312 return False
314 # ─────────────────────────────────────────────────────────────────────────
315 # Database CRUD Methods
316 # ─────────────────────────────────────────────────────────────────────────
318 async def add_attachment_to_db(
319 self,
320 attachment_request: PolicyAttachmentCreateRequest,
321 prisma_client: "PrismaClient",
322 created_by: str | None = None,
323 ) -> PolicyAttachmentDBResponse:
324 """
325 Add a policy attachment to the database.
327 Args:
328 attachment_request: The attachment creation request
329 prisma_client: The Prisma client instance
330 created_by: User who created the attachment
332 Returns:
333 PolicyAttachmentDBResponse with the created attachment
334 """
335 try:
336 created_attachment: Final[LiteLLM_PolicyAttachmentTable] = await PolicyAttachmentRepository(
337 prisma_client
338 ).table.create(
339 data={
340 "policy_name": attachment_request.policy_name,
341 "scope": attachment_request.scope,
342 "teams": attachment_request.teams or [],
343 "keys": attachment_request.keys or [],
344 "models": attachment_request.models or [],
345 "tags": attachment_request.tags or [],
346 "priority": attachment_request.priority,
347 "is_default": attachment_request.default,
348 "created_at": datetime.now(timezone.utc),
349 "updated_at": datetime.now(timezone.utc),
350 "created_by": created_by,
351 "updated_by": created_by,
352 }
353 )
355 # Also add to in-memory registry
356 attachment: Final = PolicyAttachment(
357 policy=attachment_request.policy_name,
358 scope=attachment_request.scope,
359 teams=attachment_request.teams,
360 keys=attachment_request.keys,
361 models=attachment_request.models,
362 tags=attachment_request.tags,
363 priority=attachment_request.priority,
364 default=attachment_request.default,
365 )
366 self.add_attachment(attachment)
368 return PolicyAttachmentDBResponse(
369 attachment_id=created_attachment.attachment_id,
370 policy_name=created_attachment.policy_name,
371 scope=created_attachment.scope,
372 teams=created_attachment.teams or [],
373 keys=created_attachment.keys or [],
374 models=created_attachment.models or [],
375 tags=created_attachment.tags or [],
376 priority=created_attachment.priority,
377 default=created_attachment.is_default,
378 created_at=created_attachment.created_at,
379 updated_at=created_attachment.updated_at,
380 created_by=created_attachment.created_by,
381 updated_by=created_attachment.updated_by,
382 )
383 except Exception as e:
384 verbose_proxy_logger.exception("Error adding attachment to DB: %s", e)
385 raise Exception(f"Error adding attachment to DB: {e}")
387 async def delete_attachment_from_db(
388 self,
389 attachment_id: str,
390 prisma_client: "PrismaClient",
391 ) -> dict[str, str]:
392 """
393 Delete a policy attachment from the database.
395 Args:
396 attachment_id: The ID of the attachment to delete
397 prisma_client: The Prisma client instance
399 Returns:
400 Dict with success message
401 """
402 try:
403 # Get attachment before deleting
404 attachment: Final[LiteLLM_PolicyAttachmentTable | None] = await PolicyAttachmentRepository(
405 prisma_client
406 ).table.find_unique(where={"attachment_id": attachment_id})
408 if attachment is None: 408 ↛ 409line 408 didn't jump to line 409 because the condition on line 408 was never true
409 raise Exception(f"Attachment with ID {attachment_id} not found")
411 # Delete from DB
412 await PolicyAttachmentRepository(prisma_client).table.delete(where={"attachment_id": attachment_id})
414 # Note: In-memory attachments don't have IDs, so we need to sync from DB
415 # to properly update in-memory state
416 await self.sync_attachments_from_db(prisma_client)
418 return {"message": f"Attachment {attachment_id} deleted successfully"}
419 except Exception as e:
420 verbose_proxy_logger.exception("Error deleting attachment from DB: %s", e)
421 raise Exception(f"Error deleting attachment from DB: {e}")
423 async def get_attachment_by_id_from_db(
424 self,
425 attachment_id: str,
426 prisma_client: "PrismaClient",
427 ) -> PolicyAttachmentDBResponse | None:
428 """
429 Get a policy attachment by ID from the database.
431 Args:
432 attachment_id: The ID of the attachment to retrieve
433 prisma_client: The Prisma client instance
435 Returns:
436 PolicyAttachmentDBResponse if found, None otherwise
437 """
438 try:
439 attachment: Final[LiteLLM_PolicyAttachmentTable | None] = await PolicyAttachmentRepository(
440 prisma_client
441 ).table.find_unique(where={"attachment_id": attachment_id})
443 if attachment is None:
444 return None
446 return PolicyAttachmentDBResponse(
447 attachment_id=attachment.attachment_id,
448 policy_name=attachment.policy_name,
449 scope=attachment.scope,
450 teams=attachment.teams or [],
451 keys=attachment.keys or [],
452 models=attachment.models or [],
453 tags=attachment.tags or [],
454 priority=attachment.priority,
455 default=attachment.is_default,
456 created_at=attachment.created_at,
457 updated_at=attachment.updated_at,
458 created_by=attachment.created_by,
459 updated_by=attachment.updated_by,
460 )
461 except Exception as e:
462 verbose_proxy_logger.exception("Error getting attachment from DB: %s", e)
463 raise Exception(f"Error getting attachment from DB: {e}")
465 async def get_all_attachments_from_db(
466 self,
467 prisma_client: "PrismaClient",
468 ) -> list[PolicyAttachmentDBResponse]:
469 """
470 Get all policy attachments from the database.
472 Args:
473 prisma_client: The Prisma client instance
475 Returns:
476 List of PolicyAttachmentDBResponse objects
477 """
478 try:
479 attachments: Final[Sequence[LiteLLM_PolicyAttachmentTable]] = await PolicyAttachmentRepository(
480 prisma_client
481 ).table.find_many(
482 order={"created_at": "desc"},
483 )
485 return [
486 PolicyAttachmentDBResponse(
487 attachment_id=a.attachment_id,
488 policy_name=a.policy_name,
489 scope=a.scope,
490 teams=a.teams or [],
491 keys=a.keys or [],
492 models=a.models or [],
493 tags=a.tags or [],
494 priority=a.priority,
495 default=a.is_default,
496 created_at=a.created_at,
497 updated_at=a.updated_at,
498 created_by=a.created_by,
499 updated_by=a.updated_by,
500 )
501 for a in attachments
502 ]
503 except Exception as e:
504 verbose_proxy_logger.exception("Error getting attachments from DB: %s", e)
505 raise Exception(f"Error getting attachments from DB: {e}")
507 async def sync_attachments_from_db(
508 self,
509 prisma_client: "PrismaClient",
510 ) -> None:
511 """
512 Sync policy attachments from the database to in-memory registry.
513 Config-loaded attachments are preserved.
515 Args:
516 prisma_client: The Prisma client instance
517 """
518 try:
519 attachments: Final = await self.get_all_attachments_from_db(prisma_client)
521 db_attachments: Final = [
522 PolicyAttachment(
523 policy=attachment_response.policy_name,
524 scope=attachment_response.scope,
525 teams=(attachment_response.teams if attachment_response.teams else None),
526 keys=attachment_response.keys if attachment_response.keys else None,
527 models=(attachment_response.models if attachment_response.models else None),
528 tags=attachment_response.tags if attachment_response.tags else None,
529 priority=attachment_response.priority,
530 default=attachment_response.default,
531 )
532 for attachment_response in attachments
533 ]
534 self._attachments = [*self._config_attachments, *db_attachments]
536 self._initialized = True
537 verbose_proxy_logger.info(
538 "Synced %s attachments from DB to in-memory registry (%s config-defined attachments preserved)",
539 len(attachments),
540 len(self._config_attachments),
541 )
542 except Exception as e:
543 verbose_proxy_logger.exception("Error syncing attachments from DB: %s", e)
544 raise Exception(f"Error syncing attachments from DB: {e}")
547# Global singleton instance
548_attachment_registry: AttachmentRegistry | None = None
551def get_attachment_registry() -> AttachmentRegistry:
552 """
553 Get the global AttachmentRegistry singleton.
555 Returns:
556 The global AttachmentRegistry instance
557 """
558 global _attachment_registry
559 if _attachment_registry is None:
560 _attachment_registry = AttachmentRegistry()
561 return _attachment_registry