diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index 17c469a9d63..b18ad21c47b 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -232,17 +232,37 @@ def developer_message_with_valid_cache_control(message: AllMessageValues) -> All return with_cache_control(message, validated_cache_control) -def _cache_control_values(value: object) -> tuple[object, ...]: +def _own_cache_control(value: object) -> tuple[object, ...]: fields: Final = as_str_mapping(value) - if fields is not None: - own: Final = () if fields.get("cache_control") is None else (fields["cache_control"],) - return own + _nested_cache_control_values(field for key, field in fields.items() if key != "cache_control") - items: Final = as_object_list(value) - return () if items is None else _nested_cache_control_values(items) + return () if fields is None or fields.get("cache_control") is None else (fields["cache_control"],) -def _nested_cache_control_values(values: Iterable[object]) -> tuple[object, ...]: - return tuple(itertools.chain.from_iterable(map(_cache_control_values, values))) +def _each_cache_control( + cache_controls_of: Callable[[object], tuple[object, ...]], values: object +) -> tuple[object, ...]: + return tuple(itertools.chain.from_iterable(map(cache_controls_of, as_object_list(values) or ()))) + + +def _message_cache_controls(message: object) -> tuple[object, ...]: + fields: Final = as_str_mapping(message) + if fields is None: + return () + return _own_cache_control(fields) + tuple( + itertools.chain.from_iterable(map(_own_cache_control, _mark_carrying_blocks(fields))) + ) + + +def _tool_cache_controls(tool: object) -> tuple[object, ...]: + fields: Final = as_str_mapping(tool) + return () if fields is None else _own_cache_control(fields) + _own_cache_control(fields.get("function")) + + +def _request_cache_controls(fields: Mapping[str, object]) -> tuple[object, ...]: + extra_body: Final = as_str_mapping(fields.get("extra_body")) or MappingProxyType({}) + top_level: Final = extra_body.get("cache_control", fields.get("cache_control")) + return (() if top_level is None else (top_level,)) + _each_cache_control( + _tool_cache_controls, extra_body["tools"] if "tools" in extra_body else fields.get("tools") + ) def _cache_control_ttl(cache_control: object) -> object: @@ -250,8 +270,12 @@ def _cache_control_ttl(cache_control: object) -> object: return None if fields is None else fields.get("ttl") or "5m" -def cache_control_ttls(value: object) -> tuple[object, ...]: - return tuple(map(_cache_control_ttl, _cache_control_values(value))) +def messages_cache_control_ttls(messages: object) -> tuple[object, ...]: + return tuple(map(_cache_control_ttl, _each_cache_control(_message_cache_controls, messages))) + + +def request_cache_control_ttls(fields: Mapping[str, object]) -> tuple[object, ...]: + return tuple(map(_cache_control_ttl, _request_cache_controls(fields))) def tool_config_point_ttls(points: object) -> tuple[object, ...]: @@ -271,8 +295,8 @@ def request_keeps_developer_cache_control( shares one TTL within the cap.""" return keeps_one_ttl_within_the_cap( ( - *cache_control_ttls(messages), - *cache_control_ttls(other_request_fields), + *messages_cache_control_ttls(messages), + *request_cache_control_ttls(other_request_fields), *tool_config_point_ttls(other_request_fields.get("cache_control_injection_points")), ) ) @@ -1321,7 +1345,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): continue if targeted_role == "developer" and not keeps_one_ttl_within_the_cap( - (*cache_control_ttls(messages), *outside_ttls, _cache_control_ttl(control)) + (*messages_cache_control_ttls(messages), *outside_ttls, _cache_control_ttl(control)) ): continue @@ -1906,7 +1930,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): model_info, ) outside_ttls: Final = ( - cache_control_ttls( + request_cache_control_ttls( MappingProxyType( { "tools": cast(list[object] | None, tools), # cast-ok: the seed types tools as a bare list diff --git a/tests/unit/integrations/test_anthropic_cache_control_hook.py b/tests/unit/integrations/test_anthropic_cache_control_hook.py index 183e2340e52..86245e4a079 100644 --- a/tests/unit/integrations/test_anthropic_cache_control_hook.py +++ b/tests/unit/integrations/test_anthropic_cache_control_hook.py @@ -24,10 +24,10 @@ from litellm._logging import verbose_logger from litellm.integrations.anthropic_cache_control_hook import ( CARRY_UNMATCHED_MESSAGE_POINTS, EXTERNAL_BREAKPOINTS_STAMP, - PROMPT_CACHING_STAMP, AnthropicCacheControlHook, configured_injection_points, developer_messages_for_a_prompt_caching_model, + developer_messages_within_the_cache_limits, supports_openai_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -5272,28 +5272,35 @@ def test_system_injection_point_on_a_developer_prompt_starts_no_provider_sign_in pytest.param("lemonade", "Qwen3-8B-GGUF", id="lemonade"), ], ) -def test_developer_prompt_caching_check_sends_no_model_lookup(custom_llm_provider: str, model: str): - messages: Final = cast( - list[AllMessageValues], - [ - {"role": "developer", "content": "Answer in French.", "cache_control": {"type": "ephemeral"}}, - {"role": "user", "content": "Hi"}, - ], - ) +def test_developer_prompt_caching_check_answers_without_a_model_lookup(custom_llm_provider: str, model: str): + developer: Final = {"role": "developer", "content": "Answer in French."} + question: Final = {"role": "user", "content": "Hi"} params: Final[dict[str, object]] = {"cache_control_injection_points": [{"location": "message", "role": "system"}]} with respx.mock(assert_all_called=False) as router: lookups: Final = router.route().mock(return_value=Response(500)) - developer_messages_for_a_prompt_caching_model( - messages, model=model, custom_llm_provider=custom_llm_provider, model_info=None, base_model=None + checked: Final = developer_messages_for_a_prompt_caching_model( + cast(list[AllMessageValues], [{**developer, "cache_control": {"type": "ephemeral"}}, question]), + model=model, + custom_llm_provider=custom_llm_provider, + model_info=None, + base_model=None, ) AnthropicCacheControlHook.maybe_seed_default_injection_points( - params, list(messages), model, custom_llm_provider + params, cast(list[AllMessageValues], [dict(developer), dict(question)]), model, custom_llm_provider + ) + _, injected, _ = AnthropicCacheControlHook().get_chat_completion_prompt( + model=model, + messages=cast(list[AllMessageValues], [dict(developer), dict(question)]), + non_default_params=params, + prompt_id=None, + prompt_variables=None, + dynamic_callback_params={}, ) assert lookups.called is False - seeded: Final = cast(list[dict[str, object]], params["cache_control_injection_points"]) - assert all(PROMPT_CACHING_STAMP in point for point in seeded) + assert list(checked) == [developer, question] + assert injected == [developer, question] _SHARED_FOUNDRY_MODEL: Final = "azure_ai/qa-shared-gpt-deployment" _SIBLING_BREAKPOINT_FLAGS: Final = MappingProxyType( @@ -6185,6 +6192,90 @@ def test_system_injection_point_leaves_a_developer_prompt_unmarked_when_its_mark ) +_TOOL_WITH_A_CACHE_CONTROL_ARGUMENT: Final = MappingProxyType( + { + "type": "function", + "function": { + "name": "set_cache_policy", + "parameters": {"type": "object", "properties": {"cache_control": {"type": "string"}}}, + }, + } +) + + +@pytest.mark.parametrize( + ("developer_mark", "points"), + [ + pytest.param(dict(_ONE_HOUR), None, id="client_mark"), + pytest.param(None, [{"location": "message", "role": "system", "control": dict(_ONE_HOUR)}], id="injection_point"), + ], +) +def test_one_hour_developer_prompt_stays_cached_beside_a_tool_argument_named_cache_control( + developer_mark: dict[str, str] | None, + points: list[dict[str, object]] | None, + respx_mock: respx.MockRouter, +) -> None: + route: Final = _converse_route(respx_mock) + developer: Final = {"role": "developer", "content": _DEVELOPER_PREFIX} + + litellm.completion( + model=_CLAUDE_ON_CONVERSE, + messages=[ + developer if developer_mark is None else {**developer, "cache_control": developer_mark}, + {"role": "user", "content": "Which rule covers backups?"}, + ], + tools=[dict(_TOOL_WITH_A_CACHE_CONTROL_ARGUMENT)], + cache_control_injection_points=points, + **_OFFLINE_AWS_CREDENTIALS, + ) + + assert _sent_converse_body(route).system == ( + {"text": _DEVELOPER_PREFIX}, + {"cachePoint": {"type": "default", "ttl": "1h"}}, + ) + + +_FIVE_MINUTES: Final = MappingProxyType({"type": "ephemeral", "ttl": "5m"}) +_FIVE_MINUTE_TOOL: Final = MappingProxyType({**_WEATHER_TOOL, "cache_control": dict(_FIVE_MINUTES)}) + + +@pytest.mark.parametrize( + ("outside_fields", "keeps_the_mark"), + [ + pytest.param({"tools": [dict(_TOOL_WITH_A_CACHE_CONTROL_ARGUMENT)]}, True, id="tool_argument_named_cache_control"), + pytest.param( + {"tools": [{**_WEATHER_TOOL, "function": {**_WEATHER_TOOL["function"], "cache_control": dict(_FIVE_MINUTES)}}]}, + False, + id="mark_under_the_tool_function", + ), + pytest.param({"cache_control": dict(_FIVE_MINUTES)}, False, id="top_level_mark"), + pytest.param({"extra_body": {"cache_control": dict(_FIVE_MINUTES)}}, False, id="envelope_top_level_mark"), + pytest.param( + {"tools": [dict(_WEATHER_TOOL)], "extra_body": {"tools": [dict(_FIVE_MINUTE_TOOL)]}}, + False, + id="envelope_tool_mark", + ), + pytest.param( + {"tools": [dict(_FIVE_MINUTE_TOOL)], "extra_body": {"tools": [dict(_WEATHER_TOOL)]}}, + True, + id="tool_mark_the_envelope_replaces", + ), + ], +) +def test_one_hour_developer_mark_follows_the_ttls_of_the_marks_sent_outside_the_messages( + outside_fields: dict[str, object], keeps_the_mark: bool +) -> None: + developer: Final = {"role": "developer", "content": _DEVELOPER_PREFIX} + question: Final = {"role": "user", "content": "Which rule covers backups?"} + marked_developer: Final = {**developer, "cache_control": dict(_ONE_HOUR)} + + kept: Final = developer_messages_within_the_cache_limits( + cast(list[AllMessageValues], [marked_developer, question]), MappingProxyType(outside_fields) + ) + + assert list(kept) == [marked_developer if keeps_the_mark else developer, question] + + def test_system_injection_point_marks_a_developer_prompt_for_openai_prompt_cache_breakpoints( respx_mock: respx.MockRouter, ) -> None: