From a8c835ffc940012b9a0581a381dd2b6efab7a433 Mon Sep 17 00:00:00 2001 From: victoruce <161634297+victoruce@users.noreply.github.com> Date: Tue, 9 Jun 2026 13:39:02 +0200 Subject: [PATCH] fix(vertex_ai): use REP host for context caching on eu/us multi-region endpoints (#29573) Context caching built the cachedContents URL as https://{location}-aiplatform.googleapis.com, which is an invalid host for the eu/us multi-region endpoints and returns 404. The inference path already resolves these to the REP host (https://aiplatform.{geo}.rep.googleapis.com) via get_vertex_base_url(); reuse that helper in _get_token_and_url_context_caching so caching uses the same host as inference. Adds tests covering the eu/us multi-region cachedContents URLs (v1 and v1beta1). Fixes #29571 --- .../vertex_ai_context_caching.py | 14 +-- .../test_vertex_ai_context_caching.py | 92 ++++++++++++++++--- 2 files changed, 83 insertions(+), 23 deletions(-) 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 e9f08f403f9..103801a1e8d 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 @@ -19,7 +19,7 @@ from litellm.types.llms.vertex_ai import ( VertexAICachedContentResponseObject, ) -from ..common_utils import VertexAIError +from ..common_utils import VertexAIError, get_vertex_base_url from ..vertex_llm_base import VertexBase from .transformation import ( separate_cached_messages, @@ -69,17 +69,13 @@ class ContextCachingEndpoints(VertexBase): elif custom_llm_provider == "vertex_ai": auth_header = vertex_auth_header endpoint = "cachedContents" - if vertex_location == "global": - url = f"https://aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" - else: - url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" + base_url = get_vertex_base_url(vertex_location) + url = f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" else: auth_header = vertex_auth_header endpoint = "cachedContents" - if vertex_location == "global": - url = f"https://aiplatform.googleapis.com/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" - else: - url = f"https://{vertex_location}-aiplatform.googleapis.com/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" + base_url = get_vertex_base_url(vertex_location) + url = f"{base_url}/v1beta1/projects/{vertex_project}/locations/{vertex_location}/{endpoint}" return self._check_custom_proxy( api_base=api_base, 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 74888e6cd9e..24359f0921b 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 @@ -86,7 +86,7 @@ class TestContextCachingEndpoints: cached_content = "cached_content_123" optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -129,7 +129,7 @@ class TestContextCachingEndpoints: mock_separate.return_value = ([], self.sample_messages) # No cached messages optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -177,7 +177,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -254,7 +254,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -324,7 +324,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute and Assert with pytest.raises(VertexAIError) as exc_info: @@ -364,7 +364,7 @@ class TestContextCachingEndpoints: cached_content = "cached_content_123" optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -404,7 +404,7 @@ class TestContextCachingEndpoints: mock_separate.return_value = ([], self.sample_messages) optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -453,7 +453,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -535,7 +535,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -606,7 +606,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute and Assert with pytest.raises(VertexAIError) as exc_info: @@ -648,7 +648,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Mock the check_cache to return existing cache so we don't make HTTP calls with patch.object( @@ -694,7 +694,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = self.context_caching.check_and_create_cache( @@ -735,7 +735,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Execute result = await self.context_caching.async_check_and_create_cache( @@ -778,7 +778,7 @@ class TestContextCachingEndpoints: optional_params = self.sample_optional_params.copy() original_tools = optional_params["tools"].copy() test_project = "test_project" - test_location = "test_location" + test_location = "us-central1" # Mock the async_check_cache to return existing cache so we don't make HTTP calls with patch.object( @@ -1881,6 +1881,70 @@ class TestVertexAIGlobalLocation: "global-aiplatform" not in url ), "URL should not contain 'global-aiplatform' prefix" + @pytest.mark.parametrize("vertex_location", ["eu", "us"]) + def test_multi_region_location_url_construction_v1(self, vertex_location): + """Multi-region locations (eu/us) must use the REP host for v1 API. + + Regression test: previously context caching built + ``https://{location}-aiplatform.googleapis.com`` for every non-global + location, producing the invalid host ``eu-aiplatform.googleapis.com`` + (404) for the multi-region endpoints. Inference already resolved these + to ``aiplatform.{geo}.rep.googleapis.com`` via ``get_vertex_base_url``; + context caching now uses the same helper. + """ + caching = ContextCachingEndpoints() + + with patch.object( + caching, + "_check_custom_proxy", + side_effect=lambda **kwargs: (kwargs.get("auth_header"), kwargs.get("url")), + ): + auth_header, url = caching._get_token_and_url_context_caching( + gemini_api_key=None, + custom_llm_provider="vertex_ai", + api_base=None, + vertex_project="test-project", + vertex_location=vertex_location, + vertex_auth_header="Bearer test-token", + ) + + expected_url = ( + f"https://aiplatform.{vertex_location}.rep.googleapis.com" + f"/v1/projects/test-project/locations/{vertex_location}/cachedContents" + ) + assert url == expected_url, f"Expected {expected_url}, got {url}" + assert ( + f"{vertex_location}-aiplatform" not in url + ), "URL must not use the invalid '-aiplatform' host for multi-region" + + @pytest.mark.parametrize("vertex_location", ["eu", "us"]) + def test_multi_region_location_url_construction_v1beta1(self, vertex_location): + """Multi-region locations (eu/us) must use the REP host for v1beta1 API.""" + caching = ContextCachingEndpoints() + + with patch.object( + caching, + "_check_custom_proxy", + side_effect=lambda **kwargs: (kwargs.get("auth_header"), kwargs.get("url")), + ): + auth_header, url = caching._get_token_and_url_context_caching( + gemini_api_key=None, + custom_llm_provider="vertex_ai_beta", + api_base=None, + vertex_project="test-project", + vertex_location=vertex_location, + vertex_auth_header="Bearer test-token", + ) + + expected_url = ( + f"https://aiplatform.{vertex_location}.rep.googleapis.com" + f"/v1beta1/projects/test-project/locations/{vertex_location}/cachedContents" + ) + assert url == expected_url, f"Expected {expected_url}, got {url}" + assert ( + f"{vertex_location}-aiplatform" not in url + ), "URL must not use the invalid '-aiplatform' host for multi-region" + 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.