From 3e04e2020542e2f31f3eb85da5602ced42625985 Mon Sep 17 00:00:00 2001 From: Aarish Alam Date: Sat, 31 Jan 2026 23:22:19 +0530 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20Bug=20Fix=20#19642=20:=20bug=20i?= =?UTF-8?q?n=20Vertex=20AI=20context=20caching=20=20(#19657)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * add vertex tests * add uperbound * add pagination tests --- .../vertex_ai_context_caching.py | 172 +++++---- .../test_vertex_ai_context_caching.py | 353 ++++++++++++++++++ 2 files changed, 460 insertions(+), 65 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 289963e917..ed4d2d6a74 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 @@ -27,6 +27,8 @@ local_cache_obj = Cache( type=LiteLLMCacheType.LOCAL ) # only used for calling 'get_cache_key' function +MAX_PAGINATION_PAGES = 100 # Reasonable upper bound for pagination + class ContextCachingEndpoints(VertexBase): """ @@ -115,7 +117,7 @@ class ContextCachingEndpoints(VertexBase): - None """ - _, url = self._get_token_and_url_context_caching( + _, base_url = self._get_token_and_url_context_caching( gemini_api_key=api_key, custom_llm_provider=custom_llm_provider, api_base=api_base, @@ -123,43 +125,63 @@ class ContextCachingEndpoints(VertexBase): vertex_location=vertex_location, vertex_auth_header=vertex_auth_header ) - try: - ## LOGGING - logging_obj.pre_call( - input="", - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": url, - "headers": headers, - }, - ) - resp = client.get(url=url, headers=headers) - resp.raise_for_status() - except httpx.HTTPStatusError as e: - if e.response.status_code == 403: + page_token: Optional[str] = None + + # Iterate through all pages + for _ in range(MAX_PAGINATION_PAGES): + # Build URL with pagination token if present + if page_token: + separator = "&" if "?" in base_url else "?" + url = f"{base_url}{separator}pageToken={page_token}" + else: + url = base_url + + try: + ## LOGGING + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": url, + "headers": headers, + }, + ) + + resp = client.get(url=url, headers=headers) + resp.raise_for_status() + except httpx.HTTPStatusError as e: + if e.response.status_code == 403: + return None + raise VertexAIError( + status_code=e.response.status_code, message=e.response.text + ) + except Exception as e: + raise VertexAIError(status_code=500, message=str(e)) + + raw_response = resp.json() + logging_obj.post_call(original_response=raw_response) + + if "cachedContents" not in raw_response: return None - raise VertexAIError( - status_code=e.response.status_code, message=e.response.text - ) - except Exception as e: - raise VertexAIError(status_code=500, message=str(e)) - raw_response = resp.json() - logging_obj.post_call(original_response=raw_response) - if "cachedContents" not in raw_response: - return None + all_cached_items = CachedContentListAllResponseBody(**raw_response) - all_cached_items = CachedContentListAllResponseBody(**raw_response) + if "cachedContents" not in all_cached_items: + return None - if "cachedContents" not in all_cached_items: - return None + # Check current page for matching cache_key + for cached_item in all_cached_items["cachedContents"]: + display_name = cached_item.get("displayName") + if display_name is not None and display_name == cache_key: + return cached_item.get("name") - for cached_item in all_cached_items["cachedContents"]: - display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + # Check if there are more pages + page_token = all_cached_items.get("nextPageToken") + if not page_token: + # No more pages, cache not found + break return None @@ -187,7 +209,7 @@ class ContextCachingEndpoints(VertexBase): - None """ - _, url = self._get_token_and_url_context_caching( + _, base_url = self._get_token_and_url_context_caching( gemini_api_key=api_key, custom_llm_provider=custom_llm_provider, api_base=api_base, @@ -195,43 +217,63 @@ class ContextCachingEndpoints(VertexBase): vertex_location=vertex_location, vertex_auth_header=vertex_auth_header ) - try: - ## LOGGING - logging_obj.pre_call( - input="", - api_key="", - additional_args={ - "complete_input_dict": {}, - "api_base": url, - "headers": headers, - }, - ) - resp = await client.get(url=url, headers=headers) - resp.raise_for_status() - except httpx.HTTPStatusError as e: - if e.response.status_code == 403: + page_token: Optional[str] = None + + # Iterate through all pages + for _ in range(MAX_PAGINATION_PAGES): + # Build URL with pagination token if present + if page_token: + separator = "&" if "?" in base_url else "?" + url = f"{base_url}{separator}pageToken={page_token}" + else: + url = base_url + + try: + ## LOGGING + logging_obj.pre_call( + input="", + api_key="", + additional_args={ + "complete_input_dict": {}, + "api_base": url, + "headers": headers, + }, + ) + + resp = await client.get(url=url, headers=headers) + resp.raise_for_status() + except httpx.HTTPStatusError as e: + if e.response.status_code == 403: + return None + raise VertexAIError( + status_code=e.response.status_code, message=e.response.text + ) + except Exception as e: + raise VertexAIError(status_code=500, message=str(e)) + + raw_response = resp.json() + logging_obj.post_call(original_response=raw_response) + + if "cachedContents" not in raw_response: return None - raise VertexAIError( - status_code=e.response.status_code, message=e.response.text - ) - except Exception as e: - raise VertexAIError(status_code=500, message=str(e)) - raw_response = resp.json() - logging_obj.post_call(original_response=raw_response) - if "cachedContents" not in raw_response: - return None + all_cached_items = CachedContentListAllResponseBody(**raw_response) - all_cached_items = CachedContentListAllResponseBody(**raw_response) + if "cachedContents" not in all_cached_items: + return None - if "cachedContents" not in all_cached_items: - return None + # Check current page for matching cache_key + for cached_item in all_cached_items["cachedContents"]: + display_name = cached_item.get("displayName") + if display_name is not None and display_name == cache_key: + return cached_item.get("name") - for cached_item in all_cached_items["cachedContents"]: - display_name = cached_item.get("displayName") - if display_name is not None and display_name == cache_key: - return cached_item.get("name") + # Check if there are more pages + page_token = all_cached_items.get("nextPageToken") + if not page_token: + # No more pages, cache not found + break return None @@ -501,4 +543,4 @@ class ContextCachingEndpoints(VertexBase): pass async def async_get_cache(self): - pass + pass \ No newline at end of file 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 e9d14d4e18..a47d026c16 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 @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.common_utils import VertexAIError from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + MAX_PAGINATION_PAGES, ContextCachingEndpoints, ) @@ -787,6 +788,358 @@ class TestContextCachingEndpoints: assert original_tools == self.sample_tools +class TestCheckCachePagination: + """Test pagination logic in check_cache and async_check_cache methods.""" + + def setup_method(self): + """Setup for each test method""" + self.context_caching = ContextCachingEndpoints() + self.mock_logging = MagicMock(spec=Logging) + self.mock_client = MagicMock(spec=HTTPHandler) + self.mock_async_client = MagicMock(spec=AsyncHTTPHandler) + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_finds_cache_on_second_page( + self, mock_get_token_url, custom_llm_provider + ): + """Test that check_cache correctly handles pagination and finds cache on second page""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "target_cache_key" + + # Mock first page response (no match, has nextPageToken) + first_page_response = MagicMock() + first_page_response.json.return_value = { + "cachedContents": [ + {"name": "cache_1", "displayName": "cache_key_1"}, + {"name": "cache_2", "displayName": "cache_key_2"}, + ], + "nextPageToken": "token_page_2", + } + + # Mock second page response (has match, no nextPageToken) + second_page_response = MagicMock() + second_page_response.json.return_value = { + "cachedContents": [ + {"name": "cache_3", "displayName": cache_key_to_find}, + {"name": "cache_4", "displayName": "cache_key_4"}, + ] + } + + # Setup mock client to return different responses + self.mock_client.get.side_effect = [first_page_response, second_page_response] + + # Execute + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result == "cache_3" + assert self.mock_client.get.call_count == 2 + # Check that second call includes pageToken + second_call_url = self.mock_client.get.call_args_list[1].kwargs["url"] + assert "pageToken=token_page_2" in second_call_url + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_stops_when_no_next_token( + self, mock_get_token_url, custom_llm_provider + ): + """Test that check_cache stops pagination when no nextPageToken is present""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + # Mock response without nextPageToken + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": "cache_1", "displayName": "cache_key_1"}, + {"name": "cache_2", "displayName": "cache_key_2"}, + ] + } + + self.mock_client.get.return_value = response + + # Execute + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result is None + assert self.mock_client.get.call_count == 1 + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_multiple_pages( + self, mock_get_token_url, custom_llm_provider + ): + """Test that check_cache correctly iterates through multiple pages""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "target_cache_key" + + # Mock three pages + page1 = MagicMock() + page1.json.return_value = { + "cachedContents": [{"name": "cache_1", "displayName": "cache_key_1"}], + "nextPageToken": "token_page_2", + } + + page2 = MagicMock() + page2.json.return_value = { + "cachedContents": [{"name": "cache_2", "displayName": "cache_key_2"}], + "nextPageToken": "token_page_3", + } + + page3 = MagicMock() + page3.json.return_value = { + "cachedContents": [{"name": "cache_3", "displayName": cache_key_to_find}], + } + + self.mock_client.get.side_effect = [page1, page2, page3] + + # Execute + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result == "cache_3" + assert self.mock_client.get.call_count == 3 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_finds_cache_on_second_page( + self, mock_get_token_url, custom_llm_provider + ): + """Test that async_check_cache correctly handles pagination and finds cache on second page""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "target_cache_key" + + # Mock first page response (no match, has nextPageToken) + first_page_response = MagicMock() + first_page_response.json.return_value = { + "cachedContents": [ + {"name": "cache_1", "displayName": "cache_key_1"}, + {"name": "cache_2", "displayName": "cache_key_2"}, + ], + "nextPageToken": "token_page_2", + } + + # Mock second page response (has match, no nextPageToken) + second_page_response = MagicMock() + second_page_response.json.return_value = { + "cachedContents": [ + {"name": "cache_3", "displayName": cache_key_to_find}, + {"name": "cache_4", "displayName": "cache_key_4"}, + ] + } + + # Setup mock async client to return different responses + self.mock_async_client.get = AsyncMock( + side_effect=[first_page_response, second_page_response] + ) + + # Execute + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result == "cache_3" + assert self.mock_async_client.get.call_count == 2 + # Check that second call includes pageToken + second_call_url = self.mock_async_client.get.call_args_list[1].kwargs["url"] + assert "pageToken=token_page_2" in second_call_url + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_stops_when_no_next_token( + self, mock_get_token_url, custom_llm_provider + ): + """Test that async_check_cache stops pagination when no nextPageToken is present""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + # Mock response without nextPageToken + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": "cache_1", "displayName": "cache_key_1"}, + {"name": "cache_2", "displayName": "cache_key_2"}, + ] + } + + self.mock_async_client.get = AsyncMock(return_value=response) + + # Execute + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert + assert result is None + assert self.mock_async_client.get.call_count == 1 + + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + """Test that pagination stops after MAX_PAGINATION_PAGES iterations""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + # Create mock response that always has nextPageToken (infinite pagination scenario) + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + # Create MAX_PAGINATION_PAGES responses, each with a nextPageToken + self.mock_client.get.side_effect = [ + create_page_response(i) for i in range(MAX_PAGINATION_PAGES) + ] + + # Execute + result = self.context_caching.check_cache( + cache_key=cache_key_to_find, + client=self.mock_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert - should return None after exhausting all pages without finding match + assert result is None + # Verify exactly MAX_PAGINATION_PAGES API calls were made (not more) + assert self.mock_client.get.call_count == MAX_PAGINATION_PAGES + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] + ) + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + async def test_async_check_cache_pagination_max_pages_limit( + self, mock_get_token_url, custom_llm_provider + ): + """Test that async pagination stops after MAX_PAGINATION_PAGES iterations""" + # Setup + mock_get_token_url.return_value = ("token", "https://test-url.com") + cache_key_to_find = "nonexistent_cache_key" + + # Create mock response that always has nextPageToken (infinite pagination scenario) + def create_page_response(page_num): + response = MagicMock() + response.json.return_value = { + "cachedContents": [ + {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} + ], + "nextPageToken": f"token_page_{page_num + 1}", + } + return response + + # Create MAX_PAGINATION_PAGES responses, each with a nextPageToken + self.mock_async_client.get = AsyncMock( + side_effect=[create_page_response(i) for i in range(MAX_PAGINATION_PAGES)] + ) + + # Execute + result = await self.context_caching.async_check_cache( + cache_key=cache_key_to_find, + client=self.mock_async_client, + headers={"Authorization": "Bearer token"}, + api_key="test_key", + api_base=None, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="us-central1", + vertex_auth_header="Bearer test-token", + ) + + # Assert - should return None after exhausting all pages without finding match + assert result is None + # Verify exactly MAX_PAGINATION_PAGES async API calls were made (not more) + assert self.mock_async_client.get.call_count == MAX_PAGINATION_PAGES + + class TestVertexAIGlobalLocation: """Test global location handling in context caching."""