diff --git a/litellm/proxy/guardrails/guardrail_hooks/atr/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/atr/__init__.py index 6754876f0e7..33fea7d1c3d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/atr/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/atr/__init__.py @@ -17,6 +17,7 @@ def initialize_guardrail( _cb = ATRGuardrail( rules_path=litellm_params.rules_path, severity_threshold=litellm_params.severity_threshold, + include_tags=getattr(litellm_params, "include_tags", None), guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, default_on=litellm_params.default_on, diff --git a/litellm/proxy/guardrails/guardrail_hooks/atr/atr.py b/litellm/proxy/guardrails/guardrail_hooks/atr/atr.py index 2c245dec249..eba1814028b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/atr/atr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/atr/atr.py @@ -75,6 +75,7 @@ class ATRGuardrail(CustomGuardrail): self, rules_path: Optional[str] = None, severity_threshold: Optional[str] = None, + include_tags: Optional[List[str]] = None, **kwargs: Any, ) -> None: try: @@ -97,6 +98,7 @@ class ATRGuardrail(CustomGuardrail): f"Must be one of: {sorted(_SEVERITY_RANK)}" ) self.severity_threshold = threshold + self.include_tags: Optional[List[str]] = include_tags or None self.engine = ATREngine() resolved_path = rules_path or os.environ.get("ATR_RULES_PATH") @@ -288,14 +290,26 @@ class ATRGuardrail(CustomGuardrail): ) matches = self.engine.evaluate(event) threshold_rank = _SEVERITY_RANK[self.severity_threshold] - return [ - m - for m in matches - if _SEVERITY_RANK.get( - getattr(m, "severity", "low").lower(), len(_SEVERITY_RANK) - ) - <= threshold_rank - ] + + result = [] + for m in matches: + # include_tags filter: skip rules whose tags don't intersect the allow-list + if self.include_tags is not None: + tags = getattr(m, "tags", {}) or {} + tag_values: set = ( + set(tags.values()) if isinstance(tags, dict) else set() + ) + if not tag_values.intersection(self.include_tags): + continue + + # Treat None or unrecognised severity conservatively (rank 0 = critical) + raw_severity = getattr(m, "severity", None) + severity_str = (raw_severity or "").lower() if raw_severity is not None else "" + rank = _SEVERITY_RANK.get(severity_str, 0) + if rank <= threshold_rank: + result.append(m) + + return result @staticmethod def _summarize_match(match: Any) -> dict: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_atr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_atr.py index 7ec164dd108..3d52fbfa06f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_atr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_atr.py @@ -383,3 +383,82 @@ async def test_post_call_scans_text_completion_response(fake_pyatr, tmp_path): assert excinfo.value.status_code == 400 assert excinfo.value.detail["matched_rules"][0]["rule_id"] == "ATR-400" + + +def test_scan_include_tags_filters_rules(fake_pyatr, tmp_path): + """include_tags restricts scanning to rules with matching tag values.""" + _, engine = fake_pyatr + ATRGuardrail, _, _ = _import_guardrail() + + rules_dir = tmp_path / "rules" + rules_dir.mkdir() + + injection_match = MagicMock( + rule_id="ATR-500", + title="Injection", + severity="high", + tags={"category": "prompt_injection"}, + ) + exfil_match = MagicMock( + rule_id="ATR-501", + title="Exfil", + severity="high", + tags={"category": "context_exfiltration"}, + ) + engine.evaluate.return_value = [injection_match, exfil_match] + + guard = ATRGuardrail( + rules_path=str(rules_dir), + severity_threshold="high", + include_tags=["prompt_injection"], + guardrail_name="atr-test", + ) + + matches = guard._scan("hello world", event_type="llm_input") + rule_ids = [m.rule_id for m in matches] + assert rule_ids == ["ATR-500"] + assert "ATR-501" not in rule_ids + + +def test_scan_none_severity_treated_conservatively(fake_pyatr, tmp_path): + """A match with severity=None is treated as critical (always included).""" + _, engine = fake_pyatr + ATRGuardrail, _, _ = _import_guardrail() + + rules_dir = tmp_path / "rules" + rules_dir.mkdir() + + none_severity_match = MagicMock(rule_id="ATR-600", title="Unknown sev", severity=None) + engine.evaluate.return_value = [none_severity_match] + + guard = ATRGuardrail( + rules_path=str(rules_dir), + severity_threshold="low", + guardrail_name="atr-test", + ) + + matches = guard._scan("some content", event_type="llm_input") + assert len(matches) == 1 + assert matches[0].rule_id == "ATR-600" + + +def test_scan_unknown_severity_treated_conservatively(fake_pyatr, tmp_path): + """A match with an unrecognised severity string is treated as critical.""" + _, engine = fake_pyatr + ATRGuardrail, _, _ = _import_guardrail() + + rules_dir = tmp_path / "rules" + rules_dir.mkdir() + + unknown_match = MagicMock(rule_id="ATR-601", title="Future sev", severity="informational") + engine.evaluate.return_value = [unknown_match] + + guard = ATRGuardrail( + rules_path=str(rules_dir), + severity_threshold="low", + guardrail_name="atr-test", + ) + + matches = guard._scan("some content", event_type="llm_input") + assert len(matches) == 1 + assert matches[0].rule_id == "ATR-601"