fix(utils.py): fix get_api_base to use vertexai_anthropic

This commit is contained in:
Krrish Dholakia 2024-07-16 19:16:32 -07:00
parent 6f55da69a5
commit 155ba055ee

View file

@ -3931,22 +3931,48 @@ def get_api_base(
if dynamic_api_base is not None: if dynamic_api_base is not None:
return dynamic_api_base return dynamic_api_base
stream: bool = getattr(optional_params, "stream", False)
if ( if (
_optional_params.vertex_location is not None _optional_params.vertex_location is not None
and _optional_params.vertex_project is not None and _optional_params.vertex_project is not None
): ):
from litellm.llms.vertex_ai_anthropic import create_vertex_anthropic_url
if "claude" in model:
_api_base = create_vertex_anthropic_url(
vertex_location=_optional_params.vertex_location,
vertex_project=_optional_params.vertex_project,
model=model,
stream=stream,
)
else:
if stream:
_api_base = "{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:streamGenerateContent".format( _api_base = "{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:streamGenerateContent".format(
_optional_params.vertex_location, _optional_params.vertex_location,
_optional_params.vertex_project, _optional_params.vertex_project,
_optional_params.vertex_location, _optional_params.vertex_location,
model, model,
) )
else:
_api_base = "{}-aiplatform.googleapis.com/v1/projects/{}/locations/{}/publishers/google/models/{}:generateContent".format(
_optional_params.vertex_location,
_optional_params.vertex_project,
_optional_params.vertex_location,
model,
)
return _api_base return _api_base
if custom_llm_provider is None: if custom_llm_provider is None:
return None return None
if custom_llm_provider == "gemini": if custom_llm_provider == "gemini":
if stream:
_api_base = "https://generativelanguage.googleapis.com/v1beta/models/{}:streamGenerateContent".format(
model
)
else:
_api_base = "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent".format( _api_base = "https://generativelanguage.googleapis.com/v1beta/models/{}:generateContent".format(
model model
) )