diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 0bf3715f798..5e6dd8ae525 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from typing import List, Literal, Optional, Tuple, Union import httpx @@ -50,20 +51,24 @@ class ContextCachingEndpoints(VertexBase): vertex_location: Optional[str], vertex_auth_header: Optional[str], model: Optional[str] = None, - ) -> Tuple[Optional[str], str]: + ) -> Tuple[Union[str, Mapping[str, str], None], str]: """ Internal function. Returns the token and url for the call. Handles logic if it's google ai studio vs. vertex ai. + For Google AI Studio the credential is a header mapping, since the API key is sent as + `x-goog-api-key` instead of a bearer token. + Returns token, url """ 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) + gemini_auth_header = {"x-goog-api-key": gemini_api_key} if gemini_api_key is not None else None + if api_base: + return gemini_auth_header, "{}/cachedContents".format(api_base.rstrip("/")) + return gemini_auth_header, "https://generativelanguage.googleapis.com/v1beta/cachedContents" elif custom_llm_provider == "vertex_ai": auth_header = vertex_auth_header endpoint = "cachedContents" @@ -339,7 +344,7 @@ class ContextCachingEndpoints(VertexBase): headers = { "Content-Type": "application/json", } - if isinstance(token, dict): + if isinstance(token, Mapping): headers.update(token) elif token is not None: headers["Authorization"] = f"Bearer {token}" @@ -490,7 +495,7 @@ class ContextCachingEndpoints(VertexBase): headers = { "Content-Type": "application/json", } - if isinstance(token, dict): + if isinstance(token, Mapping): headers.update(token) elif token is not None: headers["Authorization"] = f"Bearer {token}" diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py index cf75964ddb7..e5e12e5d15d 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py @@ -1881,26 +1881,30 @@ 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/genai/v1beta", "https://my-proxy.example.com/genai/v1beta/"] + ) + def test_gemini_context_caching_with_custom_api_base_uses_collection_endpoint(self, api_base): + """Gemini context caching with custom api_base must hit the `cachedContents` collection. - Regression test for https://github.com/BerriAI/litellm/issues/23846 - Previously model was hardcoded to None, causing ValueError when api_base was set. + Regression test for https://github.com/BerriAI/litellm/issues/34872 (and #23846, which + required the call to not raise when api_base is set). `cachedContents` is not a model + action, so the model must not be spliced into the url. """ 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-2.5-flash", ) - assert "models/gemini-1.5-pro" in url - assert url.startswith("https://my-proxy.example.com/") + assert url == "https://my-proxy.example.com/genai/v1beta/cachedContents" + assert auth_header == {"x-goog-api-key": "test-key"} def test_gemini_context_caching_without_api_base_ignores_model(self): """Without custom api_base, model param is not needed (default URL is used).""" @@ -1971,3 +1975,101 @@ class TestContextCachingMultiRegionUrls: assert url.startswith("https://aiplatform.googleapis.com/") assert "/locations/global/cachedContents" in url + + +class TestGeminiCustomApiBaseCacheRequests: + """Regression coverage for #34872: with a custom Gemini `api_base`, cache list/create + requests must go to `{api_base}/cachedContents` with the api key in `x-goog-api-key`, + never a stringified dict in an `Authorization` header.""" + + EXPECTED_URL = "https://my-proxy.example.com/genai/v1beta/cachedContents" + + def setup_method(self): + self.caching = ContextCachingEndpoints() + self.logging_obj = MagicMock(spec=Logging) + self.messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Stable cacheable prefix. " * 800, + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": "Reply with ok."}, + ] + + @staticmethod + def _empty_list_response(): + response = MagicMock() + response.json.return_value = {"cachedContents": []} + return response + + @staticmethod + def _create_response(): + response = MagicMock() + response.json.return_value = { + "name": "cachedContents/abc123", + "model": "models/gemini-2.5-flash", + } + return response + + def _assert_requests(self, get_call, post_call): + assert get_call.kwargs["url"] == self.EXPECTED_URL + assert post_call.kwargs["url"] == self.EXPECTED_URL + assert post_call.kwargs["json"]["model"] == "models/gemini-2.5-flash" + for headers in (get_call.kwargs["headers"], post_call.kwargs["headers"]): + assert headers["x-goog-api-key"] == "sk-gemini-key" + assert headers["x-custom"] == "1" + assert "Authorization" not in headers + + def test_sync_cache_lifecycle_uses_collection_url_and_api_key_header(self): + client = MagicMock(spec=HTTPHandler) + client.get.return_value = self._empty_list_response() + client.post.return_value = self._create_response() + + _, _, cache_name = self.caching.check_and_create_cache( + messages=self.messages, + optional_params={}, + api_key="sk-gemini-key", + api_base="https://my-proxy.example.com/genai/v1beta", + model="gemini-2.5-flash", + client=client, + timeout=30.0, + logging_obj=self.logging_obj, + custom_llm_provider="gemini", + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + extra_headers={"x-custom": "1"}, + ) + + assert cache_name == "cachedContents/abc123" + self._assert_requests(client.get.call_args, client.post.call_args) + + @pytest.mark.asyncio + async def test_async_cache_lifecycle_uses_collection_url_and_api_key_header(self): + client = MagicMock(spec=AsyncHTTPHandler) + client.get = AsyncMock(return_value=self._empty_list_response()) + client.post = AsyncMock(return_value=self._create_response()) + + _, _, cache_name = await self.caching.async_check_and_create_cache( + messages=self.messages, + optional_params={}, + api_key="sk-gemini-key", + api_base="https://my-proxy.example.com/genai/v1beta", + model="gemini-2.5-flash", + client=client, + timeout=30.0, + logging_obj=self.logging_obj, + custom_llm_provider="gemini", + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + extra_headers={"x-custom": "1"}, + ) + + assert cache_name == "cachedContents/abc123" + self._assert_requests(client.get.call_args, client.post.call_args)