mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
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.
This commit is contained in:
ryan-crabbe-berri
parent
3acf5f514a
commit
ab7883fd57
7 changed files with 270 additions and 17 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"}}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue