diff --git a/tests/test_litellm/proxy/policy_engine/test_condition_evaluator.py b/tests/test_litellm/proxy/policy_engine/test_condition_evaluator.py index 8d335691d02..292f6e8f7da 100644 --- a/tests/test_litellm/proxy/policy_engine/test_condition_evaluator.py +++ b/tests/test_litellm/proxy/policy_engine/test_condition_evaluator.py @@ -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 diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_resolver.py b/tests/test_litellm/proxy/policy_engine/test_policy_resolver.py index bf1e7455def..9d672e018af 100644 --- a/tests/test_litellm/proxy/policy_engine/test_policy_resolver.py +++ b/tests/test_litellm/proxy/policy_engine/test_policy_resolver.py @@ -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