fix: Content Safety not found throw error

microsoft · Feb 22, 2024 · 6ac0810 · 6ac0810
1 parent 737566f
commit 6ac0810
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/helpers/llm_worker.py b/helpers/llm_worker.py
@@ -234,7 +234,7 @@ def _contentsafety_category_test(
     if score == 0:
         return True  # No need to check severity
 
-    detection = next(item for item in res if item.category == category)
+    detection = next((item for item in res if item.category == category), None)
 
     if detection and detection.severity and detection.severity > score:
         _logger.debug(f"Matched {category} with severity {detection.severity}")