diff --git a/docs/my-website/docs/providers/gemini.md b/docs/my-website/docs/providers/gemini.md index 9d5eb298b8e..44d744866cb 100644 --- a/docs/my-website/docs/providers/gemini.md +++ b/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 diff --git a/litellm/llms/gemini.py b/litellm/llms/gemini.py index 7e98345b375..2db27aebaf2 100644 --- a/litellm/llms/gemini.py +++ b/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