fix(gemini): pass model to context caching URL builder for custom api_base

_get_token_and_url_context_caching() was hardcoding model=None when
calling _check_custom_proxy(), which raises ValueError when api_base
is set because Gemini proxy URLs need the model name:
{api_base}/models/{model}:cachedContents

Fixes #23846
This commit is contained in:
Chesars
2026-03-17 23:21:24 -03:00
parent 2405e0d400
commit 8828f002be
2 changed files with 42 additions and 2 deletions
@@ -51,6 +51,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
model: Optional[str] = None,
) -> Tuple[Optional[str], str]:
"""
Internal function. Returns the token and url for the call.
@@ -89,7 +90,7 @@ class ContextCachingEndpoints(VertexBase):
stream=None,
auth_header=auth_header,
url=url,
model=None,
model=model,
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_api_version="v1beta1"
@@ -342,6 +343,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
headers = {
@@ -488,6 +490,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
headers = {
@@ -1317,4 +1317,41 @@ class TestVertexAIGlobalLocation:
# Assert correct URL format for global with beta API
expected_url = "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents"
assert url == expected_url, f"Expected {expected_url}, got {url}"
assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix"
assert "global-aiplatform" not in url, "URL should not contain 'global-aiplatform' prefix"
def test_gemini_context_caching_with_custom_api_base_passes_model(self):
"""Gemini context caching with custom api_base must pass model to _check_custom_proxy.
Regression test for https://github.com/BerriAI/litellm/issues/23846
Previously model was hardcoded to None, causing ValueError when api_base was set.
"""
caching = ContextCachingEndpoints()
auth_header, url = caching._get_token_and_url_context_caching(
gemini_api_key="test-key",
custom_llm_provider="gemini",
api_base="https://my-proxy.example.com",
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
model="gemini-1.5-pro",
)
assert "models/gemini-1.5-pro" in url
assert url.startswith("https://my-proxy.example.com/")
def test_gemini_context_caching_without_api_base_ignores_model(self):
"""Without custom api_base, model param is not needed (default URL is used)."""
caching = ContextCachingEndpoints()
auth_header, url = caching._get_token_and_url_context_caching(
gemini_api_key="test-key",
custom_llm_provider="gemini",
api_base=None,
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
)
assert "generativelanguage.googleapis.com" in url
assert "cachedContents" in url