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

Merge pull request #23928 from Chesars/fix/gemini-context-caching-custom-api-base

Browse source 
fix(gemini): pass model to context caching URL builder for custom api_base
This commit is contained in:
Cesar Garcia 2026-03-18 00:43:16 -03:00 • committed by GitHub
commit 3c7e37799a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 48 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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"
@ -109,6 +110,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
model: Optional[str] = None,
) -> Optional[str]:
"""
Checks if content already cached.
@ -128,6 +130,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
page_token: Optional[str] = None
@ -201,6 +204,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project: Optional[str],
vertex_location: Optional[str],
vertex_auth_header: Optional[str],
model: Optional[str] = None,
) -> Optional[str]:
"""
Checks if content already cached.
@ -220,6 +224,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
page_token: Optional[str] = None
@ -342,6 +347,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
headers = {
@ -377,6 +383,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
if google_cache_name:
return non_cached_messages, optional_params, google_cache_name
@ -488,6 +495,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
headers = {
@ -520,6 +528,7 @@ class ContextCachingEndpoints(VertexBase):
vertex_project=vertex_project,
vertex_location=vertex_location,
vertex_auth_header=vertex_auth_header,
model=model,
)
if google_cache_name:

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

@ -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