scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-15 04:19:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 122

fix(vertex): preserve tool_choice through Gemini context-caching

Browse source 
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
This commit is contained in:
felipebridge 2026-05-28 11:37:50 -03:00
parent ea015332d8
commit eac9912e33
2 changed files with 133 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -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(

125
tests/test_litellm/llms/vertex_ai/context_caching/test_context_caching_ttl.py
View file

@ -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"])