mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-15 18:24:07 +00:00
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:
@@ -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 = {
|
||||
|
||||
+38
-1
@@ -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
|
||||
Reference in New Issue
Block a user