feat(utils.py): support global vertex ai safety settings param

This commit is contained in:
Krrish Dholakia 2024-08-21 17:36:59 -07:00
parent 9078f075f9
commit 3c99ad19fa
3 changed files with 44 additions and 7 deletions

View file

@ -3145,7 +3145,6 @@ def get_optional_params(
or model in litellm.vertex_embedding_models
or model in litellm.vertex_vision_models
):
print_verbose(f"(start) INSIDE THE VERTEX AI OPTIONAL PARAM BLOCK")
## check if unsupported param passed in
supported_params = get_supported_openai_params(
model=model, custom_llm_provider=custom_llm_provider
@ -3157,9 +3156,8 @@ def get_optional_params(
optional_params=optional_params,
)
print_verbose(
f"(end) INSIDE THE VERTEX AI OPTIONAL PARAM BLOCK - optional_params: {optional_params}"
)
if litellm.vertex_ai_safety_settings is not None:
optional_params["safety_settings"] = litellm.vertex_ai_safety_settings
elif custom_llm_provider == "gemini":
supported_params = get_supported_openai_params(
model=model, custom_llm_provider=custom_llm_provider
@ -3170,7 +3168,7 @@ def get_optional_params(
optional_params=optional_params,
model=model,
)
elif custom_llm_provider == "vertex_ai_beta" or custom_llm_provider == "gemini":
elif custom_llm_provider == "vertex_ai_beta":
supported_params = get_supported_openai_params(
model=model, custom_llm_provider=custom_llm_provider
)
@ -3185,6 +3183,8 @@ def get_optional_params(
else False
),
)
if litellm.vertex_ai_safety_settings is not None:
optional_params["safety_settings"] = litellm.vertex_ai_safety_settings
elif (
custom_llm_provider == "vertex_ai" and model in litellm.vertex_anthropic_models
):