mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
TestPolicyResolverInheritance
This commit is contained in:
Ishaan Jaffer
parent
cadbae5722
commit
b27c6e13fb
2 changed files with 127 additions and 419 deletions
|
|
@ -1,289 +1,113 @@
|
|||
"""
|
||||
Unit tests for ConditionEvaluator - tests AWS IAM-style condition evaluation.
|
||||
Unit tests for ConditionEvaluator - tests model condition evaluation.
|
||||
|
||||
Tests:
|
||||
- Condition operators (equals, in, prefix, not_equals, not_in)
|
||||
- PolicyCondition evaluation against request context
|
||||
- Metadata condition evaluation
|
||||
- Exact model match
|
||||
- Regex pattern match
|
||||
- List of models
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.policy_engine.condition_evaluator import ConditionEvaluator
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
ConditionOperator,
|
||||
PolicyCondition,
|
||||
PolicyMatchContext,
|
||||
)
|
||||
|
||||
|
||||
class TestConditionOperatorEvaluation:
|
||||
"""Test individual condition operator evaluation."""
|
||||
class TestConditionEvaluator:
|
||||
"""Test condition evaluation."""
|
||||
|
||||
def test_equals_operator_match(self):
|
||||
"""Test equals operator matches exact value."""
|
||||
operator = ConditionOperator(equals="gpt-4")
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-4") is True
|
||||
|
||||
def test_equals_operator_no_match(self):
|
||||
"""Test equals operator does not match different value."""
|
||||
operator = ConditionOperator(equals="gpt-4")
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-3.5") is False
|
||||
|
||||
def test_in_operator_match(self):
|
||||
"""Test in operator matches value in list."""
|
||||
operator = ConditionOperator(in_=["gpt-4", "gpt-4-turbo", "gpt-4o"])
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-4") is True
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-4-turbo") is True
|
||||
|
||||
def test_in_operator_no_match(self):
|
||||
"""Test in operator does not match value not in list."""
|
||||
operator = ConditionOperator(in_=["gpt-4", "gpt-4-turbo"])
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-3.5") is False
|
||||
|
||||
def test_prefix_operator_match(self):
|
||||
"""Test prefix operator matches value starting with prefix."""
|
||||
operator = ConditionOperator(prefix="bedrock/")
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "bedrock/claude-3") is True
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "bedrock/llama") is True
|
||||
|
||||
def test_prefix_operator_no_match(self):
|
||||
"""Test prefix operator does not match value not starting with prefix."""
|
||||
operator = ConditionOperator(prefix="bedrock/")
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "openai/gpt-4") is False
|
||||
|
||||
def test_not_equals_operator_match(self):
|
||||
"""Test not_equals operator matches when value is different."""
|
||||
operator = ConditionOperator(not_equals="gpt-3.5")
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-4") is True
|
||||
|
||||
def test_not_equals_operator_no_match(self):
|
||||
"""Test not_equals operator does not match when value is same."""
|
||||
operator = ConditionOperator(not_equals="gpt-4")
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-4") is False
|
||||
|
||||
def test_not_in_operator_match(self):
|
||||
"""Test not_in operator matches when value not in list."""
|
||||
operator = ConditionOperator(not_in=["gpt-3.5", "gpt-3.5-turbo"])
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-4") is True
|
||||
|
||||
def test_not_in_operator_no_match(self):
|
||||
"""Test not_in operator does not match when value in list."""
|
||||
operator = ConditionOperator(not_in=["gpt-4", "gpt-4-turbo"])
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "gpt-4") is False
|
||||
|
||||
def test_none_value_with_positive_operators(self):
|
||||
"""Test None value does not match positive operators."""
|
||||
assert ConditionEvaluator.evaluate_operator(
|
||||
ConditionOperator(equals="gpt-4"), None
|
||||
) is False
|
||||
assert ConditionEvaluator.evaluate_operator(
|
||||
ConditionOperator(in_=["gpt-4"]), None
|
||||
) is False
|
||||
assert ConditionEvaluator.evaluate_operator(
|
||||
ConditionOperator(prefix="gpt"), None
|
||||
) is False
|
||||
|
||||
def test_none_value_with_negative_operators(self):
|
||||
"""Test None value matches negative operators (None is not equal to anything)."""
|
||||
assert ConditionEvaluator.evaluate_operator(
|
||||
ConditionOperator(not_equals="gpt-4"), None
|
||||
) is True
|
||||
assert ConditionEvaluator.evaluate_operator(
|
||||
ConditionOperator(not_in=["gpt-4"]), None
|
||||
) is True
|
||||
|
||||
def test_empty_operator_matches_any(self):
|
||||
"""Test empty operator (no conditions) matches any value."""
|
||||
operator = ConditionOperator()
|
||||
assert ConditionEvaluator.evaluate_operator(operator, "anything") is True
|
||||
assert ConditionEvaluator.evaluate_operator(operator, None) is True
|
||||
|
||||
|
||||
class TestPolicyConditionEvaluation:
|
||||
"""Test PolicyCondition evaluation against request context."""
|
||||
|
||||
def test_model_condition_match(self):
|
||||
"""Test model condition matches."""
|
||||
condition = PolicyCondition(
|
||||
model=ConditionOperator(in_=["gpt-4", "gpt-4-turbo"])
|
||||
)
|
||||
def test_no_condition_always_matches(self):
|
||||
"""Test that None condition always matches."""
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
|
||||
assert ConditionEvaluator.evaluate(condition, context) is True
|
||||
|
||||
def test_model_condition_no_match(self):
|
||||
"""Test model condition does not match."""
|
||||
condition = PolicyCondition(
|
||||
model=ConditionOperator(in_=["gpt-4", "gpt-4-turbo"])
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-3.5")
|
||||
assert ConditionEvaluator.evaluate(condition, context) is False
|
||||
|
||||
def test_team_condition_match(self):
|
||||
"""Test team condition matches."""
|
||||
condition = PolicyCondition(
|
||||
team=ConditionOperator(prefix="healthcare-")
|
||||
)
|
||||
context = PolicyMatchContext(
|
||||
team_alias="healthcare-research", key_alias="key", model="gpt-4"
|
||||
)
|
||||
assert ConditionEvaluator.evaluate(condition, context) is True
|
||||
|
||||
def test_team_condition_no_match(self):
|
||||
"""Test team condition does not match."""
|
||||
condition = PolicyCondition(
|
||||
team=ConditionOperator(prefix="healthcare-")
|
||||
)
|
||||
context = PolicyMatchContext(
|
||||
team_alias="finance-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
assert ConditionEvaluator.evaluate(condition, context) is False
|
||||
|
||||
def test_key_condition_match(self):
|
||||
"""Test key condition matches."""
|
||||
condition = PolicyCondition(
|
||||
key=ConditionOperator(equals="production-key")
|
||||
)
|
||||
context = PolicyMatchContext(
|
||||
team_alias="team", key_alias="production-key", model="gpt-4"
|
||||
)
|
||||
assert ConditionEvaluator.evaluate(condition, context) is True
|
||||
|
||||
def test_multiple_conditions_all_match(self):
|
||||
"""Test multiple conditions all must match (AND logic)."""
|
||||
condition = PolicyCondition(
|
||||
model=ConditionOperator(in_=["gpt-4", "gpt-4-turbo"]),
|
||||
team=ConditionOperator(prefix="healthcare-"),
|
||||
)
|
||||
context = PolicyMatchContext(
|
||||
team_alias="healthcare-research", key_alias="key", model="gpt-4"
|
||||
)
|
||||
assert ConditionEvaluator.evaluate(condition, context) is True
|
||||
|
||||
def test_multiple_conditions_one_fails(self):
|
||||
"""Test multiple conditions - if one fails, all fails."""
|
||||
condition = PolicyCondition(
|
||||
model=ConditionOperator(in_=["gpt-4", "gpt-4-turbo"]),
|
||||
team=ConditionOperator(prefix="healthcare-"),
|
||||
)
|
||||
# Model matches but team doesn't
|
||||
context = PolicyMatchContext(
|
||||
team_alias="finance-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
assert ConditionEvaluator.evaluate(condition, context) is False
|
||||
|
||||
def test_none_condition_always_matches(self):
|
||||
"""Test None condition always matches."""
|
||||
context = PolicyMatchContext(team_alias="any", key_alias="any", model="any")
|
||||
assert ConditionEvaluator.evaluate(None, context) is True
|
||||
|
||||
|
||||
class TestMetadataConditionEvaluation:
|
||||
"""Test metadata condition evaluation."""
|
||||
|
||||
def test_metadata_condition_match(self):
|
||||
"""Test metadata condition matches."""
|
||||
condition = PolicyCondition(
|
||||
metadata={
|
||||
"environment": ConditionOperator(equals="production"),
|
||||
}
|
||||
)
|
||||
def test_exact_model_match(self):
|
||||
"""Test exact model string match."""
|
||||
condition = PolicyCondition(model="gpt-4")
|
||||
|
||||
# Match
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
|
||||
metadata = {"environment": "production"}
|
||||
assert ConditionEvaluator.evaluate(condition, context, metadata) is True
|
||||
assert ConditionEvaluator.evaluate(condition, context) is True
|
||||
|
||||
# No match
|
||||
context_other = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-3.5")
|
||||
assert ConditionEvaluator.evaluate(condition, context_other) is False
|
||||
|
||||
def test_metadata_condition_no_match(self):
|
||||
"""Test metadata condition does not match."""
|
||||
condition = PolicyCondition(
|
||||
metadata={
|
||||
"environment": ConditionOperator(equals="production"),
|
||||
}
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
|
||||
metadata = {"environment": "staging"}
|
||||
assert ConditionEvaluator.evaluate(condition, context, metadata) is False
|
||||
def test_regex_pattern_match(self):
|
||||
"""Test regex pattern matching."""
|
||||
condition = PolicyCondition(model="gpt-4.*")
|
||||
|
||||
# Matches
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4")
|
||||
) is True
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4-turbo")
|
||||
) is True
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4o")
|
||||
) is True
|
||||
|
||||
# No match
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-3.5")
|
||||
) is False
|
||||
|
||||
def test_metadata_condition_missing_field(self):
|
||||
"""Test metadata condition with missing field does not match."""
|
||||
condition = PolicyCondition(
|
||||
metadata={
|
||||
"environment": ConditionOperator(equals="production"),
|
||||
}
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
|
||||
metadata = {"other_field": "value"}
|
||||
assert ConditionEvaluator.evaluate(condition, context, metadata) is False
|
||||
def test_list_of_models_match(self):
|
||||
"""Test list of model values."""
|
||||
condition = PolicyCondition(model=["gpt-4", "gpt-4-turbo", "claude-3"])
|
||||
|
||||
# Matches
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4")
|
||||
) is True
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="claude-3")
|
||||
) is True
|
||||
|
||||
# No match
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-3.5")
|
||||
) is False
|
||||
|
||||
def test_metadata_condition_with_model_condition(self):
|
||||
"""Test combining metadata and model conditions."""
|
||||
condition = PolicyCondition(
|
||||
model=ConditionOperator(in_=["gpt-4"]),
|
||||
metadata={
|
||||
"environment": ConditionOperator(equals="production"),
|
||||
},
|
||||
)
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
|
||||
metadata = {"environment": "production"}
|
||||
assert ConditionEvaluator.evaluate(condition, context, metadata) is True
|
||||
def test_list_with_regex_patterns(self):
|
||||
"""Test list can contain regex patterns."""
|
||||
condition = PolicyCondition(model=["gpt-4.*", "claude-.*"])
|
||||
|
||||
# Matches
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4-turbo")
|
||||
) is True
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="claude-3")
|
||||
) is True
|
||||
|
||||
# No match
|
||||
assert ConditionEvaluator.evaluate(
|
||||
condition,
|
||||
PolicyMatchContext(team_alias="t", key_alias="k", model="llama-2")
|
||||
) is False
|
||||
|
||||
# Model matches but metadata doesn't
|
||||
metadata_staging = {"environment": "staging"}
|
||||
assert ConditionEvaluator.evaluate(condition, context, metadata_staging) is False
|
||||
def test_none_model_does_not_match(self):
|
||||
"""Test that None model value doesn't match conditions."""
|
||||
condition = PolicyCondition(model="gpt-4")
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model=None)
|
||||
assert ConditionEvaluator.evaluate(condition, context) is False
|
||||
|
||||
|
||||
class TestEvaluateAllConditions:
|
||||
"""Test evaluate_all_conditions helper."""
|
||||
|
||||
def test_all_conditions_match(self):
|
||||
"""Test all conditions match returns True."""
|
||||
conditions = [
|
||||
PolicyCondition(model=ConditionOperator(equals="gpt-4")),
|
||||
PolicyCondition(team=ConditionOperator(prefix="healthcare-")),
|
||||
]
|
||||
context = PolicyMatchContext(
|
||||
team_alias="healthcare-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
assert ConditionEvaluator.evaluate_all_conditions(conditions, context) is True
|
||||
|
||||
def test_one_condition_fails(self):
|
||||
"""Test one condition fails returns False."""
|
||||
conditions = [
|
||||
PolicyCondition(model=ConditionOperator(equals="gpt-4")),
|
||||
PolicyCondition(team=ConditionOperator(prefix="healthcare-")),
|
||||
]
|
||||
context = PolicyMatchContext(
|
||||
team_alias="finance-team", key_alias="key", model="gpt-4"
|
||||
)
|
||||
assert ConditionEvaluator.evaluate_all_conditions(conditions, context) is False
|
||||
|
||||
def test_empty_conditions_returns_true(self):
|
||||
"""Test empty conditions list returns True."""
|
||||
context = PolicyMatchContext(team_alias="any", key_alias="any", model="any")
|
||||
assert ConditionEvaluator.evaluate_all_conditions([], context) is True
|
||||
|
||||
|
||||
class TestEvaluateAnyCondition:
|
||||
"""Test evaluate_any_condition helper."""
|
||||
|
||||
def test_any_condition_matches(self):
|
||||
"""Test any condition matches returns True."""
|
||||
conditions = [
|
||||
PolicyCondition(model=ConditionOperator(equals="gpt-4")),
|
||||
PolicyCondition(model=ConditionOperator(equals="gpt-3.5")),
|
||||
]
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4")
|
||||
assert ConditionEvaluator.evaluate_any_condition(conditions, context) is True
|
||||
|
||||
def test_no_condition_matches(self):
|
||||
"""Test no condition matches returns False."""
|
||||
conditions = [
|
||||
PolicyCondition(model=ConditionOperator(equals="gpt-4")),
|
||||
PolicyCondition(model=ConditionOperator(equals="gpt-3.5")),
|
||||
]
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="key", model="claude-3")
|
||||
assert ConditionEvaluator.evaluate_any_condition(conditions, context) is False
|
||||
|
||||
def test_empty_conditions_returns_true(self):
|
||||
"""Test empty conditions list returns True."""
|
||||
context = PolicyMatchContext(team_alias="any", key_alias="any", model="any")
|
||||
assert ConditionEvaluator.evaluate_any_condition([], context) is True
|
||||
def test_empty_condition_always_matches(self):
|
||||
"""Test condition with no model field always matches."""
|
||||
condition = PolicyCondition() # No model specified
|
||||
context = PolicyMatchContext(team_alias="t", key_alias="k", model="any-model")
|
||||
assert ConditionEvaluator.evaluate(condition, context) is True
|
||||
|
|
|
|||
|
|
@ -4,19 +4,17 @@ Unit tests for PolicyResolver - tests guardrail resolution.
|
|||
Tests:
|
||||
- Inheritance chain resolution
|
||||
- Inheritance with add/remove
|
||||
- Conditional statements
|
||||
- Model conditions
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
ConditionOperator,
|
||||
Policy,
|
||||
PolicyCondition,
|
||||
PolicyGuardrails,
|
||||
PolicyMatchContext,
|
||||
PolicyStatement,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -103,144 +101,66 @@ class TestPolicyResolverInheritance:
|
|||
assert resolved.inheritance_chain == ["root", "middle", "leaf"]
|
||||
|
||||
|
||||
class TestPolicyResolverWithStatements:
|
||||
"""Test resolve_policy_guardrails with conditional statements."""
|
||||
class TestPolicyResolverWithConditions:
|
||||
"""Test resolve_policy_guardrails with model conditions."""
|
||||
|
||||
def test_statement_condition_matches(self):
|
||||
"""Test statement guardrails are added when condition matches."""
|
||||
def test_condition_matches(self):
|
||||
"""Test guardrails are added when condition matches."""
|
||||
policies = {
|
||||
"conditional-policy": Policy(
|
||||
guardrails=PolicyGuardrails(add=["base_guardrail"]),
|
||||
statements=[
|
||||
PolicyStatement(
|
||||
sid="GPT4Safety",
|
||||
guardrails=["toxicity_filter"],
|
||||
condition=PolicyCondition(
|
||||
model=ConditionOperator(in_=["gpt-4", "gpt-4-turbo"])
|
||||
),
|
||||
),
|
||||
],
|
||||
"gpt4-policy": Policy(
|
||||
guardrails=PolicyGuardrails(add=["toxicity_filter"]),
|
||||
condition=PolicyCondition(model="gpt-4.*"),
|
||||
),
|
||||
}
|
||||
|
||||
# GPT-4 should get both base and statement guardrails
|
||||
# GPT-4 should get guardrails
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="k", model="gpt-4")
|
||||
resolved = PolicyResolver.resolve_policy_guardrails(
|
||||
policy_name="conditional-policy",
|
||||
policy_name="gpt4-policy",
|
||||
policies=policies,
|
||||
context=context,
|
||||
)
|
||||
|
||||
assert "base_guardrail" in resolved.guardrails
|
||||
assert "toxicity_filter" in resolved.guardrails
|
||||
|
||||
def test_statement_condition_does_not_match(self):
|
||||
"""Test statement guardrails are NOT added when condition doesn't match."""
|
||||
def test_condition_does_not_match(self):
|
||||
"""Test guardrails are NOT added when condition doesn't match."""
|
||||
policies = {
|
||||
"conditional-policy": Policy(
|
||||
guardrails=PolicyGuardrails(add=["base_guardrail"]),
|
||||
statements=[
|
||||
PolicyStatement(
|
||||
sid="GPT4Safety",
|
||||
guardrails=["toxicity_filter"],
|
||||
condition=PolicyCondition(
|
||||
model=ConditionOperator(in_=["gpt-4", "gpt-4-turbo"])
|
||||
),
|
||||
),
|
||||
],
|
||||
"gpt4-policy": Policy(
|
||||
guardrails=PolicyGuardrails(add=["toxicity_filter"]),
|
||||
condition=PolicyCondition(model="gpt-4.*"),
|
||||
),
|
||||
}
|
||||
|
||||
# GPT-3.5 should only get base guardrails, not statement guardrails
|
||||
# GPT-3.5 should NOT get guardrails
|
||||
context = PolicyMatchContext(team_alias="team", key_alias="k", model="gpt-3.5")
|
||||
resolved = PolicyResolver.resolve_policy_guardrails(
|
||||
policy_name="conditional-policy",
|
||||
policy_name="gpt4-policy",
|
||||
policies=policies,
|
||||
context=context,
|
||||
)
|
||||
|
||||
assert "base_guardrail" in resolved.guardrails
|
||||
assert "toxicity_filter" not in resolved.guardrails
|
||||
|
||||
def test_multiple_statements_some_match(self):
|
||||
"""Test multiple statements where only some match."""
|
||||
def test_no_condition_always_applies(self):
|
||||
"""Test policy without condition always applies."""
|
||||
policies = {
|
||||
"multi-statement": Policy(
|
||||
guardrails=PolicyGuardrails(add=["base"]),
|
||||
statements=[
|
||||
PolicyStatement(
|
||||
sid="GPT4Only",
|
||||
guardrails=["gpt4_guardrail"],
|
||||
condition=PolicyCondition(
|
||||
model=ConditionOperator(equals="gpt-4")
|
||||
),
|
||||
),
|
||||
PolicyStatement(
|
||||
sid="HealthcareOnly",
|
||||
guardrails=["hipaa_audit"],
|
||||
condition=PolicyCondition(
|
||||
team=ConditionOperator(prefix="healthcare-")
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
}
|
||||
|
||||
# Healthcare team with GPT-4 should get all guardrails
|
||||
context = PolicyMatchContext(
|
||||
team_alias="healthcare-team", key_alias="k", model="gpt-4"
|
||||
)
|
||||
resolved = PolicyResolver.resolve_policy_guardrails(
|
||||
policy_name="multi-statement",
|
||||
policies=policies,
|
||||
context=context,
|
||||
)
|
||||
|
||||
assert "base" in resolved.guardrails
|
||||
assert "gpt4_guardrail" in resolved.guardrails
|
||||
assert "hipaa_audit" in resolved.guardrails
|
||||
|
||||
# Finance team with GPT-4 should only get base + gpt4_guardrail
|
||||
context_finance = PolicyMatchContext(
|
||||
team_alias="finance-team", key_alias="k", model="gpt-4"
|
||||
)
|
||||
resolved_finance = PolicyResolver.resolve_policy_guardrails(
|
||||
policy_name="multi-statement",
|
||||
policies=policies,
|
||||
context=context_finance,
|
||||
)
|
||||
|
||||
assert "base" in resolved_finance.guardrails
|
||||
assert "gpt4_guardrail" in resolved_finance.guardrails
|
||||
assert "hipaa_audit" not in resolved_finance.guardrails
|
||||
|
||||
def test_statement_with_no_condition_always_applies(self):
|
||||
"""Test statement with no condition always applies."""
|
||||
policies = {
|
||||
"always-policy": Policy(
|
||||
guardrails=PolicyGuardrails(add=["base"]),
|
||||
statements=[
|
||||
PolicyStatement(
|
||||
sid="AlwaysApply",
|
||||
guardrails=["always_guardrail"],
|
||||
condition=None, # No condition = always applies
|
||||
),
|
||||
],
|
||||
"global": Policy(
|
||||
guardrails=PolicyGuardrails(add=["pii_blocker"]),
|
||||
),
|
||||
}
|
||||
|
||||
context = PolicyMatchContext(team_alias="any", key_alias="any", model="any")
|
||||
resolved = PolicyResolver.resolve_policy_guardrails(
|
||||
policy_name="always-policy",
|
||||
policy_name="global",
|
||||
policies=policies,
|
||||
context=context,
|
||||
)
|
||||
|
||||
assert "base" in resolved.guardrails
|
||||
assert "always_guardrail" in resolved.guardrails
|
||||
assert "pii_blocker" in resolved.guardrails
|
||||
|
||||
def test_inheritance_with_statements(self):
|
||||
"""Test inheritance works with statements."""
|
||||
def test_inheritance_with_condition(self):
|
||||
"""Test inheritance works with conditions."""
|
||||
policies = {
|
||||
"base": Policy(
|
||||
guardrails=PolicyGuardrails(add=["pii_blocker"]),
|
||||
|
|
@ -248,62 +168,26 @@ class TestPolicyResolverWithStatements:
|
|||
"child": Policy(
|
||||
inherit="base",
|
||||
guardrails=PolicyGuardrails(add=["child_guardrail"]),
|
||||
statements=[
|
||||
PolicyStatement(
|
||||
sid="ConditionalStatement",
|
||||
guardrails=["conditional_guardrail"],
|
||||
condition=PolicyCondition(
|
||||
model=ConditionOperator(equals="gpt-4")
|
||||
),
|
||||
),
|
||||
],
|
||||
condition=PolicyCondition(model="gpt-4"),
|
||||
),
|
||||
}
|
||||
|
||||
context = PolicyMatchContext(team_alias="any-team", key_alias="k", model="gpt-4")
|
||||
resolved = PolicyResolver.resolve_policy_guardrails(
|
||||
# GPT-4 should get both base and child guardrails
|
||||
context_gpt4 = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-4")
|
||||
resolved_gpt4 = PolicyResolver.resolve_policy_guardrails(
|
||||
policy_name="child",
|
||||
policies=policies,
|
||||
context=context,
|
||||
context=context_gpt4,
|
||||
)
|
||||
assert "pii_blocker" in resolved_gpt4.guardrails
|
||||
assert "child_guardrail" in resolved_gpt4.guardrails
|
||||
|
||||
# Should have: inherited pii_blocker, child's child_guardrail, and conditional_guardrail
|
||||
assert "pii_blocker" in resolved.guardrails
|
||||
assert "child_guardrail" in resolved.guardrails
|
||||
assert "conditional_guardrail" in resolved.guardrails
|
||||
|
||||
def test_inheritance_with_remove_and_statements(self):
|
||||
"""Test inheritance with remove still works alongside statements."""
|
||||
policies = {
|
||||
"base": Policy(
|
||||
guardrails=PolicyGuardrails(add=["pii_blocker", "phi_blocker"]),
|
||||
),
|
||||
"child": Policy(
|
||||
inherit="base",
|
||||
guardrails=PolicyGuardrails(
|
||||
add=["child_guardrail"],
|
||||
remove=["phi_blocker"], # Remove phi_blocker from parent
|
||||
),
|
||||
statements=[
|
||||
PolicyStatement(
|
||||
sid="Conditional",
|
||||
guardrails=["conditional_guardrail"],
|
||||
condition=PolicyCondition(
|
||||
model=ConditionOperator(equals="gpt-4")
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
}
|
||||
|
||||
context = PolicyMatchContext(team_alias="any-team", key_alias="k", model="gpt-4")
|
||||
resolved = PolicyResolver.resolve_policy_guardrails(
|
||||
# GPT-3.5 should only get base guardrails (child condition doesn't match)
|
||||
context_gpt35 = PolicyMatchContext(team_alias="t", key_alias="k", model="gpt-3.5")
|
||||
resolved_gpt35 = PolicyResolver.resolve_policy_guardrails(
|
||||
policy_name="child",
|
||||
policies=policies,
|
||||
context=context,
|
||||
context=context_gpt35,
|
||||
)
|
||||
|
||||
assert "pii_blocker" in resolved.guardrails # Inherited
|
||||
assert "phi_blocker" not in resolved.guardrails # Removed
|
||||
assert "child_guardrail" in resolved.guardrails # Added by child
|
||||
assert "conditional_guardrail" in resolved.guardrails # From statement
|
||||
assert "pii_blocker" in resolved_gpt35.guardrails
|
||||
assert "child_guardrail" not in resolved_gpt35.guardrails
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue