scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 131

fix: count extra_body tools and cache_control in place of the direct ones

Browse source
This commit is contained in:
mateo-berri 2026-09-19 18:22:41 -07:00
parent b220178e9d
commit b0971ee0ba
2 changed files with 45 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -318,26 +318,22 @@ class AnthropicCacheControlHook(CustomPromptManagement):
A tool carries its mark at the top level (Anthropic shape) or under ``function``
(OpenAI shape). A top-level ``cache_control`` is Anthropic's automatic caching,
which places one breakpoint of its own on top of the explicit ones. Marks the
client sends through the ``extra_body`` envelope of ``request_kwargs`` reach the
wire too and count the same way. Callers pass only the tools whose mark reaches
the provider on their path.
which places one breakpoint of its own on top of the explicit ones. The
``extra_body`` envelope of ``request_kwargs`` is merged over the request on the
wire, so a ``tools`` or ``cache_control`` it carries replaces the direct value
and is counted in its place. Callers pass only the tools whose mark reaches the
provider on their path.
"""
extra_body: Final = (
_validated_object_mapping(AnthropicCacheControlHook._request_value(request_kwargs, "extra_body")) or {}
)
automatic_blocks: Final = sum(
1 for control in (cache_control, extra_body.get("cache_control")) if control is not None
)
tool_blocks: Final = sum(
1
for tool in (*(tools or ()), *(_validated_object_list(extra_body.get("tools")) or ()))
if _tool_carries_cache_breakpoint(tool)
)
wire_cache_control: Final = extra_body.get("cache_control", cache_control)
wire_tools: Final = _validated_object_list(extra_body["tools"]) if "tools" in extra_body else tools
tool_blocks: Final = sum(1 for tool in wire_tools or () if _tool_carries_cache_breakpoint(tool))
envelope_blocks: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(
_validated_object_list(extra_body.get("messages")) or (), extra_body.get("system")
)
return automatic_blocks + tool_blocks + envelope_blocks
return int(wire_cache_control is not None) + tool_blocks + envelope_blocks
@staticmethod
def _blocks_reserved_outside_messages(

36
tests/test_litellm/integrations/test_anthropic_cache_control_hook.py
View file

@ -2596,6 +2596,42 @@ class TestConfiguredInjectionPointsSurviveClientMarks:
_, result_sys = self._inject(self._marked_user_turns(3), kwargs)
assert result_sys == expected_system
@pytest.mark.parametrize(
"params,tools,marked_turns,injected",
[
({"extra_body": {"tools": [MARKED_TOOL_TOP_LEVEL]}}, [MARKED_TOOL_TOP_LEVEL], 2, 1),
({"extra_body": {"tools": [UNMARKED_TOOL]}}, [MARKED_TOOL_TOP_LEVEL], 3, 1),
({"extra_body": {"tools": [MARKED_TOOL_TOP_LEVEL]}}, [UNMARKED_TOOL], 3, 0),
({"extra_body": {"cache_control": EPHEMERAL}, "cache_control": EPHEMERAL}, None, 2, 1),
],
ids=["same_marked_tool_both_ways", "extra_body_unmarks", "extra_body_marks", "root_cache_control_both_ways"],
)
def test_chat_cap_counts_extra_body_fields_in_place_of_the_direct_ones(self, params, tools, marked_turns, injected):
"""``extra_body`` is merged over the request on the wire, so its ``tools`` and
``cache_control`` replace the direct ones rather than adding to them."""
messages = [{"role": "system", "content": "sys"}, *self._marked_user_turns(marked_turns)]
params = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED), **copy.deepcopy(params)}
self._seed(params, copy.deepcopy(messages), tools=tools)
processed = self._chat(params, copy.deepcopy(messages))
assert _count_cache_control(processed) == marked_turns + injected
@pytest.mark.parametrize(
"kwargs,tools,marked_turns,expected_system",
[
({"extra_body": {"tools": [MARKED_V1_TOOL]}}, [MARKED_V1_TOOL], 2, [{"type": "text", "text": "sys", "cache_control": EPHEMERAL}]),
({"extra_body": {"tools": [UNMARKED_V1_TOOL]}}, [MARKED_V1_TOOL], 3, [{"type": "text", "text": "sys", "cache_control": EPHEMERAL}]),
({"extra_body": {"tools": [MARKED_V1_TOOL]}}, [UNMARKED_V1_TOOL], 3, "sys"),
({"extra_body": {"cache_control": EPHEMERAL}, "cache_control": EPHEMERAL}, None, 2, [{"type": "text", "text": "sys", "cache_control": EPHEMERAL}]),
],
ids=["same_marked_tool_both_ways", "extra_body_unmarks", "extra_body_marks", "root_cache_control_both_ways"],
)
def test_v1_messages_cap_counts_extra_body_fields_in_place_of_the_direct_ones(
self, kwargs, tools, marked_turns, expected_system
):
kwargs = {"cache_control_injection_points": copy.deepcopy(self.CONFIGURED), **copy.deepcopy(kwargs)}
_, result_sys = self._inject(self._marked_user_turns(marked_turns), kwargs, tools=tools)
assert result_sys == expected_system
def test_v1_messages_automatic_defaults_stand_down_for_root_cache_control(self, monkeypatch):
monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True)
root_cache_control = {"type": "ephemeral"}