mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(policy_engine): run global policy pipelines before scoped ones (#39697)
* fix(policy_engine): run global policy pipelines before scoped ones Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(policy_engine): rank duplicate attachments by broadest scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(policy_engine): rank combined-scope attachments below single-scope ones Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: shivam <shivam@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
d51a7af655
commit
a426dc43cb
4 changed files with 141 additions and 35 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue