scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

fix(anthropic): count an injection point's index within the role it names

Browse source 
{"role": "assistant", "index": -1} read as "the last message", because the index
branch returned before the role was looked at. On a conversation ending with a user
turn it marked that turn; the assistant turn the point named went uncached.

The message it lands on may also have nowhere to write a marker. An assistant turn
that said everything through tool_calls has no content, and it is the newest
assistant turn at every step of an agent loop, so the point was spent and nothing
was written.

Count the index within the role's turns, and walk back from there to the nearest
earlier turn that accepts a marker. The walk stays inside the role the point named.
This commit is contained in:
Tan Nguyen 2026-09-23 12:31:31 +07:00
parent 0c630cef3b
commit 03e6a97639
2 changed files with 309 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

63
litellm/integrations/anthropic_cache_control_hook.py
View file

@ -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

259
tests/unit/integrations/test_anthropic_cache_control_hook.py
View file

@ -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"}]}