mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(cache): count only cache marks a provider receives in the developer TTL check
The developer TTL check read every cache_control key nested anywhere in the request, so a tool whose JSON schema has a property named cache_control dropped a 1h developer mark or made a role system point skip the developer prompt. It now reads marks on messages and their content and tool call parts, on a tool or its function, the top-level cache_control, and extra_body's cache_control and tools, which replace the request's own The lookup test now asserts the request a prompt-caching check returns instead of the stamp it sets
This commit is contained in:
mateo-berri
parent
63003a7044
commit
fa4e86f46f
2 changed files with 143 additions and 28 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue