fix: update gemini's safety_settings (#1057)

Support gemini's new safety settings
instructor-ai · Oct 12, 2024 · 8b8f48a · 8b8f48a
1 parent 8ee9c5d
commit 8b8f48a
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/instructor/utils.py b/instructor/utils.py
@@ -347,12 +347,20 @@ def update_gemini_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
     # minimize gemini safety related errors - model is highly prone to false alarms
     from google.generativeai.types import HarmCategory, HarmBlockThreshold  # type: ignore
 
-    kwargs["safety_settings"] = kwargs.get("safety_settings", {}) | {
+    kwargs["safety_settings"] = kwargs.get("safety_settings", {})
+
+    fallback_safety_settings = {
         HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_ONLY_HIGH,
         HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
         HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_ONLY_HIGH,
     }
 
+    # Update or add fallback settings, respecting stricter existing ones
+    for category, fallback_threshold in fallback_safety_settings.items():
+        current_threshold = kwargs["safety_settings"].get(category)
+        if current_threshold is None or current_threshold < fallback_threshold:
+            kwargs["safety_settings"][category] = fallback_threshold
+
     return kwargs