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 156

fix(gemini): use collection cachedContents endpoint for custom api_base

Browse source 
Gemini context caching with a custom api_base built a model-action URL ({api_base}/models/{model}:cachedContents) via _check_custom_proxy, but cachedContents is a collection endpoint with the model in the request body.

Fixes #34872
This commit is contained in:
Devin AI 2026-07-28 02:07:28 +00:00
parent 3c0b1db633
commit 78151be8bc
2 changed files with 35 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -61,9 +61,14 @@ class ContextCachingEndpoints(VertexBase):
"""
auth_header: Optional[str]
if custom_llm_provider == "gemini":
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
endpoint = "cachedContents"
url = "https://generativelanguage.googleapis.com/v1beta/{}".format(endpoint)
if api_base and gemini_api_key is None:
raise ValueError(
"Missing Gemini API key. Set the GEMINI_API_KEY or GOOGLE_API_KEY environment variable."
)
auth_header = {"x-goog-api-key": gemini_api_key} # type: ignore[assignment]
base_url = api_base.rstrip("/") if api_base else "https://generativelanguage.googleapis.com/v1beta"
return auth_header, "{}/{}".format(base_url, endpoint)
elif custom_llm_provider == "vertex_ai":
auth_header = vertex_auth_header
endpoint = "cachedContents"

36
tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
View file

@ -1881,26 +1881,46 @@ class TestVertexAIGlobalLocation:
"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.
@pytest.mark.parametrize(
"api_base",
["https://my-proxy.example.com/v1beta", "https://my-proxy.example.com/v1beta/"],
)
def test_gemini_context_caching_with_custom_api_base_uses_collection_endpoint(self, api_base):
"""Regression test for https://github.com/BerriAI/litellm/issues/34872
Regression test for https://github.com/BerriAI/litellm/issues/23846
Previously model was hardcoded to None, causing ValueError when api_base was set.
cachedContents is a collection endpoint (model goes in the request body), so a custom
Gemini api_base must not get the model-action treatment (`/models/{model}:cachedContents`),
and auth must stay the x-goog-api-key header dict rather than a stringified Bearer value.
"""
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",
api_base=api_base,
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
model="gemini-1.5-pro",
model="gemini-3-pro-preview",
)
assert "models/gemini-1.5-pro" in url
assert url.startswith("https://my-proxy.example.com/")
assert url == "https://my-proxy.example.com/v1beta/cachedContents"
assert "models/" not in url
assert auth_header == {"x-goog-api-key": "test-key"}
def test_gemini_context_caching_with_custom_api_base_requires_api_key(self):
caching = ContextCachingEndpoints()
with pytest.raises(ValueError, match="Missing Gemini API key"):
caching._get_token_and_url_context_caching(
gemini_api_key=None,
custom_llm_provider="gemini",
api_base="https://my-proxy.example.com/v1beta",
vertex_project=None,
vertex_location=None,
vertex_auth_header=None,
model="gemini-3-pro-preview",
)
def test_gemini_context_caching_without_api_base_ignores_model(self):
"""Without custom api_base, model param is not needed (default URL is used)."""