mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
fix(router): key prompt caching affinity on tools and the /v1/messages system prompt (#46411)
* fix(router): key prompt caching affinity on tools and the /v1/messages system prompt Co-authored-by: nuernber <benjamin.nuernberger@jpl.nasa.gov> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): read a 1h tool cache_control nested under function for the pin ttl --------- Co-authored-by: nuernber <benjamin.nuernberger@jpl.nasa.gov> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Yassin Kortam
GitHub
parent
fb555b00dc
commit
d56ff281de
4 changed files with 545 additions and 30 deletions
|
|
@ -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, Mapping
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Final, TypedDict, cast
|
||||
|
||||
from typing_extensions import ReadOnly
|
||||
|
|
@ -19,7 +19,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.llms.bedrock.common_utils import bedrock_cache_point_ttl_supported, bedrock_capability_model
|
||||
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_base_model_from_metadata,
|
||||
|
|
@ -94,10 +94,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,
|
||||
|
|
@ -107,19 +132,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
|
||||
),
|
||||
|
|
@ -128,7 +147,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"]
|
||||
|
|
@ -177,11 +199,23 @@ class PromptCachingDeploymentCheck(CustomLogger):
|
|||
)
|
||||
return
|
||||
|
||||
prompt_messages: Final = cast(list[AllMessageValues], messages)
|
||||
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=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,
|
||||
|
|
@ -193,7 +227,10 @@ class PromptCachingDeploymentCheck(CustomLogger):
|
|||
await cache.async_add_model_id(
|
||||
model_id=model_id,
|
||||
messages=prompt_messages,
|
||||
tools=None, # [TODO]: add tools once standard_logging_object supports it
|
||||
tools=cast( # cast-ok: AllToolParamValues widens the union for Anthropic tool shapes
|
||||
Sequence[ChatCompletionToolParam] | None,
|
||||
tools,
|
||||
),
|
||||
honors_one_hour_ttl=_honors_one_hour_cache_ttl(standard_logging_object, litellm_params),
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -85,16 +85,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, ...]:
|
||||
|
|
@ -106,7 +167,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)
|
||||
)
|
||||
|
||||
|
|
@ -118,10 +179,16 @@ def _lookback_keys(positions: tuple[PrefixPosition, ...]) -> tuple[str, ...]:
|
|||
return tuple(entry.cache_key for entry in reversed(positions) if entry.position > oldest_probed_position)
|
||||
|
||||
|
||||
def _carrier_cache_control(carrier: Mapping[str, JsonValue]) -> JsonValue | None:
|
||||
function: Final = carrier.get("function")
|
||||
nested: Final = function.get("cache_control") if isinstance(function, dict) else None
|
||||
return carrier.get("cache_control", nested)
|
||||
|
||||
|
||||
def _breakpoint_pin_ttl_seconds(carrier: JsonValue | Mapping[str, JsonValue]) -> int | None:
|
||||
if not isinstance(carrier, Mapping):
|
||||
return None
|
||||
cache_control: Final = carrier.get("cache_control")
|
||||
cache_control: Final = _carrier_cache_control(carrier)
|
||||
if not isinstance(cache_control, dict) or cache_control.get("type") != "ephemeral":
|
||||
return None
|
||||
return PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS if cache_control.get("ttl") == "1h" else PROMPT_CACHE_PIN_TTL_SECONDS
|
||||
|
|
@ -238,6 +305,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,
|
||||
|
|
@ -256,17 +339,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
|
||||
|
|
@ -274,7 +361,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ from litellm.caching.in_memory_cache import InMemoryCache
|
|||
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,
|
||||
_honors_one_hour_cache_ttl,
|
||||
|
|
@ -1236,3 +1237,225 @@ 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,
|
||||
"custom_llm_provider": "anthropic",
|
||||
"model_map_information": {"model_map_key": model, "model_map_value": None},
|
||||
"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]]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import asyncio
|
||||
from typing import Final
|
||||
from typing import Final, cast
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -65,6 +65,11 @@ def _tool(cache_control: dict[str, str] | None) -> ChatCompletionToolParam:
|
|||
return {**tool, "cache_control": cache_control} if cache_control else tool
|
||||
|
||||
|
||||
def _function_marked_tool(cache_control: dict[str, str]) -> ChatCompletionToolParam:
|
||||
tool: Final = _tool(None)
|
||||
return {**tool, "function": {**tool["function"], "cache_control": cache_control}}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("messages", "tools", "expected_ttl"),
|
||||
[
|
||||
|
|
@ -108,6 +113,12 @@ def _tool(cache_control: dict[str, str] | None) -> ChatCompletionToolParam:
|
|||
PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS,
|
||||
id="1h-on-tools-only",
|
||||
),
|
||||
pytest.param(
|
||||
_system_block_messages({"type": "ephemeral"}),
|
||||
[_function_marked_tool({"type": "ephemeral", "ttl": "1h"})],
|
||||
PROMPT_CACHE_ONE_HOUR_PIN_TTL_SECONDS,
|
||||
id="1h-nested-under-function",
|
||||
),
|
||||
pytest.param(
|
||||
_system_block_messages({"type": "ephemeral"}),
|
||||
[_tool(None)],
|
||||
|
|
@ -450,3 +461,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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue