mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
fix(policy_engine): resolve policies once and apply fallback semantics in get_matching_policies
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
99e14efb62
commit
e647255909
2 changed files with 60 additions and 9 deletions
|
|
@ -114,7 +114,7 @@ class PolicyMatcher:
|
|||
verbose_proxy_logger.debug("AttachmentRegistry not initialized, returning empty list")
|
||||
return []
|
||||
|
||||
return registry.get_attached_policies(context)
|
||||
return registry.get_attached_policies(context, PolicyMatcher.policy_applies(context))
|
||||
|
||||
@staticmethod
|
||||
def get_matching_policies_from_registry(
|
||||
|
|
@ -137,14 +137,22 @@ class PolicyMatcher:
|
|||
policies: dict[str, Policy] | None = None,
|
||||
) -> Callable[[str], bool]:
|
||||
"""Predicate telling whether a policy exists and its condition matches the context."""
|
||||
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
|
||||
return lambda policy_name: bool(
|
||||
PolicyMatcher.get_policies_with_matching_conditions(
|
||||
policy_names=(policy_name,),
|
||||
context=context,
|
||||
policies=policies,
|
||||
policies=resolved,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _registry_policies() -> dict[str, Policy]:
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
|
||||
registry: Final = get_policy_registry()
|
||||
return registry.get_all_policies() if registry.is_initialized() else {}
|
||||
|
||||
@staticmethod
|
||||
def get_policies_with_matching_conditions(
|
||||
policy_names: Sequence[str],
|
||||
|
|
@ -167,17 +175,12 @@ class PolicyMatcher:
|
|||
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: Final = get_policy_registry()
|
||||
if not registry.is_initialized():
|
||||
return []
|
||||
policies = registry.get_all_policies()
|
||||
resolved: Final = policies if policies is not None else PolicyMatcher._registry_policies()
|
||||
|
||||
matching_policies: Final = []
|
||||
for policy_name in policy_names:
|
||||
policy = policies.get(policy_name)
|
||||
policy = resolved.get(policy_name)
|
||||
if policy is None:
|
||||
continue
|
||||
# Policy matches if it has no condition OR condition evaluates to True
|
||||
|
|
|
|||
|
|
@ -8,8 +8,11 @@ Tests:
|
|||
|
||||
import pytest
|
||||
|
||||
import litellm.proxy.policy_engine.attachment_registry as attachment_registry_module
|
||||
import litellm.proxy.policy_engine.policy_registry as policy_registry_module
|
||||
from litellm.proxy.policy_engine.attachment_registry import AttachmentRegistry
|
||||
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
|
||||
from litellm.proxy.policy_engine.policy_registry import PolicyRegistry
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
PolicyMatchContext,
|
||||
PolicyScope,
|
||||
|
|
@ -196,3 +199,48 @@ class TestPolicyMatcherWithAttachments:
|
|||
attached = registry.get_attached_policies(context)
|
||||
|
||||
assert "healthcare-policy" not in attached
|
||||
|
||||
|
||||
def _global_registries(monkeypatch):
|
||||
policies = PolicyRegistry()
|
||||
policies.load_policies(
|
||||
{
|
||||
"guardrail-y": {"guardrails": {"add": ["y"]}},
|
||||
"guardrail-x": {"guardrails": {"add": ["x"]}, "condition": {"model": "claude.*"}},
|
||||
}
|
||||
)
|
||||
attachments = AttachmentRegistry()
|
||||
attachments.load_attachments(
|
||||
[
|
||||
{"policy": "guardrail-x", "tags": ["opt-in"]},
|
||||
{"policy": "guardrail-y", "scope": "*", "default": True},
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr(policy_registry_module, "get_policy_registry", lambda: policies)
|
||||
monkeypatch.setattr(attachment_registry_module, "get_attachment_registry", lambda: attachments)
|
||||
return policies
|
||||
|
||||
|
||||
class TestGetMatchingPoliciesFallback:
|
||||
def test_condition_failing_opt_in_falls_back_to_default(self, monkeypatch):
|
||||
_global_registries(monkeypatch)
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"])
|
||||
|
||||
assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-y"]
|
||||
|
||||
def test_condition_passing_opt_in_suppresses_default(self, monkeypatch):
|
||||
_global_registries(monkeypatch)
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="claude-haiku", tags=["opt-in"])
|
||||
|
||||
assert PolicyMatcher.get_matching_policies(context=context) == ["guardrail-x"]
|
||||
|
||||
def test_policy_applies_reads_registry_once(self, monkeypatch):
|
||||
policies = _global_registries(monkeypatch)
|
||||
calls = []
|
||||
original = policies.get_all_policies
|
||||
monkeypatch.setattr(policies, "get_all_policies", lambda: calls.append(1) or original())
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-5.5", tags=["opt-in"])
|
||||
|
||||
PolicyMatcher.get_matching_policies(context=context)
|
||||
|
||||
assert len(calls) == 1
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue