mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-19 08:18:15 +00:00
Ishaan JaffGitHubClaude Opus 4.6greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>* init schema with TAGS * ui: add policy test * resolvePoliciesCall * add_policy_sources_to_metadata + headers * types Policy * preview Impact * def _describe_match_reason( * match based on TAGs * TestTagBasedAttachments * test fixes * add policy_resolve_router * add_guardrails_from_policy_engine * TestMatchAttribution * refactor * fix * fix: address Greptile review feedback on policy resolve endpoints - Track unnamed keys/teams as separate counts instead of inflating affected_keys_count with duplicate "(unnamed key)" placeholders. Added unnamed_keys_count and unnamed_teams_count to response. - Push alias pattern matching to DB via _build_alias_where() which converts exact patterns to Prisma "in" and suffix wildcards to "startsWith" filters. - Gate sync_policies_from_db/sync_attachments_from_db behind force_sync query param (default false) to avoid 2 DB round-trips on every /policies/resolve request. - Remove worktree-only conftest.py that cleared sys.modules at import time — no longer needed since code moved to main repo. - Rename MAX_ESTIMATE_IMPACT_ROWS → MAX_POLICY_ESTIMATE_IMPACT_ROWS. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * fix: eliminate duplicate DB queries and fix header delimiter ambiguity - Fetch teams table once in estimate_attachment_impact and reuse for both tag-based and alias-based lookups (was querying teams twice when both tag_patterns and team_patterns were provided). - Convert tag/team filter functions from async DB queries to sync filters that operate on pre-fetched data (_filter_keys_by_tags, _filter_teams_by_tags). - Fix comma ambiguity in x-litellm-policy-sources header: use '; ' as entry delimiter since matched_via values can contain commas. - Use '+' as the within-value separator in matched_via reason strings (e.g. "tag:healthcare+team:health-team") to avoid conflict with header delimiters. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> * Update litellm/proxy/policy_engine/policy_resolve_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
336 lines
13 KiB
Python
336 lines
13 KiB
Python
"""
|
|
Unit tests for AttachmentRegistry - tests policy attachment matching.
|
|
|
|
Tests the main entry point: get_attached_policies()
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from litellm.proxy.policy_engine.attachment_registry import (
|
|
AttachmentRegistry,
|
|
get_attachment_registry,
|
|
)
|
|
from litellm.types.proxy.policy_engine import PolicyMatchContext
|
|
|
|
|
|
class TestGetAttachedPolicies:
|
|
"""Test get_attached_policies - the main entry point."""
|
|
|
|
def test_global_scope_matches_all_requests(self):
|
|
"""Test global scope (*) matches any request context."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "global-baseline", "scope": "*"},
|
|
])
|
|
|
|
# Should match any context
|
|
context = PolicyMatchContext(
|
|
team_alias="any-team", key_alias="any-key", model="any-model"
|
|
)
|
|
attached = registry.get_attached_policies(context)
|
|
assert "global-baseline" in attached
|
|
|
|
def test_team_specific_attachment(self):
|
|
"""Test team-specific attachment matches only that team."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "healthcare-policy", "teams": ["healthcare-team"]},
|
|
])
|
|
|
|
# Match
|
|
context = PolicyMatchContext(
|
|
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
|
)
|
|
assert "healthcare-policy" in registry.get_attached_policies(context)
|
|
|
|
# No match - different team
|
|
context_other = PolicyMatchContext(
|
|
team_alias="finance-team", key_alias="key", model="gpt-4"
|
|
)
|
|
assert "healthcare-policy" not in registry.get_attached_policies(context_other)
|
|
|
|
def test_key_wildcard_pattern_attachment(self):
|
|
"""Test key pattern attachment with wildcard."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "dev-policy", "keys": ["dev-key-*"]},
|
|
])
|
|
|
|
# Match - key starts with dev-key-
|
|
context = PolicyMatchContext(
|
|
team_alias="team", key_alias="dev-key-123", model="gpt-4"
|
|
)
|
|
assert "dev-policy" in registry.get_attached_policies(context)
|
|
|
|
# No match - different prefix
|
|
context_prod = PolicyMatchContext(
|
|
team_alias="team", key_alias="prod-key-123", model="gpt-4"
|
|
)
|
|
assert "dev-policy" not in registry.get_attached_policies(context_prod)
|
|
|
|
def test_model_specific_attachment(self):
|
|
"""Test model-specific attachment."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "gpt4-policy", "models": ["gpt-4", "gpt-4-turbo"]},
|
|
])
|
|
|
|
# Match
|
|
context = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="gpt-4"
|
|
)
|
|
assert "gpt4-policy" in registry.get_attached_policies(context)
|
|
|
|
# No match
|
|
context_other = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="gpt-3.5"
|
|
)
|
|
assert "gpt4-policy" not in registry.get_attached_policies(context_other)
|
|
|
|
def test_model_wildcard_pattern(self):
|
|
"""Test model wildcard pattern like bedrock/*."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "bedrock-policy", "models": ["bedrock/*"]},
|
|
])
|
|
|
|
# Match
|
|
context = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="bedrock/claude-3"
|
|
)
|
|
assert "bedrock-policy" in registry.get_attached_policies(context)
|
|
|
|
# No match
|
|
context_other = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="openai/gpt-4"
|
|
)
|
|
assert "bedrock-policy" not in registry.get_attached_policies(context_other)
|
|
|
|
def test_multiple_attachments_match_same_context(self):
|
|
"""Test multiple attachments can match the same context."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "global-baseline", "scope": "*"},
|
|
{"policy": "healthcare-policy", "teams": ["healthcare-team"]},
|
|
{"policy": "gpt4-policy", "models": ["gpt-4"]},
|
|
])
|
|
|
|
context = PolicyMatchContext(
|
|
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
|
)
|
|
attached = registry.get_attached_policies(context)
|
|
|
|
# All three should match
|
|
assert "global-baseline" in attached
|
|
assert "healthcare-policy" in attached
|
|
assert "gpt4-policy" in attached
|
|
assert len(attached) == 3
|
|
|
|
def test_same_policy_multiple_attachments_no_duplicates(self):
|
|
"""Test same policy attached multiple ways doesn't duplicate."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "multi-policy", "scope": "*"},
|
|
{"policy": "multi-policy", "teams": ["healthcare-team"]},
|
|
])
|
|
|
|
context = PolicyMatchContext(
|
|
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
|
)
|
|
attached = registry.get_attached_policies(context)
|
|
|
|
# Should only appear once
|
|
assert attached.count("multi-policy") == 1
|
|
|
|
def test_no_attachments_returns_empty(self):
|
|
"""Test empty attachments returns empty list."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([])
|
|
|
|
context = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="gpt-4"
|
|
)
|
|
attached = registry.get_attached_policies(context)
|
|
assert attached == []
|
|
|
|
def test_no_matching_attachments_returns_empty(self):
|
|
"""Test no matching attachments returns empty list."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "healthcare-policy", "teams": ["healthcare-team"]},
|
|
])
|
|
|
|
context = PolicyMatchContext(
|
|
team_alias="finance-team", key_alias="key", model="gpt-4"
|
|
)
|
|
attached = registry.get_attached_policies(context)
|
|
assert attached == []
|
|
|
|
def test_combined_team_and_model_attachment(self):
|
|
"""Test attachment with both team and model constraints."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "strict-policy", "teams": ["healthcare-team"], "models": ["gpt-4"]},
|
|
])
|
|
|
|
# Match - both team and model match
|
|
context = PolicyMatchContext(
|
|
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
|
)
|
|
assert "strict-policy" in registry.get_attached_policies(context)
|
|
|
|
# No match - team matches but model doesn't
|
|
context_wrong_model = PolicyMatchContext(
|
|
team_alias="healthcare-team", key_alias="key", model="gpt-3.5"
|
|
)
|
|
assert "strict-policy" not in registry.get_attached_policies(context_wrong_model)
|
|
|
|
# No match - model matches but team doesn't
|
|
context_wrong_team = PolicyMatchContext(
|
|
team_alias="finance-team", key_alias="key", model="gpt-4"
|
|
)
|
|
assert "strict-policy" not in registry.get_attached_policies(context_wrong_team)
|
|
|
|
|
|
class TestTagBasedAttachments:
|
|
"""Test tag-based policy attachment matching."""
|
|
|
|
def test_tag_matching_and_wildcards(self):
|
|
"""Test tag matching: exact match, wildcard match, and no-match cases."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "hipaa-policy", "tags": ["healthcare"]},
|
|
{"policy": "health-policy", "tags": ["health-*"]},
|
|
])
|
|
|
|
# Exact tag match
|
|
context = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="gpt-4",
|
|
tags=["healthcare"],
|
|
)
|
|
attached = registry.get_attached_policies(context)
|
|
assert "hipaa-policy" in attached
|
|
assert "health-policy" not in attached # "healthcare" doesn't match "health-*"
|
|
|
|
# Wildcard tag match
|
|
context_wildcard = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="gpt-4",
|
|
tags=["health-prod"],
|
|
)
|
|
attached_wildcard = registry.get_attached_policies(context_wildcard)
|
|
assert "health-policy" in attached_wildcard
|
|
assert "hipaa-policy" not in attached_wildcard
|
|
|
|
# No match — wrong tag
|
|
context_no_match = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="gpt-4",
|
|
tags=["finance"],
|
|
)
|
|
assert registry.get_attached_policies(context_no_match) == []
|
|
|
|
# No match — no tags on context
|
|
context_no_tags = PolicyMatchContext(
|
|
team_alias="team", key_alias="key", model="gpt-4",
|
|
tags=None,
|
|
)
|
|
assert registry.get_attached_policies(context_no_tags) == []
|
|
|
|
def test_tag_combined_with_team(self):
|
|
"""Test attachment with both tags and teams requires BOTH to match (AND logic)."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "strict-policy", "teams": ["team-a"], "tags": ["healthcare"]},
|
|
])
|
|
|
|
# Match — both team and tag match
|
|
context = PolicyMatchContext(
|
|
team_alias="team-a", key_alias="key", model="gpt-4",
|
|
tags=["healthcare"],
|
|
)
|
|
assert "strict-policy" in registry.get_attached_policies(context)
|
|
|
|
# No match — tag matches but team doesn't
|
|
context_wrong_team = PolicyMatchContext(
|
|
team_alias="team-b", key_alias="key", model="gpt-4",
|
|
tags=["healthcare"],
|
|
)
|
|
assert "strict-policy" not in registry.get_attached_policies(context_wrong_team)
|
|
|
|
# No match — team matches but tag doesn't
|
|
context_wrong_tag = PolicyMatchContext(
|
|
team_alias="team-a", key_alias="key", model="gpt-4",
|
|
tags=["finance"],
|
|
)
|
|
assert "strict-policy" not in registry.get_attached_policies(context_wrong_tag)
|
|
|
|
|
|
class TestMatchAttribution:
|
|
"""Test get_attached_policies_with_reasons — the attribution logic that
|
|
powers response headers and the Policy Simulator UI."""
|
|
|
|
def test_reasons_for_global_tag_team_attachments(self):
|
|
"""Test that match reasons correctly describe WHY each policy matched."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "global-baseline", "scope": "*"},
|
|
{"policy": "hipaa-policy", "tags": ["healthcare"]},
|
|
{"policy": "team-policy", "teams": ["health-team"]},
|
|
])
|
|
|
|
context = PolicyMatchContext(
|
|
team_alias="health-team", key_alias="key", model="gpt-4",
|
|
tags=["healthcare"],
|
|
)
|
|
results = registry.get_attached_policies_with_reasons(context)
|
|
reasons = {r["policy_name"]: r["matched_via"] for r in results}
|
|
|
|
assert reasons["global-baseline"] == "scope:*"
|
|
assert "tag:healthcare" in reasons["hipaa-policy"]
|
|
assert "team:health-team" in reasons["team-policy"]
|
|
|
|
def test_tags_only_attachment_matches_any_team_key_model(self):
|
|
"""Test the primary use case: tags-only attachment with no team/key/model
|
|
constraint matches any request that carries the tag."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "hipaa-guardrails", "tags": ["healthcare"]},
|
|
])
|
|
|
|
# Should match regardless of team/key/model
|
|
context = PolicyMatchContext(
|
|
team_alias="random-team", key_alias="random-key", model="claude-3",
|
|
tags=["healthcare"],
|
|
)
|
|
attached = registry.get_attached_policies(context)
|
|
assert "hipaa-guardrails" in attached
|
|
|
|
# Should not match without the tag
|
|
context_no_tag = PolicyMatchContext(
|
|
team_alias="random-team", key_alias="random-key", model="claude-3",
|
|
)
|
|
assert registry.get_attached_policies(context_no_tag) == []
|
|
|
|
def test_attachment_with_no_scope_matches_everything(self):
|
|
"""Test that an attachment with no scope/teams/keys/models/tags
|
|
matches everything because teams/keys/models default to ['*']."""
|
|
registry = AttachmentRegistry()
|
|
registry.load_attachments([
|
|
{"policy": "catch-all"},
|
|
])
|
|
|
|
context = PolicyMatchContext(
|
|
team_alias="any-team", key_alias="any-key", model="gpt-4",
|
|
)
|
|
attached = registry.get_attached_policies(context)
|
|
assert "catch-all" in attached
|
|
|
|
|
|
class TestAttachmentRegistrySingleton:
|
|
"""Test global singleton behavior."""
|
|
|
|
def test_get_attachment_registry_returns_same_instance(self):
|
|
"""Test get_attachment_registry returns same instance."""
|
|
registry1 = get_attachment_registry()
|
|
registry2 = get_attachment_registry()
|
|
assert registry1 is registry2
|