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 b677cf3b1ec..bd2aace185f 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 @@ -26,9 +26,7 @@ from .transformation import ( transform_openai_messages_to_gemini_context_caching, ) -local_cache_obj = Cache( - type=LiteLLMCacheType.LOCAL -) # only used for calling 'get_cache_key' function +local_cache_obj = Cache(type=LiteLLMCacheType.LOCAL) # only used for calling 'get_cache_key' function MAX_PAGINATION_PAGES = 100 # Reasonable upper bound for pagination @@ -64,9 +62,7 @@ class ContextCachingEndpoints(VertexBase): if custom_llm_provider == "gemini": auth_header = None endpoint = "cachedContents" - url = "https://generativelanguage.googleapis.com/v1beta/{}?key={}".format( - endpoint, gemini_api_key - ) + url = "https://generativelanguage.googleapis.com/v1beta/{}?key={}".format(endpoint, gemini_api_key) elif custom_llm_provider == "vertex_ai": auth_header = vertex_auth_header endpoint = "cachedContents" @@ -93,9 +89,7 @@ class ContextCachingEndpoints(VertexBase): model=model, vertex_project=vertex_project, vertex_location=vertex_location, - vertex_api_version="v1beta1" - if custom_llm_provider == "vertex_ai_beta" - else "v1", + vertex_api_version="v1beta1" if custom_llm_provider == "vertex_ai_beta" else "v1", ) def check_cache( @@ -161,9 +155,7 @@ class ContextCachingEndpoints(VertexBase): 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 - ) + raise VertexAIError(status_code=e.response.status_code, message=e.response.text) except Exception as e: raise VertexAIError(status_code=500, message=str(e)) @@ -255,9 +247,7 @@ class ContextCachingEndpoints(VertexBase): 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 - ) + raise VertexAIError(status_code=e.response.status_code, message=e.response.text) except Exception as e: raise VertexAIError(status_code=500, message=str(e)) @@ -316,9 +306,7 @@ class ContextCachingEndpoints(VertexBase): if cached_content is not None: return messages, optional_params, cached_content - cached_messages, non_cached_messages = separate_cached_messages( - messages=messages - ) + cached_messages, non_cached_messages = separate_cached_messages(messages=messages) if len(cached_messages) == 0: return messages, optional_params, None @@ -338,6 +326,7 @@ class ContextCachingEndpoints(VertexBase): return messages, optional_params, None tools = optional_params.pop("tools", None) + tool_choice = optional_params.pop("tool_choice", None) ## AUTHORIZATION ## token, url = self._get_token_and_url_context_caching( @@ -370,7 +359,7 @@ class ContextCachingEndpoints(VertexBase): ## CHECK IF CACHED ALREADY generated_cache_key = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools, model=model + messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) google_cache_name = self.check_cache( cache_key=generated_cache_key, @@ -389,18 +378,18 @@ class ContextCachingEndpoints(VertexBase): return non_cached_messages, optional_params, google_cache_name ## TRANSFORM REQUEST - cached_content_request_body = ( - transform_openai_messages_to_gemini_context_caching( - model=model, - messages=cached_messages, - cache_key=generated_cache_key, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - ) + cached_content_request_body = transform_openai_messages_to_gemini_context_caching( + model=model, + messages=cached_messages, + cache_key=generated_cache_key, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, ) cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice ## LOGGING logging_obj.pre_call( @@ -415,7 +404,9 @@ class ContextCachingEndpoints(VertexBase): try: response = client.post( - url=url, headers=headers, json=cached_content_request_body # type: ignore + url=url, + headers=headers, + json=cached_content_request_body, # type: ignore ) response.raise_for_status() except httpx.HTTPStatusError as err: @@ -464,9 +455,7 @@ class ContextCachingEndpoints(VertexBase): if cached_content is not None: return messages, optional_params, cached_content - cached_messages, non_cached_messages = separate_cached_messages( - messages=messages - ) + cached_messages, non_cached_messages = separate_cached_messages(messages=messages) if len(cached_messages) == 0: return messages, optional_params, None @@ -486,6 +475,7 @@ class ContextCachingEndpoints(VertexBase): return messages, optional_params, None tools = optional_params.pop("tools", None) + tool_choice = optional_params.pop("tool_choice", None) ## AUTHORIZATION ## token, url = self._get_token_and_url_context_caching( @@ -507,15 +497,13 @@ class ContextCachingEndpoints(VertexBase): headers.update(extra_headers) if client is None or not isinstance(client, AsyncHTTPHandler): - client = get_async_httpx_client( - params={"timeout": timeout}, llm_provider=litellm.LlmProviders.VERTEX_AI - ) + client = get_async_httpx_client(params={"timeout": timeout}, llm_provider=litellm.LlmProviders.VERTEX_AI) else: client = client ## CHECK IF CACHED ALREADY generated_cache_key = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools, model=model + messages=cached_messages, tools=tools, tool_choice=tool_choice, model=model ) google_cache_name = await self.async_check_cache( cache_key=generated_cache_key, @@ -535,18 +523,18 @@ class ContextCachingEndpoints(VertexBase): return non_cached_messages, optional_params, google_cache_name ## TRANSFORM REQUEST - cached_content_request_body = ( - transform_openai_messages_to_gemini_context_caching( - model=model, - messages=cached_messages, - cache_key=generated_cache_key, - custom_llm_provider=custom_llm_provider, - vertex_project=vertex_project, - vertex_location=vertex_location, - ) + cached_content_request_body = transform_openai_messages_to_gemini_context_caching( + model=model, + messages=cached_messages, + cache_key=generated_cache_key, + custom_llm_provider=custom_llm_provider, + vertex_project=vertex_project, + vertex_location=vertex_location, ) cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice ## LOGGING logging_obj.pre_call( @@ -561,7 +549,9 @@ class ContextCachingEndpoints(VertexBase): try: response = await client.post( - url=url, headers=headers, json=cached_content_request_body # type: ignore + url=url, + headers=headers, + json=cached_content_request_body, # type: ignore ) response.raise_for_status() except httpx.HTTPStatusError as err: 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 11ccd34804a..26a8d6d4dad 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 @@ -1,14 +1,11 @@ import os import sys -from typing import List from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest -sys.path.insert( - 0, os.path.abspath("../../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler @@ -69,18 +66,10 @@ class TestContextCachingEndpoints: """Teardown for each test method""" self._token_check_patcher.stop() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) - def test_check_and_create_cache_with_cached_content( - self, mock_cache_obj, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") + def test_check_and_create_cache_with_cached_content(self, mock_cache_obj, mock_separate, custom_llm_provider): """Test check_and_create_cache when cached_content is provided""" # Setup cached_content = "cached_content_123" @@ -115,15 +104,9 @@ class TestContextCachingEndpoints: mock_separate.assert_not_called() mock_cache_obj.get_cache_key.assert_not_called() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - def test_check_and_create_cache_no_cached_messages( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + def test_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): """Test check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) # No cached messages @@ -153,15 +136,9 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache is None - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "check_cache") def test_check_and_create_cache_existing_cache_found( self, mock_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -201,20 +178,14 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache == "existing_cache_name" - # Verify cache key was generated with tools and model + # Verify cache key was generated with tools, tool_choice, and model mock_cache_obj.get_cache_key.assert_called_once_with( - messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro" + messages=cached_messages, tools=self.sample_tools, tool_choice=None, model="gemini-1.5-pro" ) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @@ -281,15 +252,9 @@ class TestContextCachingEndpoints: assert "tools" in call_args.kwargs["json"] assert call_args.kwargs["json"]["tools"] == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") def test_check_and_create_cache_http_error( @@ -314,9 +279,7 @@ class TestContextCachingEndpoints: mock_response = MagicMock() mock_response.status_code = 400 mock_response.text = "Bad Request" - http_error = httpx.HTTPStatusError( - "Error", request=MagicMock(), response=mock_response - ) + http_error = httpx.HTTPStatusError("Error", request=MagicMock(), response=mock_response) self.mock_client.post.side_effect = http_error optional_params = self.sample_optional_params.copy() @@ -344,15 +307,9 @@ class TestContextCachingEndpoints: assert "Bad Request" in str(exc_info.value.message) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") async def test_async_check_and_create_cache_with_cached_content( self, mock_cache_obj, mock_separate, custom_llm_provider ): @@ -387,15 +344,9 @@ class TestContextCachingEndpoints: assert returned_cache == cached_content @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - async def test_async_check_and_create_cache_no_cached_messages( - self, mock_separate, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + async def test_async_check_and_create_cache_no_cached_messages(self, mock_separate, custom_llm_provider): """Test async_check_and_create_cache when no cached messages are found""" # Setup mock_separate.return_value = ([], self.sample_messages) @@ -426,15 +377,9 @@ class TestContextCachingEndpoints: assert returned_cache is None @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "async_check_cache") async def test_async_check_and_create_cache_existing_cache_found( self, mock_async_check_cache, mock_cache_obj, mock_separate, custom_llm_provider @@ -474,29 +419,21 @@ class TestContextCachingEndpoints: assert returned_params == optional_params assert returned_cache == "existing_cache_name" - # Verify cache key was generated with tools and model + # Verify cache key was generated with tools, tool_choice, and model mock_cache_obj.get_cache_key.assert_called_once_with( - messages=cached_messages, tools=self.sample_tools, model="gemini-1.5-pro" + messages=cached_messages, tools=self.sample_tools, tool_choice=None, model="gemini-1.5-pro" ) @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" ) @patch.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" - ) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") async def test_async_check_and_create_cache_create_new_cache( self, mock_get_client, @@ -560,20 +497,12 @@ class TestContextCachingEndpoints: assert call_args.kwargs["json"]["tools"] == self.sample_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") @patch.object(ContextCachingEndpoints, "async_check_cache") @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client" - ) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") async def test_async_check_and_create_cache_timeout_error( self, mock_get_client, @@ -594,9 +523,7 @@ class TestContextCachingEndpoints: mock_get_token_url.return_value = ("token", "https://test-url.com") # Mock timeout error - self.mock_async_client.post = AsyncMock( - side_effect=httpx.TimeoutException("Timeout") - ) + self.mock_async_client.post = AsyncMock(side_effect=httpx.TimeoutException("Timeout")) optional_params = self.sample_optional_params.copy() test_project = "test_project" @@ -622,20 +549,14 @@ class TestContextCachingEndpoints: assert exc_info.value.status_code == 408 assert "Timeout error occurred" in str(exc_info.value.message) - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tools_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): """Test that tools are properly popped from optional_params when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [ - self.sample_messages[0] - ] # System message with cache_control + cached_messages = [self.sample_messages[0]] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -645,9 +566,7 @@ class TestContextCachingEndpoints: test_location = "test_location" # Mock the check_cache to return existing cache so we don't make HTTP calls - with patch.object( - self.context_caching, "check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): # Execute result = self.context_caching.check_and_create_cache( messages=self.sample_messages, @@ -670,12 +589,8 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - def test_check_and_create_cache_tools_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): """Test that tools are NOT popped from optional_params when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -711,12 +626,8 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tools_not_popped_when_no_cached_messages(self, custom_llm_provider): """Test that tools are NOT popped from optional_params in async version when there are no cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" @@ -752,20 +663,14 @@ class TestContextCachingEndpoints: assert optional_params["tools"] == original_tools @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - async def test_async_check_and_create_cache_tools_popped_from_optional_params( - self, custom_llm_provider - ): + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tools_popped_from_optional_params(self, custom_llm_provider): """Test that tools are properly popped from optional_params in async version when there are cached messages""" with patch( "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" ) as mock_separate: # Mock to return cached messages so tools get popped - cached_messages = [ - self.sample_messages[0] - ] # System message with cache_control + cached_messages = [self.sample_messages[0]] # System message with cache_control non_cached_messages = [self.sample_messages[1]] # User message mock_separate.return_value = (cached_messages, non_cached_messages) @@ -775,9 +680,7 @@ class TestContextCachingEndpoints: test_location = "test_location" # Mock the async_check_cache to return existing cache so we don't make HTTP calls - with patch.object( - self.context_caching, "async_check_cache", return_value="existing_cache" - ): + with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): # Execute result = await self.context_caching.async_check_and_create_cache( messages=self.sample_messages, @@ -800,15 +703,292 @@ class TestContextCachingEndpoints: # But original tools should still be available for comparison assert original_tools == self.sample_tools - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) - def test_check_and_create_cache_skips_when_below_min_tokens( - self, mock_separate, custom_llm_provider + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): + """Test that tool_choice is popped from optional_params when there are cached messages. + + Gemini rejects requests that include both cachedContent and toolConfig. + tool_choice must be moved into the cached content alongside tools. + """ + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) as mock_separate: + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = "auto" + + with patch.object(self.context_caching, "check_cache", return_value="existing_cache"): + result = self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + assert "tool_choice" not in optional_params + + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + def test_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages(self, custom_llm_provider): + """Test that tool_choice is NOT popped when there are no cached messages""" + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) as mock_separate: + mock_separate.return_value = ([], self.sample_messages) + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = "auto" + + result = self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + assert "tool_choice" in optional_params + assert optional_params["tool_choice"] == "auto" + + @pytest.mark.asyncio + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tool_choice_popped_from_optional_params(self, custom_llm_provider): + """Test that tool_choice is popped from optional_params in async version when there are cached messages""" + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) as mock_separate: + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = "auto" + + with patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"): + result = await self.context_caching.async_check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + assert "tool_choice" not in optional_params + + @pytest.mark.asyncio + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + async def test_async_check_and_create_cache_tool_choice_not_popped_when_no_cached_messages( + self, custom_llm_provider ): + """Test that tool_choice is NOT popped in async version when there are no cached messages""" + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" + ) as mock_separate: + mock_separate.return_value = ([], self.sample_messages) + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = "auto" + + result = await self.context_caching.async_check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + assert "tool_choice" in optional_params + assert optional_params["tool_choice"] == "auto" + + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") + @patch.object(ContextCachingEndpoints, "check_cache") + def test_cache_key_includes_tool_choice(self, mock_check_cache, mock_cache_obj, mock_separate, custom_llm_provider): + """Test that tool_choice is included in the cache key computation. + + Different tool_choice values (e.g. 'auto' vs 'required') must produce + different cache keys, otherwise the second call silently reuses cached + content created with a different toolConfig. + """ + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_check_cache.return_value = "existing_cache" + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = "required" + + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + mock_cache_obj.get_cache_key.assert_called_once_with( + messages=cached_messages, tools=self.sample_tools, tool_choice="required", model="gemini-1.5-pro" + ) + + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" + ) + @patch.object(ContextCachingEndpoints, "check_cache") + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + def test_create_new_cache_includes_tool_config_in_body( + self, + mock_get_token_url, + mock_check_cache, + mock_transform, + mock_cache_obj, + mock_separate, + custom_llm_provider, + ): + """Test that toolConfig is included in the HTTP request body when creating new cached content.""" + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_check_cache.return_value = None + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []} + + mock_response = MagicMock() + mock_response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + } + self.mock_client.post.return_value = mock_response + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = "required" + + self.context_caching.check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + call_args = self.mock_client.post.call_args + assert call_args.kwargs["json"]["tools"] == self.sample_tools + assert call_args.kwargs["json"]["toolConfig"] == "required" + + @pytest.mark.asyncio + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj") + @patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching" + ) + @patch.object(ContextCachingEndpoints, "async_check_cache") + @patch.object(ContextCachingEndpoints, "_get_token_and_url_context_caching") + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.get_async_httpx_client") + async def test_async_create_new_cache_includes_tool_config_in_body( + self, + mock_get_client, + mock_get_token_url, + mock_async_check_cache, + mock_transform, + mock_cache_obj, + mock_separate, + custom_llm_provider, + ): + """Test that toolConfig is included in the HTTP request body when creating new cached content (async).""" + cached_messages = [self.sample_messages[0]] + non_cached_messages = [self.sample_messages[1]] + mock_separate.return_value = (cached_messages, non_cached_messages) + + mock_cache_obj.get_cache_key.return_value = "test_cache_key" + mock_async_check_cache.return_value = None + mock_get_token_url.return_value = ("token", "https://test-url.com") + mock_transform.return_value = {"model": "gemini-1.5-pro", "contents": []} + + mock_response = MagicMock() + mock_response.json.return_value = { + "name": "new_cache_name", + "model": "gemini-1.5-pro", + } + self.mock_async_client.post = AsyncMock(return_value=mock_response) + + optional_params = self.sample_optional_params.copy() + optional_params["tool_choice"] = "required" + + await self.context_caching.async_check_and_create_cache( + messages=self.sample_messages, + optional_params=optional_params, + api_key="test_key", + api_base=None, + model="gemini-1.5-pro", + client=self.mock_async_client, + timeout=30.0, + logging_obj=self.mock_logging, + custom_llm_provider=custom_llm_provider, + vertex_project="test_project", + vertex_location="test_location", + vertex_auth_header="vertext_test_token", + ) + + call_args = self.mock_async_client.post.call_args + assert call_args.kwargs["json"]["tools"] == self.sample_tools + assert call_args.kwargs["json"]["toolConfig"] == "required" + + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") + def test_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): """Test that context caching is skipped when cached content is below 1024 tokens. Gemini requires a minimum of 1024 tokens for context caching. If the cached @@ -855,16 +1035,10 @@ class TestContextCachingEndpoints: # Restart the patcher so teardown_method can stop it cleanly self._token_check_patcher.start() - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) - @patch( - "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages" - ) + @pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"]) + @patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages") @pytest.mark.asyncio - async def test_async_check_and_create_cache_skips_when_below_min_tokens( - self, mock_separate, custom_llm_provider - ): + async def test_async_check_and_create_cache_skips_when_below_min_tokens(self, mock_separate, custom_llm_provider): """Test that async context caching is skipped when cached content is below 1024 tokens.""" # Stop the default mock so the real token count check runs self._token_check_patcher.stop() @@ -917,13 +1091,9 @@ class TestCheckCachePagination: 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"] - ) + @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 - ): + 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") @@ -972,13 +1142,9 @@ class TestCheckCachePagination: 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"] - ) + @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 - ): + 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") @@ -1013,13 +1179,9 @@ class TestCheckCachePagination: assert result is None assert self.mock_client.get.call_count == 1 - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @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 - ): + 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") @@ -1064,9 +1226,7 @@ class TestCheckCachePagination: assert self.mock_client.get.call_count == 3 @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @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 @@ -1096,9 +1256,7 @@ class TestCheckCachePagination: } # Setup mock async client to return different responses - self.mock_async_client.get = AsyncMock( - side_effect=[first_page_response, second_page_response] - ) + self.mock_async_client.get = AsyncMock(side_effect=[first_page_response, second_page_response]) # Execute result = await self.context_caching.async_check_cache( @@ -1122,13 +1280,9 @@ class TestCheckCachePagination: assert "pageToken=token_page_2" in second_call_url @pytest.mark.asyncio - @pytest.mark.parametrize( - "custom_llm_provider", ["gemini", "vertex_ai", "vertex_ai_beta"] - ) + @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 - ): + 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") @@ -1163,13 +1317,9 @@ class TestCheckCachePagination: 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"] - ) + @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 - ): + 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") @@ -1179,17 +1329,13 @@ class TestCheckCachePagination: def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], + "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) - ] + self.mock_client.get.side_effect = [create_page_response(i) for i in range(MAX_PAGINATION_PAGES)] # Execute result = self.context_caching.check_cache( @@ -1211,13 +1357,9 @@ class TestCheckCachePagination: 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"] - ) + @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 - ): + 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") @@ -1227,9 +1369,7 @@ class TestCheckCachePagination: def create_page_response(page_num): response = MagicMock() response.json.return_value = { - "cachedContents": [ - {"name": f"cache_{page_num}", "displayName": f"key_{page_num}"} - ], + "cachedContents": [{"name": f"cache_{page_num}", "displayName": f"key_{page_num}"}], "nextPageToken": f"token_page_{page_num + 1}", } return response @@ -1267,7 +1407,9 @@ class TestVertexAIGlobalLocation: caching = ContextCachingEndpoints() # Mock the _check_custom_proxy to return the URL unchanged - with patch.object(caching, '_check_custom_proxy', side_effect=lambda **kwargs: (kwargs.get('auth_header'), kwargs.get('url'))): + 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", @@ -1286,7 +1428,9 @@ class TestVertexAIGlobalLocation: """Test that regional location uses correct URL (with location prefix) for v1 API.""" caching = ContextCachingEndpoints() - with patch.object(caching, '_check_custom_proxy', side_effect=lambda **kwargs: (kwargs.get('auth_header'), kwargs.get('url'))): + 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", @@ -1304,7 +1448,9 @@ class TestVertexAIGlobalLocation: """Test that global location uses correct URL for v1beta1 API.""" caching = ContextCachingEndpoints() - with patch.object(caching, '_check_custom_proxy', side_effect=lambda **kwargs: (kwargs.get('auth_header'), kwargs.get('url'))): + 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", @@ -1315,7 +1461,9 @@ class TestVertexAIGlobalLocation: ) # Assert correct URL format for global with beta API - expected_url = "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/cachedContents" + 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" @@ -1354,4 +1502,4 @@ class TestVertexAIGlobalLocation: ) assert "generativelanguage.googleapis.com" in url - assert "cachedContents" in url \ No newline at end of file + assert "cachedContents" in url