diff --git a/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py b/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py index fd8f3c54ca0..9765b078eaf 100644 --- a/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py +++ b/litellm/llms/anthropic/chat/prompt_cache_breakpoints.py @@ -202,6 +202,10 @@ def _block_mark(block: object, key: str) -> object: return (_fields(block) or {}).get(key) +def _is_supported(marker: object) -> bool: + return not isinstance(_breakpoint_mark(marker), UnsupportedPromptCacheParam) + + def _asks_for_an_hour(marker: object) -> bool: mark: Final = _breakpoint_mark(marker) return not isinstance(mark, UnsupportedPromptCacheParam) and mark.get("ttl") == "1h" @@ -243,9 +247,12 @@ def tool_mark_may_reach_claude(tool: object) -> bool: return any(tool_mark_reaches_claude(tool, route) for route in _CLAUDE_ROUTES) -def _first_usable(markers: Sequence[object]) -> tuple[object, tuple[object, ...]]: - """The first breakpoint Claude can use stays, so a malformed one ``drop_params`` drops never takes its place.""" - index: Final = next((index for index, marker in enumerate(markers) if prompt_cache_breakpoint_is_usable(marker)), 0) +def _kept_breakpoint(markers: Sequence[object]) -> tuple[object, tuple[object, ...]]: + """The first breakpoint Claude takes as sent stays, else the first one ``drop_params`` can fix, so a malformed one + or one whose ttl ``drop_params`` drops never takes the place of a valid one and costs the ttl it chose.""" + supported: Final = (index for index, marker in enumerate(markers) if _is_supported(marker)) + usable: Final = (index for index, marker in enumerate(markers) if prompt_cache_breakpoint_is_usable(marker)) + index: Final = next(chain(supported, usable), 0) return markers[index], (*markers[:index], *markers[index + 1 :]) @@ -270,7 +277,7 @@ def _tool_result_hoisted(message: _Fields) -> _Hoisted: ) if own is None and not part_breakpoints: return _Hoisted(message, ()) - kept, replaced = _first_usable(part_breakpoints if own is None else (own, *part_breakpoints)) + kept, replaced = _kept_breakpoint(part_breakpoints if own is None else (own, *part_breakpoints)) return _Hoisted( {**message, "prompt_cache_breakpoint": _joined_breakpoint(kept, replaced), "content": unmarked_parts}, replaced ) @@ -314,7 +321,7 @@ def _message_breakpoint_hoisted(message: _Fields) -> _Hoisted: if block.get("cache_control") is not None: return _Hoisted(_without_own_breakpoint(message), (own,)) existing: Final = block.get("prompt_cache_breakpoint") - kept, replaced = (own, ()) if existing is None else _first_usable((existing, own)) + kept, replaced = (own, ()) if existing is None else _kept_breakpoint((existing, own)) return _Hoisted( { **_without_own_breakpoint(message), @@ -366,7 +373,7 @@ def _hoisted_tool(tool: object) -> _Hoisted: } if function.get("cache_control") is not None or fields.get("cache_control") is not None: return _Hoisted(unmarked, markers) - kept, replaced = _first_usable(markers) + kept, replaced = _kept_breakpoint(markers) return _Hoisted({**unmarked, "prompt_cache_breakpoint": _joined_breakpoint(kept, replaced)}, replaced) diff --git a/tests/benchmarks/test_prompt_cache_breakpoint_benchmarks.py b/tests/benchmarks/test_prompt_cache_breakpoint_benchmarks.py new file mode 100644 index 00000000000..948ae0d2709 --- /dev/null +++ b/tests/benchmarks/test_prompt_cache_breakpoint_benchmarks.py @@ -0,0 +1,50 @@ +from collections.abc import Sequence +from typing import Final + +import pytest + +from litellm.llms.anthropic.chat.prompt_cache_breakpoints import ClaudePromptCache, claude_prompt_cache + +_EXPLICIT: Final = {"mode": "explicit"} + +_MARKED_SYSTEM_TURNS: Final = ( + {"role": "user", "content": "hi"}, + *({"role": "system", "content": f"s{index}", "prompt_cache_breakpoint": _EXPLICIT} for index in range(2000)), +) + +_ONE_HOUR_PARTS_BESIDE_A_BREAKPOINT: Final = ( + { + "role": "user", + "content": [ + *( + {"type": "text", "text": f"p{index}", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + for index in range(5000) + ), + {"type": "text", "text": "end", "prompt_cache_breakpoint": _EXPLICIT}, + ], + }, +) + + +def _settled_under_drop_params(messages: Sequence[object]) -> None: + outcome: Final = claude_prompt_cache( + messages, + None, + {}, + drop_params=True, + tool_history_as_text=False, + tool_config_cache_control=None, + route="anthropic", + pairs_unanswered_tool_calls=False, + ) + assert isinstance(outcome, ClaudePromptCache) + + +@pytest.mark.benchmark +def test_settling_thousands_of_marked_system_turns_under_drop_params(): + _settled_under_drop_params(_MARKED_SYSTEM_TURNS) + + +@pytest.mark.benchmark +def test_settling_a_breakpoint_beside_thousands_of_one_hour_parts_under_drop_params(): + _settled_under_drop_params(_ONE_HOUR_PARTS_BESIDE_A_BREAKPOINT) 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 c1076f87c64..24d8c65f1de 100644 --- a/tests/unit/llms/anthropic/chat/test_prompt_cache_breakpoints.py +++ b/tests/unit/llms/anthropic/chat/test_prompt_cache_breakpoints.py @@ -1,4 +1,3 @@ -import time from typing import Final import pytest @@ -933,6 +932,68 @@ def test_a_tools_breakpoint_joined_with_its_functions_keeps_the_hour_the_functio assert list(result.tools or ()) == [{"type": "function", "function": {"name": "t"}, "cache_control": _ONE_HOUR}] +@pytest.mark.parametrize( + ("messages", "tools", "sent_messages", "sent_tools"), + [ + pytest.param( + [ + { + "role": "tool", + "tool_call_id": "c1", + "content": [ + _part("r0", prompt_cache_breakpoint=_THIRTY_MINUTES), + _part("r1", prompt_cache_breakpoint={"mode": "explicit", "ttl": "1h"}), + ], + } + ], + None, + [{"role": "tool", "tool_call_id": "c1", "content": [_part("r0"), _part("r1")], "cache_control": _ONE_HOUR}], + [], + id="tool-result-parts", + ), + pytest.param( + [ + { + "role": "user", + "content": [_part("u0"), _part("u1", prompt_cache_breakpoint=_THIRTY_MINUTES)], + "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}, + } + ], + None, + [{"role": "user", "content": [_part("u0"), _part("u1", cache_control=_ONE_HOUR)]}], + [], + id="message-and-its-last-part", + ), + pytest.param( + [{"role": "user", "content": "hi"}], + [ + { + "type": "function", + "prompt_cache_breakpoint": _THIRTY_MINUTES, + "function": {"name": "t", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}}, + } + ], + [{"role": "user", "content": "hi"}], + [{"type": "function", "function": {"name": "t"}, "cache_control": _ONE_HOUR}], + id="tool-and-its-function", + ), + ], +) +def test_a_breakpoint_whose_ttl_drop_params_drops_never_costs_the_hour_of_the_one_it_joins( + messages: list[object], + tools: list[object] | None, + sent_messages: list[object], + sent_tools: list[object], +): + outcome: Final = _cache(messages, tools) + result: Final = _resolved(messages, tools, drop_params=True) + + assert isinstance(outcome, UnsupportedPromptCacheParam) + assert outcome.value == _THIRTY_MINUTES + assert list(result.messages) == sent_messages + assert list(result.tools or ()) == sent_tools + + def test_a_function_breakpoint_the_tools_own_breakpoint_replaces_is_still_reported(): tools: Final = [ { @@ -948,41 +1009,6 @@ def test_a_function_breakpoint_the_tools_own_breakpoint_replaces_is_still_report assert outcome.value == _THIRTY_MINUTES -@pytest.mark.parametrize( - "messages", - [ - pytest.param( - [ - {"role": "user", "content": "hi"}, - *( - {"role": "system", "content": f"s{index}", "prompt_cache_breakpoint": _EXPLICIT} - for index in range(8000) - ), - ], - id="system-turns-with-breakpoints", - ), - pytest.param( - [ - { - "role": "user", - "content": [ - *(_part(f"p{index}", cache_control=_ONE_HOUR) for index in range(20000)), - _part("end", prompt_cache_breakpoint=_EXPLICIT), - ], - } - ], - id="one-hour-parts-beside-a-breakpoint", - ), - ], -) -def test_settling_thousands_of_marks_under_drop_params_takes_well_under_a_second_per_thousand(messages: list[object]): - started: Final = time.perf_counter() - - _resolved(messages, drop_params=True) - - assert time.perf_counter() - started < 5 - - def test_a_tool_config_cache_point_takes_one_of_claudes_four_and_is_never_dropped_under_drop_params(): messages: Final = [ {"role": "system", "content": [_part("s", prompt_cache_breakpoint=_EXPLICIT)]},