feat: Safety settings addition(only gemini for now)

src/rotator_library/client.py

@@ -104,7 +104,27 @@ class RotatingClient:
                         lib_logger.info(f"Attempting call with key ...{current_key[-4:]} (Attempt {attempt + 1}/{self.max_retries})")
                         
                         litellm_kwargs = self.all_providers.get_provider_kwargs(**kwargs.copy())
-                        
+
+                        if provider in self._provider_instances:
+                            provider_instance = self._provider_instances[provider]
+                            
+                            # Ensure safety_settings are present, defaulting to lowest if not provided
+                            if "safety_settings" not in litellm_kwargs:
+                                litellm_kwargs["safety_settings"] = {
+                                    "harassment": "BLOCK_NONE",
+                                    "hate_speech": "BLOCK_NONE",
+                                    "sexually_explicit": "BLOCK_NONE",
+                                    "dangerous_content": "BLOCK_NONE",
+                                }
+
+                            converted_settings = provider_instance.convert_safety_settings(litellm_kwargs["safety_settings"])
+                            
+                            if converted_settings is not None:
+                                litellm_kwargs["safety_settings"] = converted_settings
+                            else:
+                                # If conversion returns None, remove it to avoid sending empty settings
+                                del litellm_kwargs["safety_settings"]
+
                         if "gemma-3" in model and "messages" in litellm_kwargs:
                             new_messages = [
                                 {"role": "user", "content": m["content"]} if m.get("role") == "system" else m

src/rotator_library/providers/gemini_provider.py

@@ -1,6 +1,6 @@
 import httpx
 import logging
-from typing import List
+from typing import List, Dict, Any
 from .provider_interface import ProviderInterface
 
 lib_logger = logging.getLogger('rotator_library')
@@ -26,3 +26,27 @@ class GeminiProvider(ProviderInterface):
         except httpx.RequestError as e:
             lib_logger.error(f"Failed to fetch Gemini models: {e}")
             return []
+
+    def convert_safety_settings(self, settings: Dict[str, str]) -> List[Dict[str, Any]]:
+        """
+        Converts generic safety settings to the Gemini-specific format.
+        """
+        if not settings:
+            return []
+
+        gemini_settings = []
+        category_map = {
+            "harassment": "HARM_CATEGORY_HARASSMENT",
+            "hate_speech": "HARM_CATEGORY_HATE_SPEECH",
+            "sexually_explicit": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+            "dangerous_content": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        }
+
+        for generic_category, threshold in settings.items():
+            if generic_category in category_map:
+                gemini_settings.append({
+                    "category": category_map[generic_category],
+                    "threshold": threshold.upper()
+                })
+        
+        return gemini_settings

src/rotator_library/providers/provider_interface.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Dict, Any
 import httpx
 
 class ProviderInterface(ABC):
@@ -21,3 +21,15 @@ class ProviderInterface(ABC):
             A list of model name strings.
         """
         pass
+
+    def convert_safety_settings(self, settings: Dict[str, str]) -> List[Dict[str, Any]]:
+        """
+        Converts a generic safety settings dictionary to the provider-specific format.
+        
+        Args:
+            settings: A dictionary with generic harm categories and thresholds.
+            
+        Returns:
+            A list of provider-specific safety setting objects or None.
+        """
+        return None