mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +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:
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
|
||||
Loading…
Add table
Reference in a new issue