scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

fix(gemini): avoid duplicate model route for full api_base

Browse source
This commit is contained in:
StatPan 2026-05-01 21:11:57 +09:00
parent eab0075353
commit d8418fff08
2 changed files with 47 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -412,7 +412,14 @@ class VertexBase:
raise ValueError(
"Model parameter is required for Gemini custom API base URLs"
)
url = "{}/models/{}:{}".format(api_base, model, endpoint)
if api_base.endswith(f"/models/{model}:{endpoint}"):
url = api_base
elif api_base.endswith(f"/models/{model}"):
url = f"{api_base}:{endpoint}"
else:
url = "{}/models/{}:{}".format(
api_base.rstrip("/"), model, endpoint
)
if gemini_api_key is None:
raise ValueError(
"Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."

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

@ -891,6 +891,45 @@ class TestVertexBase:
result_url == expected_url
), f"Expected {expected_url}, got {result_url} for model {model}"
@pytest.mark.parametrize(
"api_base, endpoint, expected_url",
[
(
"https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
"generateContent",
"https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
),
(
"https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite",
"generateContent",
"https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
),
(
"https://proxy.example.com/generativelanguage.googleapis.com/v1beta/",
"generateContent",
"https://proxy.example.com/generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
),
],
)
def test_check_custom_proxy_gemini_prebuilt_route_handling(
self, api_base, endpoint, expected_url
):
"""Test that Gemini custom api_base values are not double-appended when they already include model routing."""
vertex_base = VertexBase()
_, result_url = vertex_base._check_custom_proxy(
api_base=api_base,
custom_llm_provider="gemini",
gemini_api_key="test-api-key",
endpoint=endpoint,
stream=False,
auth_header=None,
url="https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-lite:generateContent",
model="gemini-2.5-flash-lite",
)
assert result_url == expected_url
def test_check_custom_proxy_streaming_parameter(self):
"""Test that streaming parameter correctly adds ?alt=sse to URLs"""
vertex_base = VertexBase()