From ab7883fd578b70ae377c2269f30fc527ed3ead2f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 13 Oct 2026 22:04:52 -0700 Subject: [PATCH] fix(bedrock): count Converse cache breakpoints the way Converse renders them Converse emits one cachePoint per tool message and ignores a cache_control under a tool's function, so the Claude 4-check now counts sites per wire and a dropped tool-message slot clears every mark on that message. The Converse breakpoint rename, prompt_cache_options check and cap now apply to Claude models only, so Nova and other Converse models keep forwarding prompt_cache_options as before. Also freezes the e2e PromptCacheBreakpoint model and covers a block carrying both keys on every Claude route. --- .../chat/prompt_cache_breakpoints.py | 38 ++++++-- .../bedrock/chat/converse_transformation.py | 12 +-- .../e2e/llm_translation/test_cache_control.py | 27 +++++- tests/e2e/models.py | 2 +- ...laude_chat_prompt_cache_breakpoint_wire.py | 35 +++++++ .../chat/test_prompt_cache_breakpoints.py | 91 +++++++++++++++++++ .../chat/test_converse_transformation.py | 82 +++++++++++++++++ 7 files changed, 270 insertions(+), 17 deletions(-) diff --git a/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py b/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py index 94ac9e5aaab..188f85e1855 100644 --- a/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py +++ b/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py @@ -13,6 +13,7 @@ from litellm._logging import verbose_logger from litellm.constants import MAX_CACHE_CONTROL_BLOCKS from litellm.exceptions import BadRequestError, UnsupportedParamsError from litellm.litellm_core_utils.prompt_templates.common_utils import ( + message_fields, message_without_cache_control_marks, messages_with_prompt_cache_breakpoints_as_cache_control, tools_with_prompt_cache_breakpoints_as_cache_control, @@ -20,7 +21,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.types.llms.openai import AllMessageValues PROMPT_CACHE_BREAKPOINT_ON_TOOLS: Final = "_litellm_prompt_cache_breakpoint_on_tools" -_Site: TypeAlias = tuple[Literal["tool", "message", "part", "tool_call"], int, int] +_Site: TypeAlias = tuple[Literal["tool", "message", "part", "tool_call", "tool_message"], int, int] +ClaudeWire: TypeAlias = Literal["anthropic", "bedrock_converse"] @dataclass(frozen=True, slots=True) @@ -75,11 +77,24 @@ def _message_sites(index: int, message: object, key: str) -> Iterator[_Site]: yield ("tool_call", index, position) -def _sites(messages: Sequence[object], tools: Sequence[object], key: str) -> tuple[_Site, ...]: - """Every block carrying ``key`` in wire order: tool definitions first, then the messages as listed.""" +def _converse_message_sites(index: int, message: object, key: str) -> Iterator[_Site]: + """Converse emits a single ``cachePoint`` per tool message, however many of its blocks carry a mark.""" + fields: Final = _as_mapping(message) + if fields is None or fields.get("role") != "tool": + yield from _message_sites(index, message, key) + return + if fields.get(key) is not None or any(_carries(part, key) for part in _as_blocks(fields.get("content")) or ()): + yield ("tool_message", index, -1) + + +def _sites(messages: Sequence[object], tools: Sequence[object], key: str, wire: ClaudeWire) -> tuple[_Site, ...]: + """Every mark of ``key`` the ``wire`` renders, in wire order: tool definitions first, then the messages as + listed. Converse reads only a tool definition's top-level mark.""" + tool_carries: Final = _carries if wire == "bedrock_converse" else _tool_carries + message_sites: Final = _converse_message_sites if wire == "bedrock_converse" else _message_sites return ( - *(("tool", index, -1) for index, tool in enumerate(tools) if _tool_carries(tool, key)), - *chain.from_iterable(_message_sites(index, message, key) for index, message in enumerate(messages)), + *(("tool", index, -1) for index, tool in enumerate(tools) if tool_carries(tool, key)), + *chain.from_iterable(message_sites(index, message, key) for index, message in enumerate(messages)), ) @@ -105,6 +120,11 @@ def _dropped_positions(dropped: frozenset[_Site], kind: str, index: int) -> froz def _message_without(index: int, message: AllMessageValues, dropped: frozenset[_Site]) -> AllMessageValues: + if ("tool_message", index, -1) in dropped: + content: Final = _as_blocks(message_fields(message).get("content")) or () + return message_without_cache_control_marks( + message, message_mark=True, parts=frozenset(range(len(content))), tool_calls=frozenset() + ) return message_without_cache_control_marks( message, message_mark=("message", index, -1) in dropped, @@ -120,6 +140,7 @@ def claude_cache_marks( automatic_breakpoint: bool, drop_params: bool, tools_carried_breakpoint: bool = False, + wire: ClaudeWire = "anthropic", ) -> ClaudeCacheMarks | CacheBreakpointCapExceeded: """The request's ``prompt_cache_breakpoint`` marks renamed to ``cache_control`` and, when it carried any, held to Claude's cap: ``drop_params`` keeps the latest four in wire order, otherwise the excess is reported. A request @@ -127,13 +148,14 @@ def claude_cache_marks( ``automatic_breakpoint`` is the request-level ``cache_control`` (Anthropic's automatic caching), which takes one slot of its own and is never dropped. ``tools_carried_breakpoint`` covers tools already mapped to Anthropic's - shape, where the mark was renamed before this census ran.""" + shape, where the mark was renamed before this census ran. ``wire`` is the body the marks are rendered into, so + the count matches the breakpoints the provider receives.""" listed_tools: Final = _as_blocks(tools) or () - if not tools_carried_breakpoint and not _sites(messages, listed_tools, "prompt_cache_breakpoint"): + if not tools_carried_breakpoint and not _sites(messages, listed_tools, "prompt_cache_breakpoint", "anthropic"): return ClaudeCacheMarks(tuple(messages), tuple(listed_tools)) renamed_messages: Final = messages_with_prompt_cache_breakpoints_as_cache_control(messages) renamed_tools: Final = tools_with_prompt_cache_breakpoints_as_cache_control(listed_tools) - sites: Final = _sites(renamed_messages, renamed_tools, "cache_control") + sites: Final = _sites(renamed_messages, renamed_tools, "cache_control", wire) found: Final = len(sites) + int(automatic_breakpoint) if found <= MAX_CACHE_CONTROL_BLOCKS: return ClaudeCacheMarks(renamed_messages, renamed_tools) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index b156e6f7556..e1a38d0ebe9 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -2025,13 +2025,12 @@ class AmazonConverseConfig(BaseConfig): optional_params: dict[str, object], # mutable-ok: the request body is built from this out-param litellm_params: Mapping[str, object], ) -> tuple[AllMessageValues, ...]: - """OpenAI ``prompt_cache_breakpoint`` marks as ``cache_control`` so the factory emits ``cachePoint`` blocks, - held to the provider's cap, with ``prompt_cache_options`` rejected or dropped. Models that never emit - ``cachePoint`` blocks keep their marks as written, but a Claude model among them still has - ``prompt_cache_options`` rejected or dropped.""" - emits_cache_points: Final = bedrock_model_accepts_cache_points(model) - if not emits_cache_points and BaseAWSLLM.get_bedrock_invoke_provider(model) != "anthropic": + """On Claude, OpenAI ``prompt_cache_breakpoint`` marks as ``cache_control`` so the factory emits + ``cachePoint`` blocks, held to Claude's cap, with ``prompt_cache_options`` rejected or dropped. A Claude model + that never emits ``cachePoint`` blocks keeps its marks as written. Other models are left as they are.""" + if BaseAWSLLM.get_bedrock_invoke_provider(model) != "anthropic": return tuple(messages) + emits_cache_points: Final = bedrock_model_accepts_cache_points(model) drop_params: Final = litellm.drop_params or litellm_params.get("drop_params") is True reject_prompt_cache_options( model=model, options=optional_params.pop("prompt_cache_options", None), drop_params=drop_params @@ -2044,6 +2043,7 @@ class AmazonConverseConfig(BaseConfig): tools=optional_params.get("tools"), automatic_breakpoint=optional_params.get("cache_control") is not None, drop_params=drop_params, + wire="bedrock_converse", ), model=model, llm_provider=self.custom_llm_provider or "bedrock", diff --git a/tests/e2e/llm_translation/test_cache_control.py b/tests/e2e/llm_translation/test_cache_control.py index 9bc3e7dc5b2..921136b6e97 100644 --- a/tests/e2e/llm_translation/test_cache_control.py +++ b/tests/e2e/llm_translation/test_cache_control.py @@ -33,8 +33,9 @@ Each case asserts the feature actually happened, not just a 200. Coverage matrix - Claude (claude-haiku-5-5 on anthropic, Bedrock Converse and Bedrock Invoke, claude-sonnet-5-5 on Vertex, which serves no haiku-5-5): OpenAI's ``prompt_cache_breakpoint`` on a system part, a string - system message, a user part, a tool result, a tool call or a tool definition - must write the cache on the first call and read it on the second. On + system message, a user part, a tool result, a tool call or a tool definition, + or beside a ``cache_control`` on the same system part, must write the cache + on the first call and read it on the second. On anthropic and Converse, a 30m breakpoint, five breakpoints and ``prompt_cache_options`` under drop_params must still cache. @@ -780,3 +781,25 @@ class TestPromptCacheBreakpointOnClaude: ], ) _assert_claude_writes_then_reads_cache(client, scoped_key, body) + + @pytest.mark.parametrize( + "route", [_claude_route_param(route) for route in ("anthropic", "bedrock_converse", "bedrock_invoke", "vertex")] + ) + def test_a_block_carrying_both_keys_still_caches( + self, client: PassthroughClient, resources: ResourceManager, scoped_key: str, route: ClaudeRoute + ) -> None: + model: Final = _register_claude(client, resources, route) + both_keys: Final = TextContentPart( + text=_cacheable_prefix(), + cache_control=CacheControl(), + prompt_cache_breakpoint=PromptCacheBreakpoint(ttl="30m"), + ) + body: Final = ChatBody( + model=model, + max_tokens=64, + messages=[ + ChatMessage(role="system", content=[both_keys]), + ChatMessage(role="user", content="Reply with one word."), + ], + ) + _assert_claude_writes_then_reads_cache(client, scoped_key, body) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 1635ab2e890..434a771e476 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -216,7 +216,7 @@ class ImageUrl(BaseModel): url: str -class PromptCacheBreakpoint(BaseModel): +class PromptCacheBreakpoint(BaseModel, frozen=True): """OpenAI's explicit cache marker on a content part, message, tool definition or tool call.""" mode: Literal["explicit"] = "explicit" diff --git a/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_wire.py index 611491454a0..0b5ee90083f 100644 --- a/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_wire.py @@ -256,6 +256,22 @@ def test_a_client_cache_control_on_the_same_block_wins_over_the_breakpoint(gatew assert b"prompt_cache_breakpoint" not in request.body, request.body +@pytest.mark.parametrize("provider", _PROVIDERS) +def test_a_block_carrying_both_keys_keeps_one_cache_mark_and_loses_the_breakpoint( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + marks: Final = _marked(("user_part",), {"cache_control": EPHEMERAL, "prompt_cache_breakpoint": _EXPLICIT_30M}) + with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, provider) + status, _, text = post_chat(gateway, _body(model, marker, marks)) + assert status == 200, text + request: Final = _only_request(wire) + assert _labels(provider, request) == [_expected(provider, "user_part", marker)] + assert b"prompt_cache_breakpoint" not in request.body, request.body + assert b"30m" not in request.body, request.body + + @pytest.mark.parametrize("source", ("deployment", "request")) @pytest.mark.parametrize("provider", _PROVIDERS) def test_five_breakpoints_keep_the_latest_four_under_drop_params( @@ -287,6 +303,25 @@ def test_five_breakpoints_are_a_400_before_any_provider_call_without_drop_params assert received == (), [request.target for request in received] +def test_a_tool_result_with_two_marked_parts_is_one_converse_cache_point_within_the_cap(gateway: Gateway) -> None: + marker: Final = new_marker() + marked_tool_result: Final[dict[str, JsonValue]] = { + "role": "tool", + "tool_call_id": call_id(_CITY), + "content": [{"type": "text", "text": text, **_BREAKPOINT} for text in (_TOOL_RESULT_TEXT, "warm.")], + } + messages: Final = _messages(marker, _marked(("system", "user_part", "string_message"))) + with wire_server(bedrock_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(gateway, scenario, wire, "bedrock_converse") + status, _, text = post_chat(gateway, chat_body(model, [*messages[:3], marked_tool_result, *messages[4:]])) + assert status == 200, text + request: Final = _only_request(wire) + assert bedrock_labels(request) == [ + _expected("bedrock_converse", site, marker) + for site in ("system", "user_part", "tool_result_part", "string_message") + ] + + def test_injection_points_stand_down_behind_the_latest_four_breakpoints(gateway: Gateway) -> None: marker: Final = new_marker() with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: diff --git a/tests/unit/llms/anthropic/chat/test_prompt_cache_breakpoints.py b/tests/unit/llms/anthropic/chat/test_prompt_cache_breakpoints.py index b77c25a9e87..dd9d613d5b7 100644 --- a/tests/unit/llms/anthropic/chat/test_prompt_cache_breakpoints.py +++ b/tests/unit/llms/anthropic/chat/test_prompt_cache_breakpoints.py @@ -171,6 +171,97 @@ def test_tools_already_mapped_to_the_anthropic_shape_still_arm_the_cap() -> None ) +def _tool_result_with_two_marked_parts() -> list[dict[str, object]]: + return [ + {"role": "system", "content": [_text("rules", prompt_cache_breakpoint=EXPLICIT)]}, + {"role": "user", "content": [_text("hi", prompt_cache_breakpoint=EXPLICIT)]}, + {"role": "assistant", "content": "", "tool_calls": [_call("call_1")]}, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + _text("sunny", prompt_cache_breakpoint=EXPLICIT), + _text("warm", prompt_cache_breakpoint=EXPLICIT), + ], + }, + {"role": "user", "content": "done", "prompt_cache_breakpoint": EXPLICIT}, + ] + + +def test_a_tool_result_with_two_marked_parts_takes_one_converse_slot_and_two_anthropic_slots() -> None: + messages: Final = _tool_result_with_two_marked_parts() + + converse: Final = claude_cache_marks( + messages, None, automatic_breakpoint=False, drop_params=False, wire="bedrock_converse" + ) + + assert converse == ClaudeCacheMarks( + messages=( + {"role": "system", "content": [_text("rules", cache_control=EPHEMERAL)]}, + {"role": "user", "content": [_text("hi", cache_control=EPHEMERAL)]}, + {"role": "assistant", "content": "", "tool_calls": [_call("call_1")]}, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [_text("sunny", cache_control=EPHEMERAL), _text("warm", cache_control=EPHEMERAL)], + }, + {"role": "user", "content": "done", "cache_control": EPHEMERAL}, + ), + tools=(), + ) + assert claude_cache_marks(messages, None, automatic_breakpoint=False, drop_params=False) == ( + CacheBreakpointCapExceeded(found=5) + ) + + +def test_dropping_a_converse_tool_result_slot_clears_every_mark_on_that_message() -> None: + messages: Final = [ + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + _text("sunny", prompt_cache_breakpoint=EXPLICIT), + _text("warm", prompt_cache_breakpoint=EXPLICIT), + ], + "prompt_cache_breakpoint": EXPLICIT, + }, + {"role": "user", "content": [_text(f"part {index}", prompt_cache_breakpoint=EXPLICIT) for index in range(4)]}, + ] + + marks: Final = claude_cache_marks( + messages, None, automatic_breakpoint=False, drop_params=True, wire="bedrock_converse" + ) + + assert marks == ClaudeCacheMarks( + messages=( + {"role": "tool", "tool_call_id": "call_1", "content": [_text("sunny"), _text("warm")]}, + {"role": "user", "content": [_text(f"part {index}", cache_control=EPHEMERAL) for index in range(4)]}, + ), + tools=(), + ) + + +def test_a_function_level_cache_control_takes_no_converse_slot() -> None: + tools: Final = [{"type": "function", "function": {"name": "lookup", "parameters": {}, "cache_control": EPHEMERAL}}] + messages: Final = [ + {"role": "user", "content": [_text(f"part {index}", prompt_cache_breakpoint=EXPLICIT) for index in range(4)]} + ] + + converse: Final = claude_cache_marks( + messages, tools, automatic_breakpoint=False, drop_params=False, wire="bedrock_converse" + ) + + assert converse == ClaudeCacheMarks( + messages=( + {"role": "user", "content": [_text(f"part {index}", cache_control=EPHEMERAL) for index in range(4)]}, + ), + tools=(tools[0],), + ) + assert claude_cache_marks(messages, tools, automatic_breakpoint=False, drop_params=False) == ( + CacheBreakpointCapExceeded(found=5) + ) + + def test_the_cap_error_is_a_public_bad_request_naming_drop_params() -> None: with pytest.raises(litellm.BadRequestError) as caught: resolved_claude_cache_marks(CacheBreakpointCapExceeded(found=6), model=MODEL, llm_provider="anthropic") diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 298d17657b7..9927130935e 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -10362,3 +10362,85 @@ def test_five_native_cache_control_marks_are_forwarded_as_written_past_the_cap() "content": [block for index in range(5) for block in ({"text": f"part {index}"}, _CACHE_POINT)], } ] + + +async def _breakpoint_converse_request( + sync_mode: bool, model: str, messages: list[dict[str, object]], **params: dict[str, object] +) -> dict[str, object]: + config: Final = AmazonConverseConfig() + optional_params: Final = params.get("optional_params", {}) + litellm_params: Final = params.get("litellm_params", {}) + if sync_mode: + return config._transform_request( + model=model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, headers={} + ) + return await config._async_transform_request( + model=model, messages=messages, optional_params=optional_params, litellm_params=litellm_params, headers={} + ) + + +@pytest.mark.parametrize("drop_params", [False, True], ids=["strict", "drop-params"]) +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_a_tool_result_with_two_marked_parts_is_one_cache_point_within_the_cap( + sync_mode: bool, drop_params: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + messages: Final = [ + {"role": "system", "content": [{"type": "text", "text": "rules", "prompt_cache_breakpoint": _EXPLICIT}]}, + {"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": _EXPLICIT}]}, + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + { + "role": "tool", + "tool_call_id": "call_1", + "content": [ + {"type": "text", "text": "sunny", "prompt_cache_breakpoint": _EXPLICIT}, + {"type": "text", "text": "warm", "prompt_cache_breakpoint": _EXPLICIT}, + ], + }, + {"role": "user", "content": "done", "prompt_cache_breakpoint": _EXPLICIT}, + ] + + result: Final = await _breakpoint_converse_request( + sync_mode, + _BREAKPOINT_MODEL, + messages, + optional_params={"tools": [_LOOKUP_TOOL]}, + litellm_params={"drop_params": drop_params}, + ) + + assert result["system"] == [{"text": "rules"}, _CACHE_POINT] + assert result["messages"] == [ + {"role": "user", "content": [{"text": "hi"}, _CACHE_POINT]}, + {"role": "assistant", "content": [{"toolUse": {"toolUseId": "call_1", "name": "lookup", "input": {}}}]}, + { + "role": "user", + "content": [ + {"toolResult": {"content": [{"text": "sunny"}, {"text": "warm"}], "toolUseId": "call_1"}}, + _CACHE_POINT, + {"text": "done"}, + _CACHE_POINT, + ], + }, + ] + + +@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"]) +@pytest.mark.asyncio +async def test_prompt_cache_options_on_a_non_claude_converse_model_are_forwarded_as_before( + sync_mode: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "drop_params", False) + + result: Final = await _breakpoint_converse_request( + sync_mode, + "us.amazon.nova-pro-v1:0", + [{"role": "user", "content": "hi"}], + optional_params={"prompt_cache_options": {"mode": "implicit"}}, + ) + + assert result["additionalModelRequestFields"] == {"prompt_cache_options": {"mode": "implicit"}}