diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 213622cb43a2..296bfb6fc85a 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -27,6 +27,11 @@ LiteLLMLoggingObj = Any +# Anthropic (and Bedrock Claude) reject requests with more than 4 cache_control +# breakpoints: "A maximum of 4 blocks with cache_control may be provided." +MAX_CACHE_CONTROL_BLOCKS = 4 + + class AnthropicCacheControlHook(CustomPromptManagement): def get_chat_completion_prompt( self, @@ -61,16 +66,30 @@ def get_chat_completion_prompt( processed_messages = copy.deepcopy(messages) # Separate message-level and non-message-level injection points - remaining_points = [] + message_points: List[CacheControlMessageInjectionPoint] = [] + remaining_points: List[CacheControlInjectionPoint] = [] for point in injection_points: if point.get("location") == "message": - point = cast(CacheControlMessageInjectionPoint, point) - processed_messages = self._process_message_injection( - point=point, messages=processed_messages - ) + message_points.append(cast(CacheControlMessageInjectionPoint, point)) else: remaining_points.append(point) + # Non-message points (currently Bedrock tool_config) are handled in the + # provider transform, where each tool_config point appends at most one + # cachePoint to the tools. That block also counts toward Anthropic's + # limit, so reserve a slot for it here to leave room. + reserved_blocks = ( + 1 + if any(p.get("location") == "tool_config" for p in remaining_points) + else 0 + ) + + processed_messages = self._apply_message_injections( + points=message_points, + messages=processed_messages, + max_blocks=MAX_CACHE_CONTROL_BLOCKS - reserved_blocks, + ) + # Pass through non-message injection points for provider-specific handling if remaining_points: non_default_params["cache_control_injection_points"] = remaining_points @@ -78,14 +97,71 @@ def get_chat_completion_prompt( return model, processed_messages, non_default_params @staticmethod - def _process_message_injection( - point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues] + def _apply_message_injections( + points: List[CacheControlMessageInjectionPoint], + messages: List[AllMessageValues], + max_blocks: int, ) -> List[AllMessageValues]: - """Process message-level cache control injection.""" - control: ChatCompletionCachedContent = point.get( - "control", None - ) or ChatCompletionCachedContent(type="ephemeral") + """Apply message-level cache control injection points in order. + + Anthropic allows at most ``MAX_CACHE_CONTROL_BLOCKS`` cache_control + breakpoints per request. Client-supplied breakpoints count toward that + limit, so we never inject onto a message that already carries + cache_control (preserving the client's TTL) and we stop injecting once + ``max_blocks`` is reached. Injection points are honored in config order, + so earlier points win when slots are scarce. + """ + used_blocks = sum( + AnthropicCacheControlHook._count_cache_control_blocks(msg) + for msg in messages + ) + + limit_reached = False + for point in points: + if used_blocks >= max_blocks: + limit_reached = True + break + + control: ChatCompletionCachedContent = point.get( + "control", None + ) or ChatCompletionCachedContent(type="ephemeral") + + for target_index in AnthropicCacheControlHook._resolve_target_indices( + point=point, messages=messages + ): + if used_blocks >= max_blocks: + limit_reached = True + break + + if AnthropicCacheControlHook._message_has_cache_control( + messages[target_index] + ): + # Client already marked this message; don't overwrite it. + continue + + messages[target_index] = ( + AnthropicCacheControlHook._safe_insert_cache_control_in_message( + messages[target_index], control + ) + ) + used_blocks += 1 + + if limit_reached: + break + + if limit_reached: + verbose_logger.warning( + f"AnthropicCacheControlHook: Reached the Anthropic limit of " + f"{MAX_CACHE_CONTROL_BLOCKS} cache_control blocks. Skipping further injection." + ) + return messages + + @staticmethod + def _resolve_target_indices( + point: CacheControlMessageInjectionPoint, messages: List[AllMessageValues] + ) -> List[int]: + """Resolve which message indices an injection point targets.""" _targetted_index: Optional[Union[int, str]] = point.get("index", None) targetted_index: Optional[int] = None if isinstance(_targetted_index, str): @@ -96,36 +172,49 @@ def _process_message_injection( else: targetted_index = _targetted_index - targetted_role = point.get("role", None) - # Case 1: Target by specific index if targetted_index is not None: original_index = targetted_index - # Handle negative indices (convert to positive) if targetted_index < 0: targetted_index += len(messages) if 0 <= targetted_index < len(messages): - messages[targetted_index] = ( - AnthropicCacheControlHook._safe_insert_cache_control_in_message( - messages[targetted_index], control - ) - ) - else: - verbose_logger.warning( - f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. " - f"Targeted index was {targetted_index}. Skipping cache control injection for this point." - ) + return [targetted_index] + + verbose_logger.warning( + f"AnthropicCacheControlHook: Provided index {original_index} is out of bounds for message list of length {len(messages)}. " + f"Targeted index was {targetted_index}. Skipping cache control injection for this point." + ) + return [] + # Case 2: Target by role - elif targetted_role is not None: - for msg in messages: - if msg.get("role") == targetted_role: - msg = ( - AnthropicCacheControlHook._safe_insert_cache_control_in_message( - message=msg, control=control - ) - ) - return messages + targetted_role = point.get("role", None) + if targetted_role is not None: + return [ + idx + for idx, msg in enumerate(messages) + if msg.get("role") == targetted_role + ] + + return [] + + @staticmethod + def _count_cache_control_blocks(message: AllMessageValues) -> int: + """Count cache_control breakpoints on a message (message + content level).""" + count = 0 + if message.get("cache_control") is not None: + count += 1 + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("cache_control") is not None: + count += 1 + return count + + @staticmethod + def _message_has_cache_control(message: AllMessageValues) -> bool: + """Return True if the message already carries any cache_control.""" + return AnthropicCacheControlHook._count_cache_control_blocks(message) > 0 @staticmethod def _safe_insert_cache_control_in_message( 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 1a4d03528e7c..6afe5efc54d4 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1087,3 +1087,357 @@ async def test_anthropic_cache_control_hook_string_negative_index(): f"Expected cachePoint in last message content, got: {last_message_content}. " "String index '-1' was not parsed correctly (str.isdigit() returns False for negative strings)." ) + + +def _count_cache_control(messages: List[AllMessageValues]) -> int: + """Count cache_control breakpoints across messages (message + content level).""" + count = 0 + for message in messages: + if message.get("cache_control") is not None: + count += 1 + content = message.get("content") + if isinstance(content, list): + for block in content: + if isinstance(block, dict) and block.get("cache_control") is not None: + count += 1 + return count + + +def _build_injection_points(): + return [ + { + "location": "message", + "role": "system", + "control": {"type": "ephemeral", "ttl": "1h"}, + }, + { + "location": "message", + "index": -1, + "control": {"type": "ephemeral", "ttl": "5m"}, + }, + ] + + +def test_cache_control_hook_caps_at_four_blocks_with_client_cache_control(): + """Regression for LIT-3667 / Anthropic 'A maximum of 4 blocks ... Found 5'. + + A Hermes-style request already carries 4 client cache_control breakpoints on + its system messages. With both auto-inject points configured the hook must + NOT add a 5th breakpoint, and must NOT overwrite the client's existing + breakpoints (TTL must be preserved). + """ + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": f"System block {i}", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + for i in range(4) + ] + messages.append({"role": "user", "content": "hello"}) + + _, processed, _ = hook.get_chat_completion_prompt( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + non_default_params={ + "cache_control_injection_points": _build_injection_points() + }, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + assert ( + _count_cache_control(processed) == 4 + ), "Hook must cap cache_control at Anthropic's limit of 4 blocks" + + # Client TTL on system blocks must be preserved (not overwritten by config). + for i in range(4): + assert processed[i]["content"][-1]["cache_control"] == { + "type": "ephemeral", + "ttl": "1h", + } + + # The last (user) message must not receive a 5th breakpoint. + user_message = processed[-1] + assert user_message.get("cache_control") is None + user_content = user_message.get("content") + if isinstance(user_content, list): + assert all( + block.get("cache_control") is None + for block in user_content + if isinstance(block, dict) + ) + + +def test_cache_control_hook_caps_at_four_blocks_without_client_cache_control(): + """Four plain system messages + role:system + index:-1 must stay at 4 blocks. + + role:system fills all four slots, so the index:-1 point is skipped. + """ + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + {"role": "system", "content": f"System {i}"} for i in range(4) + ] + messages.append({"role": "user", "content": "hello"}) + + _, processed, _ = hook.get_chat_completion_prompt( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + non_default_params={ + "cache_control_injection_points": _build_injection_points() + }, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + assert _count_cache_control(processed) == 4 + # All four system messages cached; user message skipped (limit reached). + assert all(processed[i].get("cache_control") is not None for i in range(4)) + assert processed[-1].get("cache_control") is None + + +def test_cache_control_hook_does_not_overwrite_existing_cache_control(): + """If a targeted message already has client cache_control, do not inject.""" + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "Cached by client", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + }, + {"role": "user", "content": "hello"}, + ] + + _, processed, _ = hook.get_chat_completion_prompt( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + # Target the already-cached system message with a different TTL. + non_default_params={ + "cache_control_injection_points": [ + { + "location": "message", + "index": 0, + "control": {"type": "ephemeral", "ttl": "5m"}, + } + ] + }, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + # Client's 1h TTL must be preserved, not replaced by the config's 5m. + assert processed[0]["content"][-1]["cache_control"] == { + "type": "ephemeral", + "ttl": "1h", + } + assert _count_cache_control(processed) == 1 + + +@pytest.mark.asyncio +async def test_cache_control_hook_bedrock_payload_caps_cachepoints_at_four(): + """End-to-end: outgoing Bedrock payload must not exceed 4 cachePoint blocks. + + Reproduces the customer report where 4 client cache_control system blocks + plus auto-inject produced 5 cachePoint blocks and Bedrock returned 400. + """ + 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", + }, + ): + 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: + messages = [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": f"System block {i}", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + for i in range(4) + ] + messages.append({"role": "user", "content": "hello"}) + + await litellm.acompletion( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + max_tokens=32, + cache_control_injection_points=_build_injection_points(), + client=client, + ) + + 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 + ) + + assert cache_points <= 4, ( + f"Bedrock payload exceeded Anthropic's 4 cache_control block limit: " + f"found {cache_points} cachePoint blocks" + ) + + +def test_cache_control_hook_reserves_slot_for_tool_config_point(): + """A tool_config injection point consumes one of the 4 slots downstream. + + With role:system targeting 4 system messages plus a tool_config point, the + hook must inject at most 3 message-level blocks so the tool_config cachePoint + appended by the Bedrock transform keeps the total at 4, not 5. + """ + hook = AnthropicCacheControlHook() + + messages: List[AllMessageValues] = [ + {"role": "system", "content": f"System {i}"} for i in range(4) + ] + messages.append({"role": "user", "content": "hello"}) + + _, processed, non_default_params = hook.get_chat_completion_prompt( + model="bedrock/us.anthropic.claude-opus-4-6-v1:0", + messages=messages, + non_default_params={ + "cache_control_injection_points": [ + { + "location": "message", + "role": "system", + "control": {"type": "ephemeral", "ttl": "1h"}, + }, + {"location": "tool_config"}, + ] + }, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, + ) + + assert _count_cache_control(processed) == 3 + # The tool_config point is passed through for the provider transform. + assert non_default_params["cache_control_injection_points"] == [ + {"location": "tool_config"} + ] + + +@pytest.mark.asyncio +async def test_cache_control_hook_bedrock_payload_caps_with_tool_config_point(): + """End-to-end: message + tool_config injection must not exceed 4 cachePoints.""" + 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", + }, + ): + 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: + messages = [ + {"role": "system", "content": f"System block {i}"} for i in range(4) + ] + messages.append({"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": "message", + "role": "system", + "control": {"type": "ephemeral", "ttl": "1h"}, + }, + {"location": "tool_config"}, + ], + client=client, + ) + + 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 + + assert cache_points <= 4, ( + f"Bedrock payload exceeded Anthropic's 4 cache_control block limit " + f"when mixing message and tool_config injection: found {cache_points}" + )