diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 4ab3e0d35cf..c54b42bf0ca 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -165,6 +165,25 @@ def _index_of_block_accepting_cache_control(content: list[object], on_a_tool_mes return None +def _message_accepts_cache_control(message: object) -> bool: + """Whether a cache_control marker written on this message reaches the provider. + + Reads the same block rule as `_safe_insert_cache_control_in_message`, so the two + cannot disagree about where a marker goes. A message whose content is empty -- + an assistant turn that said everything through ``tool_calls`` -- has nowhere to + put one. + """ + if not isinstance(message, dict): + return False + on_a_tool_message: Final = message.get("role") == "tool" + content: Final = message.get("content") + if isinstance(content, str): + return content != "" or on_a_tool_message + if isinstance(content, list): + return _index_of_block_accepting_cache_control(content, on_a_tool_message) is not None + return False + + # Set by a caller whose message list is not the one that goes upstream -- today the # Responses API layer, whose `instructions` only becomes a system message further down. # Tells this hook to hand role-targeted points to the pass holding the final messages @@ -471,28 +490,46 @@ class AnthropicCacheControlHook(CustomPromptManagement): else: targetted_index = _targetted_index - # Case 1: Target by specific index - if targetted_index is not None: - original_index: Final = targetted_index - if targetted_index < 0: - targetted_index += len(messages) + # The messages the point names. A point with a role counts its index within + # that role's turns, so {role: assistant, index: -1} is the last assistant + # turn rather than the last message. + targetted_role: Final = point.get("role", None) + candidates: Final = [ + index + for index, message in enumerate(messages) + if targetted_role is None or message.get("role") == targetted_role + ] - if 0 <= targetted_index < len(messages): - return [targetted_index] + # Case 1: Target by role alone + if targetted_index is None: + if targetted_role is None: + return [] + return [index for index in candidates if _message_accepts_cache_control(messages[index])] + # Case 2: Target by index, within the role the point named + original_index: Final = targetted_index + if targetted_index < 0: + targetted_index += len(candidates) + + if not 0 <= targetted_index < len(candidates): verbose_logger.warning( "AnthropicCacheControlHook: Provided index %s is out of bounds for message list of length %s. Targeted index was %s. Skipping cache control injection for this point.", original_index, - len(messages), + len(candidates), targetted_index, ) return [] - # Case 2: Target by role - targetted_role: Final = point.get("role", None) - if targetted_role is not None: - return [idx for idx, msg in enumerate(messages) if msg.get("role") == targetted_role] - + # The message it arrives at may have nowhere to write a marker -- an assistant + # turn that said everything through tool_calls is the common one -- so walk back + # to the nearest earlier turn that does. Stopping there would spend the point + # and write nothing. + position = targetted_index + while position >= 0: + index = candidates[position] + if _message_accepts_cache_control(messages[index]): + return [index] + position -= 1 return [] @staticmethod diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 3adf0e18640..c9af51e0392 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -3802,3 +3802,262 @@ async def test_anthropic_cache_control_hook_marks_an_empty_tool_result(monkeypat } ], } + + +def _anthropic_response_mock() -> MagicMock: + mock_response = MagicMock() + mock_response.json.return_value = { + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-5", + "content": [{"type": "text", "text": "Sure."}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 10, "output_tokens": 20}, + } + mock_response.status_code = 200 + return mock_response + + +async def _marked_messages(messages, points, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("ANTHROPIC_API_KEY", "fake_anthropic_key") + monkeypatch.setattr(litellm, "callbacks", [AnthropicCacheControlHook()]) + client = AsyncHTTPHandler() + with patch.object(client, "post", return_value=_anthropic_response_mock()) as mock_post: + await litellm.acompletion( + model="anthropic/claude-sonnet-4-5", + messages=messages, + cache_control_injection_points=points, + client=client, + ) + return mock_post.call_args.kwargs["json"]["messages"] + + +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_counts_the_index_within_the_role(monkeypatch: pytest.MonkeyPatch): + """ + {"role": "assistant", "index": -1} means the last assistant turn. Counting the index + over every message instead marks whatever message happens to be last, so the turn the + point named is not cached. + """ + marked = await _marked_messages( + [ + {"role": "user", "content": "What is 2 + 2?"}, + {"role": "assistant", "content": "The answer is 4."}, + {"role": "user", "content": "Why?"}, + ], + [{"location": "message", "role": "assistant", "index": -1}], + monkeypatch, + ) + + assert marked == [ + {"role": "user", "content": [{"type": "text", "text": "What is 2 + 2?"}]}, + { + "role": "assistant", + "content": [{"type": "text", "text": "The answer is 4.", "cache_control": {"type": "ephemeral"}}], + }, + {"role": "user", "content": [{"type": "text", "text": "Why?"}]}, + ] + + +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_walks_back_off_a_tool_call_turn(monkeypatch: pytest.MonkeyPatch): + """ + An assistant turn that said everything through tool_calls has no content to mark, and + it is the newest assistant turn at every step of an agent loop. Stopping there spends + the point and writes nothing, so the walk goes back to the turn before it. + """ + marked = await _marked_messages( + [ + {"role": "user", "content": "Open the file."}, + {"role": "assistant", "content": "Opening it now."}, + {"role": "user", "content": "Thanks."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "open_file", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "file contents"}, + ], + [{"location": "message", "role": "assistant", "index": -1}], + monkeypatch, + ) + + assert marked[1] == { + "role": "assistant", + "content": [{"type": "text", "text": "Opening it now.", "cache_control": {"type": "ephemeral"}}], + } + + +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_walk_back_stays_in_the_role(monkeypatch: pytest.MonkeyPatch): + """ + The walk stays inside the role the point named, so an empty last user turn marks an + earlier user turn rather than the assistant turn between them. + """ + marked = await _marked_messages( + [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": "first answer"}, + {"role": "user", "content": ""}, + ], + [{"location": "message", "role": "user", "index": -1}], + monkeypatch, + ) + + assert marked[0] == { + "role": "user", + "content": [{"type": "text", "text": "first question", "cache_control": {"type": "ephemeral"}}], + } + assert marked[1] == {"role": "assistant", "content": [{"type": "text", "text": "first answer"}]} + + +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_bounds_the_index_by_the_role(monkeypatch: pytest.MonkeyPatch): + """ + An index is in bounds when the role has that many turns, not when the message list + does. -3 against two assistant turns names no turn, so the point marks nothing. + """ + marked = await _marked_messages( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ], + [{"location": "message", "role": "assistant", "index": -3}], + monkeypatch, + ) + + assert marked == [ + {"role": "user", "content": [{"type": "text", "text": "u1"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "a1"}]}, + {"role": "user", "content": [{"type": "text", "text": "u2"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "a2"}]}, + ] + + +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_bounds_a_positive_index_by_the_role(monkeypatch: pytest.MonkeyPatch): + """ + Index 3 names a fourth turn of the role. Two assistant turns among five messages is + out of bounds, and bounding it by the message list instead reads past the end of the + role's turns. + """ + marked = await _marked_messages( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + {"role": "user", "content": "u3"}, + ], + [{"location": "message", "role": "assistant", "index": 3}], + monkeypatch, + ) + + assert marked == [ + {"role": "user", "content": [{"type": "text", "text": "u1"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "a1"}]}, + {"role": "user", "content": [{"type": "text", "text": "u2"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "a2"}]}, + {"role": "user", "content": [{"type": "text", "text": "u3"}]}, + ] + + +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_role_alone_skips_a_turn_with_nowhere_to_write( + monkeypatch: pytest.MonkeyPatch, +): + """ + A point naming a role and no index marks every turn of that role that accepts a + marker. An empty turn would spend a breakpoint on a block the conversion replaces. + """ + marked = await _marked_messages( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": ""}, + {"role": "assistant", "content": "a2"}, + {"role": "user", "content": "u3"}, + ], + [{"location": "message", "role": "user"}], + monkeypatch, + ) + + assert marked[0] == { + "role": "user", + "content": [{"type": "text", "text": "u1", "cache_control": {"type": "ephemeral"}}], + } + assert marked[2] == { + "role": "user", + "content": [{"type": "text", "text": "[System: Empty message content sanitised to satisfy protocol]"}], + } + assert marked[4] == { + "role": "user", + "content": [{"type": "text", "text": "u3", "cache_control": {"type": "ephemeral"}}], + } + + +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_marks_a_tool_turn_that_returned_nothing( + monkeypatch: pytest.MonkeyPatch, +): + """ + A tool that returned nothing is the newest turn of an agent loop. Its empty content + reaches the provider inside the tool_result, so the marker goes on it rather than + walking back and leaving the tool_use outside the cached prefix. + """ + marked = await _marked_messages( + [ + {"role": "user", "content": "Search for it."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "search", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": ""}, + ], + [{"location": "message", "index": 2}], + monkeypatch, + ) + + assert marked[-1] == { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "call_1", + "content": "", + "cache_control": {"type": "ephemeral"}, + } + ], + } + + +@pytest.mark.asyncio +async def test_anthropic_cache_control_hook_walks_back_off_a_thinking_only_turn(monkeypatch: pytest.MonkeyPatch): + """ + A turn whose only block is a thinking block has no block that accepts a marker, so + the walk goes back to the turn before it rather than spending the point there. + """ + marked = await _marked_messages( + [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + {"role": "assistant", "content": [{"type": "thinking", "thinking": "t", "signature": "s"}]}, + ], + [{"location": "message", "role": "assistant", "index": -1}], + monkeypatch, + ) + + assert marked[1] == { + "role": "assistant", + "content": [{"type": "text", "text": "a1", "cache_control": {"type": "ephemeral"}}], + } + assert marked[3] == {"role": "assistant", "content": [{"type": "thinking", "thinking": "t", "signature": "s"}]}