diff --git a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py index 9f1558dedae..640666700c8 100644 --- a/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py +++ b/litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py @@ -4,7 +4,7 @@ Check if prompt caching is valid for a given deployment Route to previously cached model id, if valid """ -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from typing import Final, cast from litellm import verbose_logger @@ -16,7 +16,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( ) from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.token_counter import offload_token_count -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam from litellm.types.utils import CallTypes, StandardLoggingPayload from litellm.utils import get_prompt_cache_min_tokens, is_prompt_caching_valid_prompt @@ -69,10 +69,35 @@ class PromptCachingDeploymentCheck(CustomLogger): if not healthy_deployments[1:]: return healthy_deployments - if messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( - messages=messages, + request_data: Final = ( + cast( # cast-ok: request_kwargs is a bare dict; this only reads its keys + Mapping[str, object], request_kwargs + ) + if request_kwargs is not None + else None + ) + request_tools: Final = ( + cast( # cast-ok: request_kwargs is untyped; the stand-down scan duck-types every tool it reads + list[AllToolParamValues] | None, request_data.get("tools") + ) + if request_data is not None + else None + ) + prompt_messages: Final = ( + PromptCachingCache.prepend_system_prompt( + messages, request_data.get("system") if request_data is not None else None + ) + if messages is not None + else None + ) + if prompt_messages is not None and await offload_token_count(is_prompt_caching_valid_prompt)( + messages=prompt_messages, model=model, min_token_count=_get_min_token_count_for_deployments(healthy_deployments), + tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes the counter duck-types + list[ChatCompletionToolParam] | None, + request_tools, + ), ): prompt_cache: Final = PromptCachingCache( cache=self.cache, @@ -82,19 +107,13 @@ class PromptCachingDeploymentCheck(CustomLogger): ## `litellm.acompletion`, after a deployment has been picked, so the affinity key has to ## be derived from the messages as they will be sent, not as they arrive here. affinity_messages: Final = AnthropicCacheControlHook.messages_with_default_injections( - messages=cast(list[AllMessageValues], messages), + messages=prompt_messages, models=( deployment["litellm_params"]["model"] for deployment in healthy_deployments if isinstance(deployment.get("litellm_params"), dict) and deployment["litellm_params"].get("model") ), - tools=( - cast( # cast-ok: request_kwargs is untyped; the stand-down scan duck-types every tool it reads - list[AllToolParamValues] | None, request_kwargs.get("tools") - ) - if request_kwargs is not None - else None - ), + tools=request_tools, enable_prompt_caching=( request_kwargs.get("enable_prompt_caching") is True if request_kwargs is not None else None ), @@ -103,7 +122,10 @@ class PromptCachingDeploymentCheck(CustomLogger): model_id_dict: Final = await prompt_cache.async_get_model_id( messages=affinity_messages, - tools=None, + tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes + Sequence[ChatCompletionToolParam] | None, + request_tools, + ), ) if model_id_dict is not None: model_id: Final = model_id_dict["model_id"] @@ -146,18 +168,34 @@ class PromptCachingDeploymentCheck(CustomLogger): ) return + event_data: Final = cast( # cast-ok: success-event kwargs are an untyped dict + Mapping[str, object], kwargs + ) + tools: Final = cast(list[AllToolParamValues] | None, event_data.get("tools")) + prompt_messages: Final = PromptCachingCache.prepend_system_prompt( + cast(list[AllMessageValues], messages), # cast-ok: isinstance(list) checked above + event_data.get("system"), + ) + ## PROMPT CACHING - cache model id, if prompt caching valid prompt + provider if await offload_token_count(is_prompt_caching_valid_prompt)( model=model, - messages=cast(list[AllMessageValues], messages), + messages=prompt_messages, + tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes the counter duck-types + list[ChatCompletionToolParam] | None, + tools, + ), ): cache: Final = PromptCachingCache( cache=self.cache, ) await cache.async_add_model_id( model_id=model_id, - messages=messages, - tools=None, # [TODO]: add tools once standard_logging_object supports it + messages=prompt_messages, + tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes + Sequence[ChatCompletionToolParam] | None, + tools, + ), ) return diff --git a/litellm/router_utils/prompt_caching_cache.py b/litellm/router_utils/prompt_caching_cache.py index e6a8dc88086..cb01cbb9437 100644 --- a/litellm/router_utils/prompt_caching_cache.py +++ b/litellm/router_utils/prompt_caching_cache.py @@ -84,16 +84,77 @@ def _chain_digest(digest: bytes, unit: bytes) -> bytes: return hashlib.sha256(digest + unit).digest() +def _is_ephemeral_mark(mark: object) -> bool: + if not isinstance(mark, dict): + return False + mark_map: Final = cast( # cast-ok: cache_control marks are plain JSON dicts + Mapping[str, object], mark + ) + return mark_map.get("type") == "ephemeral" + + +def _tool_carries_ephemeral_mark(tool: ChatCompletionToolParam) -> bool: + tool_map: Final = cast( # cast-ok: TypedDict.get rejects keys the type does not declare + Mapping[str, object], tool + ) + if _is_ephemeral_mark(tool_map.get("cache_control")): + return True + function: Final = tool_map.get("function") + if not isinstance(function, dict): + return False + return _is_ephemeral_mark( + cast(Mapping[str, object], function).get("cache_control") # cast-ok: function is a plain JSON dict + ) + + +def _strip_tool_cache_control(tool: JsonValue) -> JsonValue: + if not isinstance(tool, dict): + return tool + function: Final = tool.get("function") + stripped_function: Final = ( + {key: value for key, value in function.items() if key != "cache_control"} + if isinstance(function, dict) and "cache_control" in function + else function + ) + return { + key: (stripped_function if key == "function" else value) + for key, value in tool.items() + if key != "cache_control" + } + + +def _pin_cache_key(digest: bytes) -> str: + return f"deployment:{digest.hex()}:prompt_caching" + + def _seed(tools: Sequence[ChatCompletionToolParam] | None) -> bytes: - if tools is None: + if not tools: return hashlib.sha256(b"").digest() return hashlib.sha256( _canonical_bytes( - _TOOLS_ADAPTER.validate_python(to_jsonable_python(tools, serialize_unknown=True, bytes_mode="base64")) + tuple( + _strip_tool_cache_control(tool) + for tool in _TOOLS_ADAPTER.validate_python( + to_jsonable_python(tools, serialize_unknown=True, bytes_mode="base64") + ) + ) ) ).digest() +def _tool_boundary_position(tools: Sequence[ChatCompletionToolParam] | None) -> tuple[PrefixPosition, ...]: + if not tools: + return () + last_marked: Final = max( + (index for index, tool in enumerate(tools) if _tool_carries_ephemeral_mark(tool)), + default=None, + ) + if last_marked is None: + return () + digest: Final = _seed(tools[: last_marked + 1]) + return (PrefixPosition(cache_key=_pin_cache_key(digest), position=0),) + + def _positions_of( prefix: tuple[Mapping[str, JsonValue], ...], tools: Sequence[ChatCompletionToolParam] | None ) -> tuple[PrefixPosition, ...]: @@ -105,7 +166,7 @@ def _positions_of( for run_type, previous in zip(run_types, (None, *run_types[:-1])) ) return tuple( - PrefixPosition(cache_key=f"deployment:{digest.hex()}:prompt_caching", position=position) + PrefixPosition(cache_key=_pin_cache_key(digest), position=position) for digest, position in zip(digests, positions) ) @@ -223,6 +284,22 @@ class PromptCachingCache: return cacheable_prefix + @staticmethod + def prepend_system_prompt( + messages: list[AllMessageValues], # mutable-ok: callers pass the request messages list through + system: object, + ) -> list[AllMessageValues]: # mutable-ok: callers bind the result back into list-typed slots + if not isinstance(system, (str, list)) or not system: + return messages + first: Final = cast( # cast-ok: every AllMessageValues member is a dict + Mapping[str, object], messages[0] if messages else {} + ) + if first.get("role") == "system" and first.get("content") == system: + return messages + return cast( # cast-ok: the prepended system message is a plain dict the AllMessageValues union covers + list[AllMessageValues], [{"role": "system", "content": system}, *messages] + ) + @staticmethod def prefix_positions( messages: list[AllMessageValues] | None, @@ -241,17 +318,21 @@ class PromptCachingCache: already replaced by their size placeholder, so a request carrying the raw image bytes derives the same keys the write side stored. """ + boundary: Final = _tool_boundary_position(tools) if not messages: - return () - return _positions_of( - _PREFIX_ADAPTER.validate_python( - to_jsonable_python( - truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)), - serialize_unknown=True, - bytes_mode="base64", - ) + return boundary + return ( + *boundary, + *_positions_of( + _PREFIX_ADAPTER.validate_python( + to_jsonable_python( + truncate_base64_in_messages(PromptCachingCache.extract_cacheable_prefix(messages)), + serialize_unknown=True, + bytes_mode="base64", + ) + ), + tools, ), - tools, ) @staticmethod @@ -259,7 +340,7 @@ class PromptCachingCache: messages: list[AllMessageValues] | None, tools: Sequence[ChatCompletionToolParam] | None, ) -> tuple[PrefixPosition, ...]: - if not messages: + if not messages and not _tool_boundary_position(tools): return () return await offload_token_count(PromptCachingCache.prefix_positions)(messages, tools) diff --git a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py index c6ab259991e..52d36e01018 100644 --- a/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_prompt_caching_deployment_check.py @@ -12,6 +12,7 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import DEFAULT_MINIMUM_PROMPT_CACHE_TOKEN_COUNT, PROMPT_CACHE_LOOKBACK_POSITIONS from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( PromptCachingDeploymentCheck, _get_min_token_count_for_deployments, @@ -1075,3 +1076,223 @@ async def test_router_with_prompt_caching(anthropic_messages): print("response=", response) assert response._hidden_params["model_id"] == initial_model_id + + +_OPENAI_TOOL: Final = { + "type": "function", + "function": { + "name": "get_weather", + "description": "word " * 3000, + "parameters": {"type": "object", "properties": {}}, + }, + "cache_control": {"type": "ephemeral"}, +} +_ANTHROPIC_TOOL: Final = { + "name": "get_weather", + "description": "word " * 3000, + "input_schema": {"type": "object", "properties": {}}, + "cache_control": {"type": "ephemeral"}, +} + + +def _success_kwargs( + *, + call_type: str, + model: str, + model_id: str, + request_messages: list, + request_kwargs: dict, +) -> dict: + """The success event shape: `messages` is what standard logging produced, `system`/`tools` stay raw.""" + return { + "standard_logging_object": { + "call_type": call_type, + "model": model, + "model_id": model_id, + "messages": StandardLoggingPayloadSetup.append_system_prompt_messages( + kwargs=request_kwargs, messages=copy.deepcopy(request_messages) + ), + }, + "system": request_kwargs.get("system"), + "tools": request_kwargs.get("tools"), + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "call_type,request_messages,request_kwargs", + [ + pytest.param( + "acompletion", + [{"role": "user", "content": "weather?"}], + {"tools": [_OPENAI_TOOL]}, + id="chat_tool_breakpoint_only", + ), + pytest.param( + "anthropic_messages", + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "word " * 3000, "cache_control": {"type": "ephemeral"}}, + ], + } + ], + {"system": "you are terse"}, + id="anthropic_messages_string_system_and_marked_user_block", + ), + pytest.param( + "anthropic_messages", + [{"role": "user", "content": "hi"}], + { + "system": [ + { + "type": "text", + "text": "word " * 3000, + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ] + }, + id="anthropic_messages_list_system_breakpoint_only", + ), + pytest.param( + "anthropic_messages", + [{"role": "user", "content": "weather?"}], + {"tools": [_ANTHROPIC_TOOL]}, + id="anthropic_messages_tool_breakpoint_only", + ), + pytest.param( + "acompletion", + [{"role": "user", "content": "weather?"}], + { + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "description": "word " * 3000, + "parameters": {"type": "object", "properties": {}}, + "cache_control": {"type": "ephemeral"}, + }, + } + ] + }, + id="chat_nested_function_cache_control_only", + ), + pytest.param( + "anthropic_messages", + [ + { + "role": "user", + "content": [ + {"type": "text", "text": "word " * 3000, "cache_control": {"type": "ephemeral"}}, + ], + } + ], + {}, + id="control_message_breakpoint_no_system", + ), + ], +) +async def test_affinity_key_covers_the_tools_and_system_prompt_the_provider_sees( + call_type, request_messages, request_kwargs +): + """ + The write side keys the pin on what standard logging logged plus the raw `system`/`tools` + kwargs, and the read side must rebuild the same prefix from the raw request. Anything the + provider caches that the key ignores (tools, the /v1/messages system prompt) either writes no + pin or writes one the lookup can never compute. + """ + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + + await check.async_log_success_event( + kwargs=_success_kwargs( + call_type=call_type, + model=AUTO_CACHING_MODEL, + model_id="dep-1", + request_messages=request_messages, + request_kwargs=request_kwargs, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=cast(list[AllMessageValues], request_messages), + request_kwargs=request_kwargs, + ) + + assert filtered == [deployments[0]] + + +@pytest.mark.asyncio +async def test_a_different_tool_set_does_not_share_the_pin(): + """Tools hash into the prefix, so a same-messages request with other tools must not read the pin.""" + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + request_messages = [{"role": "user", "content": "weather?"}] + other_tool = {**_OPENAI_TOOL, "function": {**_OPENAI_TOOL["function"], "name": "get_time"}} + + await check.async_log_success_event( + kwargs=_success_kwargs( + call_type="acompletion", + model=AUTO_CACHING_MODEL, + model_id="dep-1", + request_messages=request_messages, + request_kwargs={"tools": [_OPENAI_TOOL]}, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=cast(list[AllMessageValues], request_messages), + request_kwargs={"tools": [other_tool]}, + ) + + assert filtered == deployments + + +@pytest.mark.asyncio +async def test_a_pin_written_without_tools_is_read_when_tools_is_an_empty_list(): + """`tools=None` on the write and `tools=[]` on the read are the same empty seed.""" + cache = DualCache() + check = PromptCachingDeploymentCheck(cache=cache) + deployments = _deployments(AUTO_CACHING_MODEL, AUTO_CACHING_MODEL) + request_messages = [ + { + "role": "user", + "content": [{"type": "text", "text": "word " * 3000, "cache_control": {"type": "ephemeral"}}], + } + ] + + await check.async_log_success_event( + kwargs=_success_kwargs( + call_type="anthropic_messages", + model=AUTO_CACHING_MODEL, + model_id="dep-1", + request_messages=request_messages, + request_kwargs={}, + ), + response_obj=None, + start_time=None, + end_time=None, + ) + + filtered = await check.async_filter_deployments( + model=MODEL_GROUP_ALIAS, + healthy_deployments=deployments, + messages=cast(list[AllMessageValues], request_messages), + request_kwargs={"tools": []}, + ) + + assert filtered == [deployments[0]] diff --git a/tests/unit/router_utils/test_prompt_caching_cache.py b/tests/unit/router_utils/test_prompt_caching_cache.py index 3a1132354c4..171e10bac3f 100644 --- a/tests/unit/router_utils/test_prompt_caching_cache.py +++ b/tests/unit/router_utils/test_prompt_caching_cache.py @@ -1,9 +1,10 @@ import asyncio -from typing import Final +from typing import Final, cast import pytest from litellm.router_utils.prompt_caching_cache import PromptCachingCache +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam async def _first_non_none(read): @@ -184,3 +185,160 @@ async def test_same_cacheable_prefix_routes_to_same_deployment() -> None: model_id_3: Final = response3._hidden_params["model_id"] assert model_id_1 == model_id_2 == model_id_3 + + +def _marked_tool( + name: str, cache_control: dict | None = None, *, nested: bool = False +) -> ChatCompletionToolParam: + tool: Final = { + "type": "function", + "function": {"name": name, "description": "word " * 40, "parameters": {"type": "object", "properties": {}}}, + } + if cache_control is not None and nested: + return cast( + ChatCompletionToolParam, + {**tool, "function": {**tool["function"], "cache_control": cache_control}}, + ) + if cache_control is not None: + return cast(ChatCompletionToolParam, {**tool, "cache_control": cache_control}) + return cast(ChatCompletionToolParam, tool) + + +def test_prepend_system_prompt_dedupes_the_string_standard_logging_prepends(): + """ + The success event logs the string system prompt already prepended as a system message, and the + request still carries it in `system`. Prepending unconditionally would hash the prefix twice on + one side, so the copy already sitting at messages[0] must be recognized as the same prompt. + """ + messages = cast( + list[AllMessageValues], + [{"role": "system", "content": "be brief"}, {"role": "user", "content": "hi"}], + ) + + assert PromptCachingCache.prepend_system_prompt(messages, "be brief") is messages + + +def test_prepend_system_prompt_prepends_a_list_system_prompt(): + """ + Standard logging only prepends a string system, so a /v1/messages list system never reaches the + write side's messages. Both sides have to rebuild the same list the provider caches. + """ + system: Final = [{"type": "text", "text": "be brief", "cache_control": {"type": "ephemeral", "ttl": "1h"}}] + messages = cast(list[AllMessageValues], [{"role": "user", "content": "hi"}]) + + result = PromptCachingCache.prepend_system_prompt(messages, system) + + assert result == [{"role": "system", "content": system}, *messages] + + +@pytest.mark.parametrize("system", [None, "", [], 0, {"type": "text"}]) +def test_prepend_system_prompt_leaves_unsystem_values_untouched(system): + messages = cast(list[AllMessageValues], [{"role": "user", "content": "hi"}]) + + assert PromptCachingCache.prepend_system_prompt(messages, system) is messages + + +def test_prepend_system_prompt_never_mutates_the_input_messages(): + system: Final = [{"type": "text", "text": "be brief"}] + user_message: Final = {"role": "user", "content": "hi"} + messages = cast(list[AllMessageValues], [user_message]) + + result = PromptCachingCache.prepend_system_prompt(messages, system) + + assert result is not messages + assert messages == [user_message] + + +def test_a_tool_breakpoint_emits_a_position_zero_key_even_without_messages(): + """ + A request that marks only its tools still hits the provider prompt cache on the tool prefix, so + the pin has to exist at a position before the first message block. + """ + tools: Final = [_marked_tool("read", {"type": "ephemeral"})] + + positions = PromptCachingCache.prefix_positions([], tools) + assert len(positions) == 1 + assert positions[0].position == 0 + + with_messages = PromptCachingCache.prefix_positions( + cast( + list[AllMessageValues], + [{"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]}], + ), + tools, + ) + assert [position.position for position in with_messages] == [0, 1] + + +def test_a_retimed_or_moved_tool_breakpoint_keeps_the_same_key(): + """ + cache_control on a tool is a boundary marker, not part of the cached content, so a ttl change + or a breakpoint that leaves an earlier tool unmarked hashes to the same prefix key. + """ + marked: Final = {"type": "ephemeral"} + retimed: Final = {"type": "ephemeral", "ttl": "1h"} + tools_5m: Final = [_marked_tool("read", marked)] + tools_1h: Final = [_marked_tool("read", retimed)] + + key_5m: Final = PromptCachingCache.get_prompt_caching_cache_key([], tools_5m) + key_1h: Final = PromptCachingCache.get_prompt_caching_cache_key([], tools_1h) + + assert key_5m is not None + assert key_5m == key_1h + + both_marked: Final = [_marked_tool("read", marked), _marked_tool("write", marked)] + last_marked: Final = [_marked_tool("read"), _marked_tool("write", retimed)] + + key_both: Final = PromptCachingCache.get_prompt_caching_cache_key([], both_marked) + key_last: Final = PromptCachingCache.get_prompt_caching_cache_key([], last_marked) + assert key_both == key_last + + +def test_a_nested_function_cache_control_marks_a_boundary_and_strips_from_the_key(): + """ + LiteLLM also honors the OpenAI-style mark nested inside `function` (the Anthropic transform + maps it to the tool's cache_control), so it has to mark a boundary and hash away exactly like + the top-level form, or a client that marks that way loses affinity to a breakpoint that moved. + """ + marked: Final = {"type": "ephemeral"} + retimed: Final = {"type": "ephemeral", "ttl": "1h"} + + positions = PromptCachingCache.prefix_positions([], [_marked_tool("read", marked, nested=True)]) + assert len(positions) == 1 + assert positions[0].position == 0 + + key_nested: Final = PromptCachingCache.get_prompt_caching_cache_key( + [], [_marked_tool("read", marked, nested=True)] + ) + key_nested_retimed: Final = PromptCachingCache.get_prompt_caching_cache_key( + [], [_marked_tool("read", retimed, nested=True)] + ) + key_toplevel: Final = PromptCachingCache.get_prompt_caching_cache_key([], [_marked_tool("read", marked)]) + + assert key_nested is not None + assert key_nested == key_nested_retimed + assert key_nested == key_toplevel + + +def test_unmarked_tools_change_the_key_marked_tools_do_not_share(): + """Tools are provider-cached content, so a different tool set must never read another request's pin.""" + tools_a: Final = [_marked_tool("read", {"type": "ephemeral"})] + tools_b: Final = [_marked_tool("write", {"type": "ephemeral"})] + + key_a: Final = PromptCachingCache.get_prompt_caching_cache_key([], tools_a) + key_b: Final = PromptCachingCache.get_prompt_caching_cache_key([], tools_b) + + assert key_a is not None + assert key_a != key_b + + +def test_empty_and_missing_tools_seed_the_same_key(): + """`tools=[]` and no tools at all are the same request to the provider, so they must key alike.""" + messages = cast( + list[AllMessageValues], + [{"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]}], + ) + + key_none: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, None) + key_empty: Final = PromptCachingCache.get_prompt_caching_cache_key(messages, []) + assert key_none == key_empty