Include vertex_ai_beta in vertex_ai param mapping

This commit is contained in:
Tiger Yu 2024-06-28 10:36:58 -07:00
parent 2bd8205ef5
commit b0c1d235be
3 changed files with 13 additions and 6 deletions

View file

@ -2324,11 +2324,11 @@ def get_optional_params(
elif k == "hf_model_name" and custom_llm_provider != "sagemaker":
continue
elif (
k.startswith("vertex_") and custom_llm_provider != "vertex_ai"
k.startswith("vertex_") and custom_llm_provider != "vertex_ai" and custom_llm_provider != "vertex_ai_beta"
): # allow dynamically setting vertex ai init logic
continue
passed_params[k] = v
optional_params: Dict = {}
common_auth_dict = litellm.common_cloud_provider_auth_params
@ -2347,7 +2347,7 @@ def get_optional_params(
non_default_params=passed_params, optional_params=optional_params
)
)
elif custom_llm_provider == "vertex_ai":
elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta":
optional_params = litellm.VertexAIConfig().map_special_auth_params(
non_default_params=passed_params, optional_params=optional_params
)
@ -3826,6 +3826,11 @@ def get_supported_openai_params(
return litellm.VertexAIConfig().get_supported_openai_params()
elif request_type == "embeddings":
return litellm.VertexAITextEmbeddingConfig().get_supported_openai_params()
elif custom_llm_provider == "vertex_ai_beta":
if request_type == "chat_completion":
return litellm.VertexAIConfig().get_supported_openai_params()
elif request_type == "embeddings":
return litellm.VertexAITextEmbeddingConfig().get_supported_openai_params()
elif custom_llm_provider == "sagemaker":
return ["stream", "temperature", "max_tokens", "top_p", "stop", "n"]
elif custom_llm_provider == "aleph_alpha":