phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Include vertex_ai_beta in vertex_ai param mapping

Browse source
This commit is contained in:
Tiger Yu 2024-06-28 10:36:58 -07:00
parent c4abf87643
commit 6893c836c5
3 changed files with 13 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

2
litellm/__init__.py
View file

@ -106,7 +106,7 @@ aleph_alpha_key: Optional[str] = None
nlp_cloud_key: Optional[str] = None
common_cloud_provider_auth_params: dict = {
"params": ["project", "region_name", "token"],
"providers": ["vertex_ai", "bedrock", "watsonx", "azure"],
"providers": ["vertex_ai", "bedrock", "watsonx", "azure", "vertex_ai_beta"],
}
use_client: bool = False
ssl_verify: bool = True

4
litellm/llms/vertex_httpx.py
View file

@ -733,10 +733,12 @@ class VertexLLM(BaseLLM):
if project_id is None:
project_id = creds.project_id
else:
creds, project_id = google_auth.default(
creds, creds_project_id = google_auth.default(
quota_project_id=project_id,
scopes=["https://www.googleapis.com/auth/cloud-platform"],
)
if project_id is None:
project_id = creds_project_id
creds.refresh(Request())

9
litellm/utils.py
View file

@ -2324,7 +2324,7 @@ 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
@ -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":