scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(vertex_ai): return create_vertex_url result directly for openai-path partner models with custom api_base (#32380)

Browse source
This commit is contained in:
Mateo Wang 2026-07-07 14:49:21 -07:00 • committed by GitHub
parent ee69a62304
commit 46d9742950
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 76 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/llms/vertex_ai/vertex_llm_base.py
View file

@ -316,6 +316,9 @@ class VertexBase:
api_base=api_base,
)
if partner == VertexPartnerProvider.llama:
return default_api_base
if len(default_api_base.split(":")) > 1:
endpoint = default_api_base.split(":")[-1]
else:

73
tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py
View file

@ -15,6 +15,7 @@ sys.path.insert(
import litellm
from litellm.llms.vertex_ai.vertex_ai_aws_wif import VertexAIAwsWifAuth
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.vertex_ai import VertexPartnerProvider
def run_sync(coro):
@ -1025,6 +1026,78 @@ class TestVertexBase:
assert result_url == "https://10.96.32.8/v1/projects/test-project/locations/us-central1/endpoints/1234567890:predict"
@pytest.mark.parametrize(
"custom_api_base, stream, expected_url",
[
(
"https://aiplatform-myendpoint.p.googleapis.com",
False,
"https://aiplatform-myendpoint.p.googleapis.com/v1/projects/test-project/locations/global/endpoints/openapi/chat/completions",
),
(
"https://aiplatform-myendpoint.p.googleapis.com",
True,
"https://aiplatform-myendpoint.p.googleapis.com/v1/projects/test-project/locations/global/endpoints/openapi/chat/completions",
),
(
"https://gateway.example.com/vertex-proxy",
False,
"https://gateway.example.com/vertex-proxy/v1/projects/test-project/locations/global/endpoints/openapi/chat/completions",
),
],
ids=["psc-host", "psc-host-streaming", "api-base-with-path"],
)
def test_get_complete_vertex_url_openai_path_partner_custom_api_base(
self, custom_api_base, stream, expected_url
):
vertex_base = VertexBase()
result = vertex_base.get_complete_vertex_url(
custom_api_base=custom_api_base,
vertex_location="global",
vertex_project="test-project",
project_id="test-project",
partner=VertexPartnerProvider.llama,
stream=stream,
model="minimaxai/minimax-m2-maas",
)
assert result == expected_url
assert result.count("://") == 1
def test_get_complete_vertex_url_openai_path_partner_default_api_base(self):
vertex_base = VertexBase()
result = vertex_base.get_complete_vertex_url(
custom_api_base=None,
vertex_location="us-central1",
vertex_project="test-project",
project_id="test-project",
partner=VertexPartnerProvider.llama,
stream=True,
model="meta/llama-3.1-405b-instruct-maas",
)
assert (
result
== "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/endpoints/openapi/chat/completions"
)
def test_get_complete_vertex_url_rawpredict_partner_custom_api_base_keeps_endpoint_format(self):
vertex_base = VertexBase()
result = vertex_base.get_complete_vertex_url(
custom_api_base="https://gateway.example.com/vertex-proxy",
vertex_location="us-central1",
vertex_project="test-project",
project_id="test-project",
partner=VertexPartnerProvider.mistralai,
stream=False,
model="mistral-large-2411",
)
assert result == "https://gateway.example.com/vertex-proxy:rawPredict"
@pytest.mark.parametrize(
"api_base, custom_llm_provider, gemini_api_key, endpoint, stream, auth_header, url, model, expected_auth_header, expected_url",
[