From eac9912e33bf6655043f702232ed7282fedae980 Mon Sep 17 00:00:00 2001 From: felipebridge Date: Thu, 28 May 2026 11:37:50 -0300 Subject: [PATCH] fix(vertex): preserve tool_choice through Gemini context-caching When context-caching is active, `_transform_request_body()` guards tool- related fields behind `can_send_cache_incompatible_fields` (which is `False` whenever `cached_content` is set) and silently drops `toolConfig`. `check_and_create_cache` already pops `tools` from `optional_params` and embeds it in the `CachedContentRequestBody`; `tool_choice` received no equivalent treatment. Fix: in both the sync (`check_and_create_cache`) and async (`async_check_and_create_cache`) paths: 1. Pop `tool_choice` from `optional_params` alongside `tools` so the downstream guard in `transformation.py` never sees it. 2. Include `tool_choice` in the `get_cache_key()` call so requests that differ only in tool-choice produce distinct cache entries. 3. Set `cached_content_request_body["toolConfig"] = tool_choice` (only when non-None) so the constraint is forwarded to the Vertex AI cached-content creation endpoint. `CachedContentRequestBody` already declares the `toolConfig` field (litellm/types/llms/vertex_ai.py), so no type changes are required. Fixes BerriAI/litellm#29088 --- .../vertex_ai_context_caching.py | 10 +- .../test_context_caching_ttl.py | 125 ++++++++++++++++++ 2 files changed, 133 insertions(+), 2 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 3f945adca0d..e9f08f403f9 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 @@ -337,6 +337,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( @@ -371,7 +372,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, @@ -402,6 +403,8 @@ class ContextCachingEndpoints(VertexBase): ) cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice ## LOGGING logging_obj.pre_call( @@ -487,6 +490,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( @@ -518,7 +522,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 = await self.async_check_cache( cache_key=generated_cache_key, @@ -550,6 +554,8 @@ class ContextCachingEndpoints(VertexBase): ) cached_content_request_body["tools"] = tools + if tool_choice is not None: + cached_content_request_body["toolConfig"] = tool_choice ## LOGGING logging_obj.pre_call( diff --git a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py index 8a13baa0006..6131d133454 100644 --- a/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py +++ b/tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py @@ -389,5 +389,130 @@ class TestEdgeCases: assert ttl == "3600s" +class TestToolChoiceContextCaching: + """Tests for tool_choice handling in context caching. + + Regression tests for https://github.com/BerriAI/litellm/issues/29088 + + Before the fix, `tool_choice` was never popped from `optional_params` during + cache creation. It then reached `_transform_request_body()` where the guard + ``can_send_cache_incompatible_fields`` silently dropped it because + ``cached_content`` was already set. Result: `toolConfig` appeared in neither + the cached-content request nor the final API call. + + The fix mirrors the existing `tools` handling: pop `tool_choice` from + `optional_params` at cache-creation time, include it in the cache key, and + set it on the `CachedContentRequestBody` as `toolConfig`. + """ + + def _make_cached_messages(self) -> list: + return [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "stable prefix", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + + def test_tool_choice_popped_from_optional_params(self) -> None: + """tool_choice must be removed from optional_params so the downstream + guard in transformation.py cannot silently drop it.""" + from unittest.mock import MagicMock, patch + from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + ContextCachingEndpoints, + ) + + endpoint = ContextCachingEndpoints() + tool_choice = {"mode": "ANY", "allowedFunctionNames": ["my_func"]} + optional_params = { + "tools": [{"function_declarations": [{"name": "my_func"}]}], + "tool_choice": tool_choice, + } + + with ( + patch.object(endpoint, "_get_token_and_url_context_caching", return_value=("token", "http://url")), + patch.object(endpoint, "check_cache", return_value=None), + patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt", return_value=True), + patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.HTTPHandler") as mock_http, + ): + mock_resp = MagicMock() + mock_resp.json.return_value = {"name": "cached/123", "model": "gemini-1.5-pro"} + mock_http.return_value.post.return_value = mock_resp + + endpoint.check_and_create_cache( + messages=self._make_cached_messages(), + optional_params=optional_params, + api_key="test-key", + api_base=None, + model="gemini-1.5-pro", + client=None, + timeout=None, + logging_obj=MagicMock(), + custom_llm_provider="gemini", + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + + assert "tool_choice" not in optional_params, ( + "tool_choice must be popped from optional_params so the downstream " + "can_send_cache_incompatible_fields guard cannot drop it silently" + ) + + def test_tool_choice_included_in_cached_content_body(self) -> None: + """toolConfig must appear in the body sent to the cache-creation endpoint.""" + from unittest.mock import MagicMock, call, patch + from litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching import ( + ContextCachingEndpoints, + ) + + endpoint = ContextCachingEndpoints() + tool_choice = {"mode": "ANY", "allowedFunctionNames": ["my_func"]} + optional_params = { + "tools": [{"function_declarations": [{"name": "my_func"}]}], + "tool_choice": tool_choice, + } + + with ( + patch.object(endpoint, "_get_token_and_url_context_caching", return_value=("token", "http://url")), + patch.object(endpoint, "check_cache", return_value=None), + patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt", return_value=True), + patch("litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.HTTPHandler") as mock_http, + ): + mock_resp = MagicMock() + mock_resp.json.return_value = {"name": "cached/123", "model": "gemini-1.5-pro"} + mock_client = MagicMock() + mock_client.post.return_value = mock_resp + mock_http.return_value = mock_client + + endpoint.check_and_create_cache( + messages=self._make_cached_messages(), + optional_params=optional_params, + api_key="test-key", + api_base=None, + model="gemini-1.5-pro", + client=None, + timeout=None, + logging_obj=MagicMock(), + custom_llm_provider="gemini", + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + + post_calls = mock_client.post.call_args_list + assert len(post_calls) == 1, "exactly one POST to the cache endpoint expected" + sent_body = post_calls[0].kwargs.get("json") or post_calls[0].args[1] + assert "toolConfig" in sent_body, ( + "toolConfig must be present in the cached-content request body" + ) + assert sent_body["toolConfig"] == tool_choice + + if __name__ == "__main__": pytest.main([__file__, "-v"])