remove safety_settings for groq

SmartManoj · Sep 18, 2024 · 8fd630b · 8fd630b
1 parent 609b9f6
commit 8fd630b
Showing 1 changed file with 26 additions and 23 deletions.
diff --git a/openhands/llm/llm.py b/openhands/llm/llm.py
@@ -144,28 +144,31 @@ def __init__(
                 logger.info(f'Setting OllamaConfig.num_ctx to {total}')
 
         # This only seems to work with Google as the provider, not with OpenRouter!
-        gemini_safety_settings = (
-            [
-                {
-                    'category': 'HARM_CATEGORY_HARASSMENT',
-                    'threshold': 'BLOCK_NONE',
-                },
-                {
-                    'category': 'HARM_CATEGORY_HATE_SPEECH',
-                    'threshold': 'BLOCK_NONE',
-                },
-                {
-                    'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
-                    'threshold': 'BLOCK_NONE',
-                },
-                {
-                    'category': 'HARM_CATEGORY_DANGEROUS_CONTENT',
-                    'threshold': 'BLOCK_NONE',
-                },
-            ]
-            if self.config.model.lower().startswith('gemini')
-            else None
-        )
+        gemini_safety_settings = [
+            {
+                'category': 'HARM_CATEGORY_HARASSMENT',
+                'threshold': 'BLOCK_NONE',
+            },
+            {
+                'category': 'HARM_CATEGORY_HATE_SPEECH',
+                'threshold': 'BLOCK_NONE',
+            },
+            {
+                'category': 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+                'threshold': 'BLOCK_NONE',
+            },
+            {
+                'category': 'HARM_CATEGORY_DANGEROUS_CONTENT',
+                'threshold': 'BLOCK_NONE',
+            },
+        ]
+
+        if self.config.model.lower().startswith('gemini'):
+            kwargs = {
+                'safety_settings': gemini_safety_settings,
+            }
+        else:
+            kwargs = {}
 
         self.completion_unwrapped = partial(
             litellm_completion,
@@ -179,7 +182,7 @@ def __init__(
             temperature=self.config.temperature,
             top_p=self.config.top_p,
             caching=self.config.enable_cache,
-            safety_settings=gemini_safety_settings,
+            **kwargs,
         )
 
         if self.vision_is_active():