diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ad4687e95db..c2805d00c2e 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -173,6 +173,7 @@ _ENABLE_TEAM_STALE_ALIAS_BYPASS: bool | None = None if TYPE_CHECKING: from litellm.integrations.otel.model.destination import OtelDestination + from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry from litellm.proxy.proxy_server import ProxyConfig as _ProxyConfig from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext @@ -3141,6 +3142,7 @@ def _match_and_track_policies( context: "PolicyMatchContext", request_body_policies: Sequence[str], policies_override: dict[str, "Policy"] | None = None, + attachment_registry_override: "AttachmentRegistry | None" = None, ) -> tuple[list[str], dict[str, str]]: """ Match policies via attachments and request body, track them in metadata. @@ -3157,7 +3159,9 @@ def _match_and_track_policies( from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher # Get matching policies via attachments (with match reasons for attribution) - attachment_registry: Final = get_attachment_registry() + attachment_registry: Final = ( + attachment_registry_override if attachment_registry_override is not None else get_attachment_registry() + ) matches_with_reasons: Final = attachment_registry.get_attached_policies_with_reasons(context) matching_policy_names: Final = [m["policy_name"] for m in matches_with_reasons] policy_reasons: Final = {m["policy_name"]: m["matched_via"] for m in matches_with_reasons} @@ -3165,9 +3169,11 @@ def _match_and_track_policies( verbose_proxy_logger.debug("Policy engine: matched policies via attachments: %s", matching_policy_names) # Combine attachment-based policies with dynamic request body policies - all_policy_names: Final = set(matching_policy_names) - if request_body_policies and isinstance(request_body_policies, list): - all_policy_names.update(request_body_policies) + request_body_policies_list: Final = ( + tuple(request_body_policies) if request_body_policies and isinstance(request_body_policies, list) else () + ) + all_policy_names: Final = tuple(dict.fromkeys((*matching_policy_names, *request_body_policies_list))) + if request_body_policies_list: verbose_proxy_logger.debug("Policy engine: added dynamic policies from request body: %s", request_body_policies) if not all_policy_names: @@ -3238,16 +3244,14 @@ def _apply_resolved_guardrails_to_metadata( if not resolved_guardrails and not pipelines: return - existing_guardrails = data[metadata_variable_name].get("guardrails", []) - if not isinstance(existing_guardrails, list): - existing_guardrails = [] + existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) + existing_guardrails_list: Final = existing_guardrails if isinstance(existing_guardrails, list) else [] # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined = set(existing_guardrails) - combined.update(resolved_guardrails) - data[metadata_variable_name]["guardrails"] = list(combined) + combined: Final = list(dict.fromkeys((*existing_guardrails_list, *resolved_guardrails))) + data[metadata_variable_name]["guardrails"] = combined - verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", list(combined)) + verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) async def add_guardrails_from_policy_engine( diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index 05df44242aa..a8ead86ac36 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -30,6 +30,23 @@ class PolicyAttachmentMatch(TypedDict): matched_via: str +def _attachment_specificity(attachment: PolicyAttachment) -> tuple[int, int]: + if attachment.is_global(): + return (0, 0) + + dims: Final = tuple( + specificity + for values, specificity in ( + (attachment.teams, 1), + (attachment.keys, 2), + (attachment.tags, 3), + (attachment.models, 4), + ) + if values + ) + return (max(dims, default=0), len(dims)) + + class AttachmentRegistry: """ In-memory registry for storing and managing policy attachments. @@ -116,31 +133,26 @@ class AttachmentRegistry: """ from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher - results: Final[list[PolicyAttachmentMatch]] = [] - seen_policies: Final[set[str]] = set() + matching_attachments: Final = sorted( + ( + attachment + for attachment in self._attachments + if PolicyMatcher.scope_matches(scope=attachment.to_policy_scope(), context=context) + ), + key=_attachment_specificity, + ) + unique_attachments: Final = tuple( + next(attachment for attachment in matching_attachments if attachment.policy == policy_name) + for policy_name in dict.fromkeys(attachment.policy for attachment in matching_attachments) + ) - for attachment in self._attachments: - scope = attachment.to_policy_scope() - if PolicyMatcher.scope_matches(scope=scope, context=context): - if attachment.policy not in seen_policies: - seen_policies.add(attachment.policy) - matched_via = self._describe_match_reason(attachment, context) - results.append( - { - "policy_name": attachment.policy, - "matched_via": matched_via, - } - ) - verbose_proxy_logger.debug( - "Attachment matched: policy=%s, matched_via=%s, context=(team=%s, key=%s, model=%s)", - attachment.policy, - matched_via, - context.team_alias, - context.key_alias, - context.model, - ) - - return results + return [ + { + "policy_name": attachment.policy, + "matched_via": self._describe_match_reason(attachment, context), + } + for attachment in unique_attachments + ] @staticmethod def _describe_match_reason(attachment: PolicyAttachment, context: PolicyMatchContext) -> str: diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index cc231e383a3..87b1bc56659 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -139,6 +139,71 @@ class TestGetAttachedPolicies: assert "gpt4-policy" in attached assert len(attached) == 3 + def test_matches_are_ordered_from_broadest_to_narrowest_scope(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "model-policy", "models": ["gpt-4"]}, + {"policy": "team-policy", "teams": ["t1"]}, + {"policy": "global-policy", "scope": "*"}, + ] + ) + + context = PolicyMatchContext(team_alias="t1", model="gpt-4") + + assert registry.get_attached_policies(context) == [ + "global-policy", + "team-policy", + "model-policy", + ] + + def test_combined_team_and_model_attachment_uses_model_specificity(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "team-policy", "teams": ["t1"]}, + {"policy": "team-model-policy", "teams": ["t1"], "models": ["gpt-4"]}, + ] + ) + + context = PolicyMatchContext(team_alias="t1", model="gpt-4") + + assert registry.get_attached_policies(context) == [ + "team-policy", + "team-model-policy", + ] + + def test_duplicate_policy_uses_broadest_matching_attachment(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "shared-policy", "models": ["gpt-4"]}, + {"policy": "model-policy", "models": ["gpt-4"]}, + {"policy": "shared-policy", "scope": "*"}, + ] + ) + + context = PolicyMatchContext(model="gpt-4") + + assert registry.get_attached_policies(context) == [ + "shared-policy", + "model-policy", + ] + assert registry.get_attached_policies_with_reasons(context)[0]["matched_via"] == "scope:*" + + def test_duplicate_policy_prefers_single_scope_over_combined_scope(self): + registry = AttachmentRegistry() + registry.load_attachments( + [ + {"policy": "shared-policy", "teams": ["t1"], "models": ["gpt-4"]}, + {"policy": "shared-policy", "models": ["gpt-4"]}, + ] + ) + + context = PolicyMatchContext(team_alias="t1", model="gpt-4") + + assert registry.get_attached_policies_with_reasons(context)[0]["matched_via"] == "model:gpt-4" + def test_same_policy_multiple_attachments_no_duplicates(self): """Test same policy attached multiple ways doesn't duplicate.""" registry = AttachmentRegistry() diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 94aaa32519e..ec9025a5220 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -23,6 +23,7 @@ from litellm.proxy.litellm_pre_call_utils import ( _get_dynamic_logging_metadata, _get_enforced_params, _get_metadata_variable_name, + _match_and_track_policies, _promoted_trace_control_fields, _resolve_credential_from_model_config, _resolve_provider_from_deployment, @@ -4149,6 +4150,30 @@ async def test_add_guardrails_from_policy_engine(): attachment_registry._initialized = False +def test_match_and_track_policies_preserves_attachment_and_request_body_order(): + from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry + from litellm.types.proxy.policy_engine import Policy, PolicyMatchContext + + attachment_policy_names = [f"attachment-policy-{index}" for index in range(8)] + request_body_policy_names = ["body-policy-1", "body-policy-2"] + policy_names = [*attachment_policy_names, *request_body_policy_names] + policies = {policy_name: Policy() for policy_name in policy_names} + attachment_registry = AttachmentRegistry() + attachment_registry.load_attachments( + [{"policy": policy_name, "scope": "*"} for policy_name in attachment_policy_names] + ) + + applied_policy_names, _ = _match_and_track_policies( + data={"metadata": {}}, + context=PolicyMatchContext(model="gpt-4"), + request_body_policies=request_body_policy_names, + policies_override=policies, + attachment_registry_override=attachment_registry, + ) + + assert applied_policy_names == policy_names + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_its_pipeline_also_steps(): from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry