diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index cb1b223fb47..cb957dec57f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -77,9 +77,17 @@ class lakeraAI_Moderation(CustomGuardrail): return flagged: Final = _results[0].get("flagged", False) - category_scores: Final[dict | None] = _results[0].get("category_scores", None) + if flagged is True: + raise HTTPException( + status_code=400, + detail={ + "error": "Violated content safety policy", + "lakera_ai_response": response, + }, + ) if self.category_thresholds is not None: + category_scores: Final[dict | None] = _results[0].get("category_scores", None) if category_scores is not None: typed_cat_scores: Final = LakeraCategoryThresholds(**category_scores) if "jailbreak" in typed_cat_scores and "jailbreak" in self.category_thresholds: @@ -101,14 +109,6 @@ class lakeraAI_Moderation(CustomGuardrail): "lakera_ai_response": response, }, ) - elif flagged is True: - raise HTTPException( - status_code=400, - detail={ - "error": "Violated content safety policy", - "lakera_ai_response": response, - }, - ) return diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai.py new file mode 100644 index 00000000000..e2a0d1065ff --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -0,0 +1,55 @@ +""" +Tests for the Lakera AI v1 guardrail hook (pre-call moderation via /moderations). +""" + +import pytest +from fastapi import HTTPException + +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation + +THRESHOLDS = {"jailbreak": 0.9, "prompt_injection": 0.9} + + +def flagged_response(flagged: bool, category_scores: dict | None = None) -> dict: + result = {"flagged": flagged} + if category_scores is not None: + result["category_scores"] = category_scores + return {"results": [result]} + + +def test_flagged_true_blocks_when_category_thresholds_configured(): + guardrail = lakeraAI_Moderation(category_thresholds=THRESHOLDS) + response = flagged_response( + flagged=True, + category_scores={"jailbreak": 0.1, "prompt_injection": 0.1, "hate": 0.99}, + ) + with pytest.raises(HTTPException) as exc_info: + guardrail._check_response_flagged(response) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "Violated content safety policy" + + +def test_threshold_violation_blocks_when_not_flagged(): + guardrail = lakeraAI_Moderation(category_thresholds=THRESHOLDS) + response = flagged_response( + flagged=False, + category_scores={"jailbreak": 0.95, "prompt_injection": 0.1}, + ) + with pytest.raises(HTTPException) as exc_info: + guardrail._check_response_flagged(response) + assert exc_info.value.detail["error"] == "Violated jailbreak threshold" + + +def test_clean_response_passes_when_thresholds_configured(): + guardrail = lakeraAI_Moderation(category_thresholds=THRESHOLDS) + response = flagged_response( + flagged=False, + category_scores={"jailbreak": 0.1, "prompt_injection": 0.1}, + ) + assert guardrail._check_response_flagged(response) is None + + +def test_flagged_true_blocks_without_thresholds(): + guardrail = lakeraAI_Moderation() + with pytest.raises(HTTPException): + guardrail._check_response_flagged(flagged_response(flagged=True))