diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 6f90acade10..cff7d23935c 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -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: diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 1dfd9cf619b..2723526ae6b 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -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