Merge pull request #2013 from afbarbaro/litellm_gemini_safety_settings

Add safety_settings parameter to gemini generate_content calls

docs/my-website/docs/providers/gemini.md

@@ -16,6 +16,34 @@ response = completion(
 )
 ```
 
+## Specifying Safety Settings 
+In certain use-cases you may need to make calls to the models and pass [safety settigns](https://ai.google.dev/docs/safety_setting_gemini) different from the defaults. To do so, simple pass the `safety_settings` argument to `completion` or `acompletion`. For example:
+
+```python
+response = completion(
+    model="gemini/gemini-pro", 
+    messages=[{"role": "user", "content": "write code for saying hi from LiteLLM"}]
+    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",
+        },
+    ]
+)
+```
+
 # Gemini-Pro-Vision
 LiteLLM Supports the following image types passed in `url`
 - Images with direct links - https://storage.googleapis.com/github-repo/img/gemini/intro/landmark3.jpg

litellm/llms/gemini.py

@@ -121,6 +121,13 @@ def completion(
     ## Load Config
     inference_params = copy.deepcopy(optional_params)
     stream = inference_params.pop("stream", None)
+
+    # Handle safety settings
+    safety_settings_param = inference_params.pop("safety_settings", None)
+    safety_settings = None
+    if safety_settings_param:
+        safety_settings = [genai.types.SafetySettingDict(x) for x in safety_settings_param]
+
     config = litellm.GeminiConfig.get_config()
     for k, v in config.items():
         if (
@@ -141,11 +148,13 @@ def completion(
             response = _model.generate_content(
                 contents=prompt,
                 generation_config=genai.types.GenerationConfig(**inference_params),
+                safety_settings=safety_settings,
             )
         else:
             response = _model.generate_content(
                 contents=prompt,
                 generation_config=genai.types.GenerationConfig(**inference_params),
+                safety_settings=safety_settings,
                 stream=True,
             )
             return response