scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-27 01:22:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

fix: forward a tool_config point only while the cap has a slot left

Browse source
This commit is contained in:
mateo-berri 2026-09-19 06:28:16 -07:00
parent 171b33abfe
commit 7520925924
2 changed files with 142 additions and 37 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

67
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -209,14 +209,12 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# Create a deep copy of messages to avoid modifying the original list
processed_messages = copy.deepcopy(messages)
# Separate message-level and non-message-level injection points
message_points: Final[list[CacheControlMessageInjectionPoint]] = []
remaining_points: Final[list[CacheControlInjectionPoint]] = []
for point in injection_points:
if point.get("location") == "message":
message_points.append(cast(CacheControlMessageInjectionPoint, point))
else:
remaining_points.append(point)
message_points: Final = tuple(
cast(CacheControlMessageInjectionPoint, point)
for point in injection_points
if point.get("location") == "message"
)
remaining_points: Final = tuple(point for point in injection_points if point.get("location") != "message")
stamped_dialect: Final = injection_points[0].get("_litellm_openai_dialect")
openai_dialect: Final = (
@ -243,10 +241,9 @@ class AnthropicCacheControlHook(CustomPromptManagement):
else tuple(message_points)
)
stamped_external: Final = injection_points[0].get(EXTERNAL_BREAKPOINTS_STAMP)
external_breakpoints: Final = stamped_external if isinstance(stamped_external, int) else 0
reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages(
remaining_points,
stamped_external if isinstance(stamped_external, int) else 0,
openai_dialect,
remaining_points, external_breakpoints, openai_dialect
)
breakpoints_before: Final = AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages)
processed_messages = self._apply_message_injections(
@ -266,7 +263,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# `instructions`, which is only a system message once the bridge builds one. A later
# pass re-applies them safely: a target that already carries a mark is skipped and
# the census counts every mark on the wire, litellm's own included.
carried_points: Final[Sequence[CacheControlInjectionPoint]] = (*remaining_points, *carried_message_points)
carried_points: Final[Sequence[CacheControlInjectionPoint]] = (
*AnthropicCacheControlHook._points_with_a_slot_left(
remaining_points,
AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages) + external_breakpoints,
openai_dialect,
),
*carried_message_points,
)
if carried_points:
non_default_params["cache_control_injection_points"] = list(carried_points)
@ -331,6 +335,16 @@ class AnthropicCacheControlHook(CustomPromptManagement):
tool_config_blocks: Final = 1 if any(p.get("location") == "tool_config" for p in remaining_points) else 0
return external_breakpoints + tool_config_blocks
@staticmethod
def _points_with_a_slot_left(
remaining_points: Sequence[CacheControlInjectionPoint], breakpoints_on_wire: int, openai_dialect: bool
) -> tuple[CacheControlInjectionPoint, ...]:
"""A ``tool_config`` point becomes a cachePoint the Bedrock converse transform never
counts against the cap, so it is forwarded only while the wire still has a slot."""
if openai_dialect or breakpoints_on_wire < MAX_CACHE_CONTROL_BLOCKS:
return tuple(remaining_points)
return tuple(point for point in remaining_points if point.get("location") != "tool_config")
@staticmethod
def _apply_message_injections(
points: Sequence[CacheControlMessageInjectionPoint],
@ -529,19 +543,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
processed_messages: list[dict] = copy.deepcopy(messages)
processed_system = copy.deepcopy(system) if system is not None else None
message_points: Final[list[CacheControlMessageInjectionPoint]] = []
system_points: Final[list[CacheControlMessageInjectionPoint]] = []
remaining_points: Final[list[CacheControlInjectionPoint]] = []
for point in injection_points:
if point.get("location") == "message":
msg_point = cast(CacheControlMessageInjectionPoint, point)
if msg_point.get("role") == "system":
system_points.append(msg_point)
else:
message_points.append(msg_point)
else:
remaining_points.append(point)
role_points: Final = tuple(
cast(CacheControlMessageInjectionPoint, point)
for point in injection_points
if point.get("location") == "message"
)
system_points: Final = tuple(point for point in role_points if point.get("role") == "system")
message_points: Final = tuple(point for point in role_points if point.get("role") != "system")
remaining_points: Final = tuple(point for point in injection_points if point.get("location") != "message")
reserved_blocks: Final = AnthropicCacheControlHook._blocks_reserved_outside_messages(
remaining_points, external_breakpoints, openai_dialect
@ -581,8 +590,14 @@ class AnthropicCacheControlHook(CustomPromptManagement):
max_blocks=max_blocks - system_blocks,
openai_dialect=openai_dialect,
)
forwarded_points: Final = AnthropicCacheControlHook._points_with_a_slot_left(
remaining_points,
AnthropicCacheControlHook.count_request_cache_breakpoints(processed_messages, processed_system)
+ external_breakpoints,
openai_dialect,
)
return processed_messages, processed_system, remaining_points
return processed_messages, processed_system, list(forwarded_points)
@staticmethod
def _default_control() -> ChatCompletionCachedContent:

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

@ -1335,17 +1335,7 @@ async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point(mo
)
request_body = json.loads(mock_post.call_args.kwargs["data"])
cache_points = sum(
1 for block in request_body.get("system", []) if isinstance(block, dict) and "cachePoint" in block
)
for msg in request_body.get("messages", []):
content = msg.get("content", [])
if isinstance(content, list):
cache_points += sum(1 for block in content if isinstance(block, dict) and "cachePoint" in block)
for tool in request_body.get("toolConfig", {}).get("tools", []):
if isinstance(tool, dict) and "cachePoint" in tool:
cache_points += 1
cache_points = _count_converse_cache_points(request_body)
assert cache_points <= 4, (
f"Bedrock payload exceeded Anthropic's 4 cache_control block limit "
@ -1353,6 +1343,89 @@ async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point(mo
)
def _count_converse_cache_points(request_body: dict) -> int:
system_points = sum(
1 for block in request_body.get("system", []) if isinstance(block, dict) and "cachePoint" in block
)
message_points = sum(
1
for msg in request_body.get("messages", [])
if isinstance(msg.get("content"), list)
for block in msg["content"]
if isinstance(block, dict) and "cachePoint" in block
)
tool_points = sum(
1
for tool in request_body.get("toolConfig", {}).get("tools", [])
if isinstance(tool, dict) and "cachePoint" in tool
)
return system_points + message_points + tool_points
@pytest.mark.asyncio
async def test_cache_control_hook_bedrock_tool_config_point_stands_down_when_client_marks_fill_the_cap(
monkeypatch: pytest.MonkeyPatch,
):
"""The client's own four marks fill the cap, so the configured tool_config point must
not land as a fifth cachePoint in the converse payload."""
with patch.dict(
os.environ,
{
"AWS_ACCESS_KEY_ID": "fake_access_key_id",
"AWS_SECRET_ACCESS_KEY": "fake_secret_access_key",
"AWS_REGION_NAME": "us-east-1",
},
):
monkeypatch.setattr(litellm, "callbacks", [AnthropicCacheControlHook()])
mock_response = MagicMock()
mock_response.json.return_value = {
"output": {"message": {"role": "assistant", "content": "ok"}},
"stopReason": "end_turn",
"usage": {"inputTokens": 100, "outputTokens": 4, "totalTokens": 104},
}
mock_response.status_code = 200
client = AsyncHTTPHandler()
with patch.object(client, "post", return_value=mock_response) as mock_post:
marked = {"type": "ephemeral"}
messages = [
{"role": "system", "content": [{"type": "text", "text": "sys", "cache_control": marked}]},
*(
{"role": "user", "content": [{"type": "text", "text": f"turn {i}", "cache_control": marked}]}
for i in range(3)
),
{"role": "user", "content": "What is the weather?"},
]
await litellm.acompletion(
model="bedrock/us.anthropic.claude-opus-4-6-v1:0",
messages=messages,
max_tokens=32,
tools=[
{
"type": "function",
"function": {
"name": "get_weather",
"description": "Get weather for a location",
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
},
}
],
cache_control_injection_points=[{"location": "tool_config"}],
client=client,
)
request_body = json.loads(mock_post.call_args.kwargs["data"])
assert _count_converse_cache_points(request_body) == 4
assert not any("cachePoint" in tool for tool in request_body["toolConfig"]["tools"])
class TestApplyToAnthropicMessagesRequest:
"""Tests for apply_to_anthropic_messages_request (v1/messages cache control)."""
@ -2091,6 +2164,7 @@ class TestConfiguredInjectionPointsSurviveClientMarks:
CONFIGURED = [{"location": "message", "role": "system"}]
TAIL_POINT = [{"location": "message", "index": -1}]
TOOL_CONFIG_POINT = [{"location": "tool_config"}]
EPHEMERAL = {"type": "ephemeral"}
CLEAN_MESSAGES: List[AllMessageValues] = [
@ -2220,6 +2294,22 @@ class TestConfiguredInjectionPointsSurviveClientMarks:
processed = self._chat(params, copy.deepcopy(messages))
assert _count_cache_control(processed) == 4
@pytest.mark.parametrize("marked_turns,forwarded", [(3, ["tool_config"]), (4, [])], ids=["slot_left", "cap_full"])
def test_chat_forwards_tool_config_point_only_while_a_slot_is_left(self, marked_turns, forwarded):
"""A forwarded tool_config point becomes a Bedrock cachePoint unconditionally, so
it stands down once the client's own marks fill the cap."""
messages = [{"role": "system", "content": "sys"}, *self._marked_user_turns(marked_turns)]
params = {"cache_control_injection_points": copy.deepcopy(self.TOOL_CONFIG_POINT)}
self._seed(params, copy.deepcopy(messages), tools=[self.UNMARKED_TOOL])
self._chat(params, copy.deepcopy(messages))
assert [p["location"] for p in params.get("cache_control_injection_points", [])] == forwarded
@pytest.mark.parametrize("marked_turns,forwarded", [(3, ["tool_config"]), (4, [])], ids=["slot_left", "cap_full"])
def test_v1_messages_forwards_tool_config_point_only_while_a_slot_is_left(self, marked_turns, forwarded):
kwargs = {"cache_control_injection_points": copy.deepcopy(self.TOOL_CONFIG_POINT)}
self._inject(self._marked_user_turns(marked_turns), kwargs, tools=[self.UNMARKED_V1_TOOL])
assert [p["location"] for p in kwargs.get("cache_control_injection_points", [])] == forwarded
@pytest.mark.parametrize("marked_turns,injected", [(2, 1), (3, 0)])
def test_chat_root_cache_control_reserves_a_slot(self, marked_turns, injected):
"""Anthropic's automatic caching (a top-level ``cache_control``) places one