From 501b8fa445fad3dce952742b3a6ecc209ad829b4 Mon Sep 17 00:00:00 2001 From: Enzo Moreira Date: Mon, 13 Apr 2026 21:07:52 -0300 Subject: [PATCH] fix(gemini): pop tool_choice from optional_params when creating cached content When using context caching with Gemini, `check_and_create_cache()` correctly pops `tools` from `optional_params` and includes them in the cached content. However, `tool_choice` was not being popped, causing it to remain as `toolConfig` in the final `GenerateContent` request alongside `cachedContent`. The Gemini API rejects requests that include both `cachedContent` and `toolConfig` with error 400: "CachedContent can not be used with GenerateContent request setting system_instruction, tools or tool_config". This fix pops `tool_choice` from `optional_params` (both sync and async paths) and includes it as `toolConfig` in the cached content request body, matching the existing pattern for `tools`. Verified against the Gemini API: the CachedContent creation endpoint accepts `toolConfig` alongside `tools` and `system_instruction`. --- .../vertex_ai_context_caching.py | 86 ++-- .../test_vertex_ai_context_caching.py | 429 +++++++++--------- 2 files changed, 247 insertions(+), 268 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 b677cf3b1ec..d1b15a4e036 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( @@ -369,9 +358,7 @@ class ContextCachingEndpoints(VertexBase): client = client ## CHECK IF CACHED ALREADY - generated_cache_key = local_cache_obj.get_cache_key( - messages=cached_messages, tools=tools, model=model - ) + generated_cache_key = local_cache_obj.get_cache_key(messages=cached_messages, tools=tools, model=model) google_cache_name = self.check_cache( cache_key=generated_cache_key, client=client, @@ -389,18 +376,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 +402,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 +453,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 +473,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,16 +495,12 @@ 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 - ) + generated_cache_key = local_cache_obj.get_cache_key(messages=cached_messages, tools=tools, model=model) google_cache_name = await self.async_check_cache( cache_key=generated_cache_key, client=client, @@ -535,18 +519,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 +545,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..7b575638c44 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 @@ -206,15 +183,9 @@ class TestContextCachingEndpoints: messages=cached_messages, tools=self.sample_tools, 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 @@ -480,23 +425,15 @@ class TestContextCachingEndpoints: ) @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,137 @@ 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") + 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 +880,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 +936,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 +987,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 +1024,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 +1071,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 +1101,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 +1125,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 +1162,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 +1174,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 +1202,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 +1214,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 +1252,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 +1273,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 +1293,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 +1306,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 +1347,4 @@ class TestVertexAIGlobalLocation: ) assert "generativelanguage.googleapis.com" in url - assert "cachedContents" in url \ No newline at end of file + assert "cachedContents" in url