tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-07 02:25:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add safety_settings parameter to gemini generate_content calls

Browse Source
This commit is contained in:
Andres Barbaro
2024-02-17 19:13:04 -08:00
committed by ishaan-jaff
parent 03da603f1f
commit 0400e95433
1 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/gemini.py
+9
View File
@@ -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