From 58aa4e9e163ee57fe35d8bb0678725718f254c53 Mon Sep 17 00:00:00 2001
From: Varshith Puli <pulivarshit@gmail.com>
Date: Thu, 24 Sep 2026 15:58:17 +0000
Subject: [PATCH] fix(lakera_ai): keep flagged verdict as blocking floor with
 category_thresholds

Previously, configuring category_thresholds moved the top-level flagged
verdict into an elif branch that never ran, so a Lakera flagged=True on a
category without a configured threshold passed unblocked. Setting
thresholds to be stricter could silently weaken blocking. The flagged
verdict is now checked unconditionally, and category thresholds remain
as purely additive tightening.

Fixes BerriAI/litellm#30727.

Signed-off-by: Varshith Puli <pulivarshit@gmail.com>
---
 .../guardrails/guardrail_hooks/lakera_ai.py   | 18 +++---
 .../guardrail_hooks/test_lakera_ai.py         | 55 +++++++++++++++++++
 2 files changed, 64 insertions(+), 9 deletions(-)
 create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_lakera_ai.py

diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py
index c69f90282c3..aac18203a25 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))