mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
test(anthropic): cover prompt_cache_breakpoint to cache_control on every Claude route
Unit tests for the block-level rename, the wire-order cap, the drop_params sources, prompt_cache_options rejection and the untouched native cache_control path on the Anthropic, Vertex AI, Bedrock invoke, Bedrock Converse and Azure AI transforms, plus a providers wire test that drives all five routes through a live proxy against scripted peers, including a proxy started with litellm_settings drop_params
This commit is contained in:
ryan-crabbe-berri
parent
432c8b5aa5
commit
e671037110
10 changed files with 1311 additions and 1 deletions
|
|
@ -0,0 +1,368 @@
|
|||
from collections.abc import Iterator, Mapping
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, assert_never
|
||||
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, gateway_from_environment
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.vertex import service_account_json
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.providers._cache_control_marks_support import (
|
||||
ANTHROPIC_MODEL,
|
||||
ASK,
|
||||
ASK_LABEL,
|
||||
BEDROCK_MODEL,
|
||||
CITIES,
|
||||
EPHEMERAL,
|
||||
POINTS,
|
||||
PROVIDER_KEY,
|
||||
SYSTEM,
|
||||
SYSTEM_LABEL,
|
||||
TOOL,
|
||||
Mark,
|
||||
anthropic_labels,
|
||||
anthropic_marks,
|
||||
anthropic_peer,
|
||||
bedrock_labels,
|
||||
bedrock_peer,
|
||||
call_id,
|
||||
chat_body,
|
||||
final_label,
|
||||
final_text,
|
||||
gateway_injected,
|
||||
new_marker,
|
||||
owned_config,
|
||||
post_chat,
|
||||
tool_call,
|
||||
tool_use_label,
|
||||
)
|
||||
from pydantic import BaseModel, ConfigDict, JsonValue
|
||||
|
||||
Provider = Literal["anthropic", "bedrock_converse", "bedrock_invoke", "vertex_ai", "azure_ai"]
|
||||
Site = Literal["tool", "system", "user_part", "tool_call", "tool_result_part", "string_message"]
|
||||
|
||||
_PROVIDERS: Final[tuple[Provider, ...]] = ("anthropic", "bedrock_converse", "bedrock_invoke", "vertex_ai", "azure_ai")
|
||||
_SITES: Final[tuple[Site, ...]] = ("tool", "system", "user_part", "tool_call", "tool_result_part", "string_message")
|
||||
_FIVE_SITES: Final[tuple[Site, ...]] = ("tool", "system", "user_part", "tool_call", "string_message")
|
||||
_EXPLICIT: Final[dict[str, JsonValue]] = {"mode": "explicit"}
|
||||
_EXPLICIT_30M: Final[dict[str, JsonValue]] = {"mode": "explicit", "ttl": "30m"}
|
||||
_BREAKPOINT: Final[Mapping[str, JsonValue]] = MappingProxyType({"prompt_cache_breakpoint": _EXPLICIT})
|
||||
_NATIVE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache_control": EPHEMERAL})
|
||||
_CITY: Final = CITIES[0]
|
||||
_TOOL_RESULT_TEXT: Final = "sunny."
|
||||
_BEDROCK_TOKEN: Final = "synthetic-bedrock-bearer"
|
||||
_VERTEX_PROJECT: Final = "scripted-project"
|
||||
_VERTEX_LOCATION: Final = "us-east5"
|
||||
_TARGETS: Final[Mapping[Provider, str]] = MappingProxyType(
|
||||
{
|
||||
"anthropic": "/v1/messages",
|
||||
"bedrock_converse": f"/model/{BEDROCK_MODEL}/converse",
|
||||
"bedrock_invoke": f"/model/{BEDROCK_MODEL}/invoke",
|
||||
"vertex_ai": (
|
||||
f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}"
|
||||
f"/publishers/anthropic/models/{ANTHROPIC_MODEL}:rawPredict"
|
||||
),
|
||||
"azure_ai": "/anthropic/v1/messages",
|
||||
}
|
||||
)
|
||||
_CAP_HINT: Final = "Claude accepts at most 4 cache breakpoints, found 5"
|
||||
|
||||
|
||||
class _Nested(BaseModel):
|
||||
model_config = ConfigDict(extra="allow", frozen=True)
|
||||
type: str
|
||||
text: str | None = None
|
||||
cache_control: JsonValue = None
|
||||
|
||||
|
||||
class _Block(BaseModel):
|
||||
model_config = ConfigDict(extra="allow", frozen=True)
|
||||
type: str
|
||||
tool_use_id: str | None = None
|
||||
content: str | tuple[_Nested, ...] | None = None
|
||||
|
||||
|
||||
class _Message(BaseModel):
|
||||
model_config = ConfigDict(extra="allow", frozen=True)
|
||||
role: str
|
||||
content: str | tuple[_Block, ...]
|
||||
|
||||
|
||||
class _Body(BaseModel):
|
||||
model_config = ConfigDict(extra="allow", frozen=True)
|
||||
messages: tuple[_Message, ...] = ()
|
||||
|
||||
|
||||
def _nested_blocks(body: _Body) -> Iterator[tuple[_Block, _Nested]]:
|
||||
for message in body.messages:
|
||||
for block in message.content if isinstance(message.content, tuple) else ():
|
||||
if isinstance(block.content, tuple):
|
||||
yield from ((block, inner) for inner in block.content)
|
||||
|
||||
|
||||
def _nested_tool_result_labels(request: Request) -> list[str]:
|
||||
return [
|
||||
f"user:tool_result:{block.tool_use_id}:{inner.type}:{inner.text}"
|
||||
for block, inner in _nested_blocks(_Body.model_validate_json(request.body))
|
||||
if inner.cache_control is not None
|
||||
]
|
||||
|
||||
|
||||
def _labels(provider: Provider, request: Request) -> list[str]:
|
||||
if provider == "bedrock_converse":
|
||||
return bedrock_labels(request)
|
||||
return [*anthropic_labels(request), *_nested_tool_result_labels(request)]
|
||||
|
||||
|
||||
def _respond(provider: Provider, request: Request) -> Reply:
|
||||
return bedrock_peer(request) if provider == "bedrock_converse" else anthropic_peer(request)
|
||||
|
||||
|
||||
def _litellm_params(gateway: Gateway, wire: Wire, provider: Provider) -> dict[str, JsonValue]:
|
||||
match provider:
|
||||
case "anthropic":
|
||||
return {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_base": wire.url, "api_key": PROVIDER_KEY}
|
||||
case "bedrock_converse":
|
||||
return {
|
||||
"model": f"bedrock/converse/{BEDROCK_MODEL}",
|
||||
"api_key": PROVIDER_KEY,
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_bedrock_runtime_endpoint": wire.url,
|
||||
}
|
||||
case "bedrock_invoke":
|
||||
return {
|
||||
"model": f"bedrock/invoke/{BEDROCK_MODEL}",
|
||||
"api_key": _BEDROCK_TOKEN,
|
||||
"aws_region_name": "us-east-1",
|
||||
"api_base": wire.url,
|
||||
"aws_bedrock_runtime_endpoint": wire.url,
|
||||
}
|
||||
case "vertex_ai":
|
||||
return {
|
||||
"model": f"vertex_ai/{ANTHROPIC_MODEL}",
|
||||
"api_base": wire.url,
|
||||
"api_key": None,
|
||||
"vertex_project": _VERTEX_PROJECT,
|
||||
"vertex_location": _VERTEX_LOCATION,
|
||||
"vertex_credentials": service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")),
|
||||
}
|
||||
case "azure_ai":
|
||||
return {"model": f"azure_ai/{ANTHROPIC_MODEL}", "api_base": wire.url, "api_key": PROVIDER_KEY}
|
||||
case _:
|
||||
assert_never(provider)
|
||||
|
||||
|
||||
def _deployment(gateway: Gateway, scenario: Scenario, wire: Wire, provider: Provider, **fields: JsonValue) -> str:
|
||||
return scenario.model(model_info=None, **_litellm_params(gateway, wire, provider), **fields)
|
||||
|
||||
|
||||
def _messages(marker: str, marks: Mapping[Site, Mapping[str, JsonValue]]) -> list[JsonValue]:
|
||||
def at(site: Site) -> Mapping[str, JsonValue]:
|
||||
return marks.get(site, {})
|
||||
|
||||
return [
|
||||
{"role": "system", "content": [{"type": "text", "text": SYSTEM, **at("system")}]},
|
||||
{"role": "user", "content": [{"type": "text", "text": ASK, **at("user_part")}]},
|
||||
{"role": "assistant", "content": "", "tool_calls": [tool_call(_CITY, **at("tool_call"))]},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": call_id(_CITY),
|
||||
"content": [{"type": "text", "text": _TOOL_RESULT_TEXT, **at("tool_result_part")}],
|
||||
},
|
||||
{"role": "user", "content": final_text(marker), **at("string_message")},
|
||||
]
|
||||
|
||||
|
||||
def _body(
|
||||
model: str, marker: str, marks: Mapping[Site, Mapping[str, JsonValue]], **fields: JsonValue
|
||||
) -> dict[str, JsonValue]:
|
||||
return {**chat_body(model, _messages(marker, marks), **fields), "tools": [{**TOOL, **marks.get("tool", {})}]}
|
||||
|
||||
|
||||
def _marked(
|
||||
sites: tuple[Site, ...], mark: Mapping[str, JsonValue] = _BREAKPOINT
|
||||
) -> Mapping[Site, Mapping[str, JsonValue]]:
|
||||
return MappingProxyType({site: mark for site in sites})
|
||||
|
||||
|
||||
def _expected(provider: Provider, site: Site, marker: str) -> str:
|
||||
converse: Final = provider == "bedrock_converse"
|
||||
match site:
|
||||
case "tool":
|
||||
return "tool:lookup_weather"
|
||||
case "system":
|
||||
return f"system:text:{SYSTEM}" if converse else SYSTEM_LABEL
|
||||
case "user_part":
|
||||
return ASK_LABEL
|
||||
case "tool_call":
|
||||
return f"assistant:toolUse:{call_id(_CITY)}" if converse else tool_use_label(_CITY)
|
||||
case "tool_result_part":
|
||||
if converse:
|
||||
return f"user:toolResult:{call_id(_CITY)}"
|
||||
return f"user:tool_result:{call_id(_CITY)}:text:{_TOOL_RESULT_TEXT}"
|
||||
case "string_message":
|
||||
return final_label(marker)
|
||||
case _:
|
||||
assert_never(site)
|
||||
|
||||
|
||||
def _latest_four(provider: Provider, marker: str) -> list[str]:
|
||||
return [_expected(provider, site, marker) for site in _FIVE_SITES[1:]]
|
||||
|
||||
|
||||
def _only_request(wire: Wire) -> Request:
|
||||
received: Final = wire.drain()
|
||||
assert len(received) == 1, [request.target for request in received]
|
||||
return received[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("site", _SITES)
|
||||
@pytest.mark.parametrize("provider", _PROVIDERS)
|
||||
def test_breakpoint_becomes_a_cache_mark_on_the_same_block(gateway: Gateway, provider: Provider, site: Site) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, provider)
|
||||
status, _, text = post_chat(gateway, _body(model, marker, _marked((site,))))
|
||||
assert status == 200, text
|
||||
request: Final = _only_request(wire)
|
||||
assert request.target == _TARGETS[provider], request.target
|
||||
assert b"prompt_cache_breakpoint" not in request.body, request.body
|
||||
assert _labels(provider, request) == [_expected(provider, site, marker)]
|
||||
|
||||
|
||||
def test_breakpoint_ttl_is_not_forwarded(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
marks: Final = _marked(("user_part",), {"prompt_cache_breakpoint": _EXPLICIT_30M})
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, "anthropic")
|
||||
status, _, text = post_chat(gateway, _body(model, marker, marks))
|
||||
assert status == 200, text
|
||||
request: Final = _only_request(wire)
|
||||
assert anthropic_marks(request) == (Mark(ASK_LABEL, None),)
|
||||
assert b"30m" not in request.body, request.body
|
||||
|
||||
|
||||
def test_a_client_cache_control_on_the_same_block_wins_over_the_breakpoint(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
hour: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"}
|
||||
marks: Final = _marked(("user_part",), {"cache_control": hour, "prompt_cache_breakpoint": _EXPLICIT_30M})
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, "anthropic")
|
||||
status, _, text = post_chat(gateway, _body(model, marker, marks))
|
||||
assert status == 200, text
|
||||
request: Final = _only_request(wire)
|
||||
assert anthropic_marks(request) == (Mark(ASK_LABEL, "1h"),)
|
||||
assert b"prompt_cache_breakpoint" not in request.body, request.body
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ("deployment", "request"))
|
||||
@pytest.mark.parametrize("provider", _PROVIDERS)
|
||||
def test_five_breakpoints_keep_the_latest_four_under_drop_params(
|
||||
gateway: Gateway, provider: Provider, source: str
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
deployment: Final[dict[str, JsonValue]] = {"drop_params": True} if source == "deployment" else {}
|
||||
request_fields: Final[dict[str, JsonValue]] = {"drop_params": True} if source == "request" else {}
|
||||
with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, provider, **deployment)
|
||||
status, _, text = post_chat(gateway, _body(model, marker, _marked(_FIVE_SITES), **request_fields))
|
||||
assert status == 200, text
|
||||
request: Final = _only_request(wire)
|
||||
assert _labels(provider, request) == _latest_four(provider, marker)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _PROVIDERS)
|
||||
def test_five_breakpoints_are_a_400_before_any_provider_call_without_drop_params(
|
||||
gateway: Gateway, provider: Provider
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, provider)
|
||||
status, _, text = post_chat(gateway, _body(model, marker, _marked(_FIVE_SITES)))
|
||||
received: Final = wire.drain()
|
||||
assert status == 400, text
|
||||
assert _CAP_HINT in text, text
|
||||
assert "drop_params" in text, text
|
||||
assert received == (), [request.target for request in received]
|
||||
|
||||
|
||||
def test_injection_points_stand_down_behind_the_latest_four_breakpoints(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(
|
||||
gateway, scenario, wire, "anthropic", cache_control_injection_points=POINTS, drop_params=True
|
||||
)
|
||||
status, response_id, text = post_chat(gateway, _body(model, marker, _marked(_FIVE_SITES)))
|
||||
assert status == 200, text
|
||||
assert anthropic_labels(_only_request(wire)) == _latest_four("anthropic", marker)
|
||||
assert gateway_injected(response_id) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _PROVIDERS)
|
||||
def test_prompt_cache_options_are_a_400_before_any_provider_call_without_drop_params(
|
||||
gateway: Gateway, provider: Provider
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, provider)
|
||||
status, _, text = post_chat(
|
||||
gateway, _body(model, marker, _marked(("user_part",)), prompt_cache_options=_EXPLICIT_30M)
|
||||
)
|
||||
received: Final = wire.drain()
|
||||
assert status == 400, text
|
||||
assert "prompt_cache_options" in text, text
|
||||
assert "drop_params" in text, text
|
||||
assert received == (), [request.target for request in received]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _PROVIDERS)
|
||||
def test_prompt_cache_options_are_dropped_under_drop_params(gateway: Gateway, provider: Provider) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, provider, drop_params=True)
|
||||
status, _, text = post_chat(
|
||||
gateway, _body(model, marker, _marked(("user_part",)), prompt_cache_options=_EXPLICIT_30M)
|
||||
)
|
||||
assert status == 200, text
|
||||
request: Final = _only_request(wire)
|
||||
assert b"prompt_cache_options" not in request.body, request.body
|
||||
assert _labels(provider, request) == [ASK_LABEL]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", _PROVIDERS)
|
||||
def test_five_native_cache_control_marks_are_forwarded_and_the_provider_error_reaches_the_caller(
|
||||
gateway: Gateway, provider: Provider
|
||||
) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _deployment(gateway, scenario, wire, provider)
|
||||
status, _, text = post_chat(gateway, _body(model, marker, _marked(_FIVE_SITES, _NATIVE)))
|
||||
received: Final = wire.drain()
|
||||
assert status == 400, text
|
||||
assert "Found 5" in text, text
|
||||
assert _CAP_HINT not in text, text
|
||||
assert [len(_labels(provider, request)) for request in received] == [5], [request.target for request in received]
|
||||
|
||||
|
||||
def _drop_params_config(directory: Path) -> Path:
|
||||
return owned_config(directory, [], litellm_settings={"drop_params": True})
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def drop_params_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]:
|
||||
"""A proxy running with `litellm_settings: drop_params: true`, the setting the cap error names"""
|
||||
with gateway_from_environment() as lane:
|
||||
directory: Final = tmp_path_factory.mktemp("drop_params")
|
||||
with owned_proxy(lane, directory, {}, config=_drop_params_config(directory)) as owned:
|
||||
yield owned
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_proxy_level_drop_params_keeps_the_latest_four_breakpoints(drop_params_gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(anthropic_peer) as wire, drop_params_gateway.scenario() as scenario:
|
||||
model: Final = _deployment(drop_params_gateway, scenario, wire, "anthropic")
|
||||
status, _, text = post_chat(drop_params_gateway, _body(model, marker, _marked(_FIVE_SITES)))
|
||||
assert status == 200, text
|
||||
assert anthropic_labels(_only_request(wire)) == _latest_four("anthropic", marker)
|
||||
|
|
@ -1161,7 +1161,7 @@ def test_cache_control_hook_counts_tool_call_marks_except_answered_server_tool_c
|
|||
},
|
||||
}
|
||||
|
||||
assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 2
|
||||
assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 3
|
||||
|
||||
|
||||
@pytest.mark.parametrize("tool_call_type", ["function", "custom", None])
|
||||
|
|
@ -4279,3 +4279,24 @@ class TestMalformedInjectionPointsAreIgnored:
|
|||
kwargs: dict = {"cache_control_injection_points": copy.deepcopy(value)}
|
||||
_stamp_injection_points_with_dialect(kwargs, "gpt-5.6", "openai")
|
||||
assert kwargs == {"cache_control_injection_points": value}
|
||||
|
||||
|
||||
def test_cache_control_hook_counts_tool_call_prompt_cache_breakpoints():
|
||||
message: Final[AllMessageValues] = {
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": cast(
|
||||
list[ChatCompletionAssistantToolCall],
|
||||
[
|
||||
{
|
||||
"id": f"call_{index}",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
"prompt_cache_breakpoint": {"mode": "explicit"},
|
||||
}
|
||||
for index in range(3)
|
||||
],
|
||||
),
|
||||
}
|
||||
|
||||
assert AnthropicCacheControlHook.count_request_cache_breakpoints([message]) == 3
|
||||
|
|
|
|||
|
|
@ -23,12 +23,18 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
|
|||
hoist_images_from_tool_messages,
|
||||
is_encrypted_reasoning_block,
|
||||
merge_consecutive_system_messages,
|
||||
message_with_prompt_cache_breakpoints_as_cache_control,
|
||||
message_without_cache_control_marks,
|
||||
messages_with_prompt_cache_breakpoints_as_cache_control,
|
||||
parse_tool_call_arguments,
|
||||
prompt_cache_breakpoint_as_cache_control,
|
||||
responses_reasoning_items_from_thinking_blocks,
|
||||
salvage_concatenated_tool_arguments,
|
||||
split_concatenated_json_objects,
|
||||
strip_encrypted_reasoning_from_messages,
|
||||
system_messages_first,
|
||||
tool_with_prompt_cache_breakpoint_as_cache_control,
|
||||
tools_with_prompt_cache_breakpoints_as_cache_control,
|
||||
update_messages_with_model_file_ids,
|
||||
with_prompt_cache_breakpoint,
|
||||
)
|
||||
|
|
@ -2306,3 +2312,129 @@ def test_a_breakpoint_the_client_wrote_survives_a_per_part_cache_control() -> No
|
|||
"text": "rules",
|
||||
"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"},
|
||||
}
|
||||
|
||||
|
||||
_EPHEMERAL: Final = {"type": "ephemeral"}
|
||||
_EXPLICIT: Final = {"mode": "explicit"}
|
||||
_EXPLICIT_30M: Final = {"mode": "explicit", "ttl": "30m"}
|
||||
|
||||
|
||||
def _lookup_call(identity: str, **fields: object) -> dict[str, object]:
|
||||
return {"id": identity, "type": "function", "function": {"name": "lookup", "arguments": "{}"}, **fields}
|
||||
|
||||
|
||||
def test_prompt_cache_breakpoint_becomes_an_ephemeral_cache_control_without_the_ttl() -> None:
|
||||
marked: Final = {"type": "text", "text": "rules", "prompt_cache_breakpoint": _EXPLICIT_30M}
|
||||
|
||||
assert prompt_cache_breakpoint_as_cache_control(marked) == {
|
||||
"type": "text",
|
||||
"text": "rules",
|
||||
"cache_control": _EPHEMERAL,
|
||||
}
|
||||
|
||||
|
||||
def test_a_cache_control_the_client_wrote_wins_over_the_breakpoint_on_the_same_block() -> None:
|
||||
both: Final = {
|
||||
"type": "text",
|
||||
"text": "rules",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
"prompt_cache_breakpoint": _EXPLICIT_30M,
|
||||
}
|
||||
|
||||
assert prompt_cache_breakpoint_as_cache_control(both) == {
|
||||
"type": "text",
|
||||
"text": "rules",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
}
|
||||
|
||||
|
||||
def test_a_block_without_a_breakpoint_is_returned_untouched() -> None:
|
||||
unmarked: Final = {"type": "text", "text": "rules", "cache_control": _EPHEMERAL}
|
||||
|
||||
assert prompt_cache_breakpoint_as_cache_control(unmarked) is unmarked
|
||||
|
||||
|
||||
def test_breakpoints_are_renamed_at_the_message_level_on_parts_and_on_tool_calls() -> None:
|
||||
message: Final = {
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "a", "prompt_cache_breakpoint": _EXPLICIT}, {"type": "text", "text": "b"}],
|
||||
"tool_calls": [_lookup_call("call_1", prompt_cache_breakpoint=_EXPLICIT_30M), _lookup_call("call_2")],
|
||||
"prompt_cache_breakpoint": _EXPLICIT,
|
||||
}
|
||||
|
||||
assert message_with_prompt_cache_breakpoints_as_cache_control(message) == {
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "a", "cache_control": _EPHEMERAL}, {"type": "text", "text": "b"}],
|
||||
"tool_calls": [_lookup_call("call_1", cache_control=_EPHEMERAL), _lookup_call("call_2")],
|
||||
"cache_control": _EPHEMERAL,
|
||||
}
|
||||
|
||||
|
||||
def test_a_message_level_breakpoint_on_string_content_is_renamed_and_other_messages_are_kept() -> None:
|
||||
unmarked: Final = {"role": "user", "content": "unmarked"}
|
||||
|
||||
renamed: Final = messages_with_prompt_cache_breakpoints_as_cache_control(
|
||||
[{"role": "user", "content": "hi", "prompt_cache_breakpoint": _EXPLICIT}, unmarked]
|
||||
)
|
||||
|
||||
assert renamed == ({"role": "user", "content": "hi", "cache_control": _EPHEMERAL}, unmarked)
|
||||
assert renamed[1] is unmarked
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tool",
|
||||
[
|
||||
{"type": "function", "function": {"name": "lookup", "parameters": {}}, "prompt_cache_breakpoint": _EXPLICIT},
|
||||
{"type": "function", "function": {"name": "lookup", "parameters": {}, "prompt_cache_breakpoint": _EXPLICIT}},
|
||||
],
|
||||
ids=["top-level", "under-function"],
|
||||
)
|
||||
def test_a_tool_breakpoint_becomes_a_top_level_cache_control(tool: dict[str, object]) -> None:
|
||||
assert tool_with_prompt_cache_breakpoint_as_cache_control(tool) == {
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "parameters": {}},
|
||||
"cache_control": _EPHEMERAL,
|
||||
}
|
||||
|
||||
|
||||
def test_a_tool_cache_control_the_client_wrote_wins_over_the_breakpoint() -> None:
|
||||
tool: Final = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
"prompt_cache_breakpoint": _EXPLICIT_30M,
|
||||
},
|
||||
}
|
||||
|
||||
assert tool_with_prompt_cache_breakpoint_as_cache_control(tool) == {
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "cache_control": {"type": "ephemeral", "ttl": "1h"}},
|
||||
}
|
||||
|
||||
|
||||
def test_tools_that_are_not_a_list_carry_no_marks_and_unmarked_tools_are_kept() -> None:
|
||||
unmarked: Final = {"type": "function", "function": {"name": "lookup"}}
|
||||
|
||||
assert tools_with_prompt_cache_breakpoints_as_cache_control(None) == ()
|
||||
assert tools_with_prompt_cache_breakpoints_as_cache_control([unmarked, "opaque"]) == (unmarked, "opaque")
|
||||
|
||||
|
||||
def test_message_without_cache_control_marks_strips_only_the_listed_positions() -> None:
|
||||
message: Final = {
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "a", "cache_control": _EPHEMERAL},
|
||||
{"type": "text", "text": "b", "cache_control": _EPHEMERAL},
|
||||
],
|
||||
"tool_calls": [_lookup_call("call_1", cache_control=_EPHEMERAL), _lookup_call("call_2", cache_control=_EPHEMERAL)],
|
||||
"cache_control": _EPHEMERAL,
|
||||
}
|
||||
|
||||
assert message_without_cache_control_marks(
|
||||
message, message_mark=True, parts=frozenset({0}), tool_calls=frozenset({1})
|
||||
) == {
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b", "cache_control": _EPHEMERAL}],
|
||||
"tool_calls": [_lookup_call("call_1", cache_control=_EPHEMERAL), _lookup_call("call_2")],
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8735,3 +8735,274 @@ def test_completion_surfaces_response_citations_in_provider_specific_fields(resp
|
|||
assert response.choices[0].message.provider_specific_fields["citations"] == [
|
||||
[{**citation, "supported_text": "the grass is green"}]
|
||||
]
|
||||
|
||||
|
||||
EXPLICIT: Final = {"mode": "explicit"}
|
||||
EXPLICIT_30M: Final = {"mode": "explicit", "ttl": "30m"}
|
||||
_WEATHER_FUNCTION: Final = {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"parameters": WEATHER_PARAMETERS,
|
||||
}
|
||||
_ANTHROPIC_WEATHER_TOOL: Final = {
|
||||
"name": "get_current_weather",
|
||||
"description": "Get the current weather in a given location",
|
||||
"input_schema": WEATHER_PARAMETERS,
|
||||
"type": "custom",
|
||||
}
|
||||
_ANTHROPIC_REPLY: Final = {
|
||||
"id": "msg_breakpoints",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "text", "text": "Hello!"}],
|
||||
"model": PROMPT_CACHING_MODEL,
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 12, "output_tokens": 6},
|
||||
}
|
||||
|
||||
|
||||
def _five_marked_parts() -> list[dict[str, object]]:
|
||||
return [{"type": "text", "text": f"part {index}", "prompt_cache_breakpoint": EXPLICIT} for index in range(5)]
|
||||
|
||||
|
||||
def _latest_four_marked() -> list[dict[str, object]]:
|
||||
return [
|
||||
{"type": "text", "text": "part 0"},
|
||||
*({"type": "text", "text": f"part {index}", "cache_control": EPHEMERAL} for index in range(1, 5)),
|
||||
]
|
||||
|
||||
|
||||
async def test_prompt_cache_breakpoints_on_system_user_and_tool_definition_become_cache_control(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request: Final = await _send_prompt_caching_request(
|
||||
respx_mock,
|
||||
monkeypatch,
|
||||
messages=[
|
||||
{"role": "system", "content": [{"type": "text", "text": "rules", "prompt_cache_breakpoint": EXPLICIT}]},
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": EXPLICIT_30M}]},
|
||||
],
|
||||
tools=[{"type": "function", "function": {**_WEATHER_FUNCTION, "prompt_cache_breakpoint": EXPLICIT}}],
|
||||
)
|
||||
|
||||
assert json.loads(request.content) == {
|
||||
"model": PROMPT_CACHING_MODEL,
|
||||
"system": [{"type": "text", "text": "rules", "cache_control": EPHEMERAL}],
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": EPHEMERAL}]}],
|
||||
"tools": [{**_ANTHROPIC_WEATHER_TOOL, "cache_control": EPHEMERAL}],
|
||||
"max_tokens": litellm.get_max_tokens(PROMPT_CACHING_MODEL),
|
||||
}
|
||||
|
||||
|
||||
async def test_prompt_cache_breakpoints_on_tool_calls_tool_results_and_string_messages_become_cache_control(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request: Final = await _send_prompt_caching_request(
|
||||
respx_mock,
|
||||
monkeypatch,
|
||||
messages=[
|
||||
{"role": "user", "content": "weather?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_current_weather", "arguments": '{"location": "Boston"}'},
|
||||
"prompt_cache_breakpoint": EXPLICIT,
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": [{"type": "text", "text": "sunny", "prompt_cache_breakpoint": EXPLICIT}],
|
||||
},
|
||||
{"role": "assistant", "content": "It is sunny."},
|
||||
{"role": "user", "content": "thanks", "prompt_cache_breakpoint": EXPLICIT_30M},
|
||||
],
|
||||
tools=[{"type": "function", "function": _WEATHER_FUNCTION}],
|
||||
)
|
||||
|
||||
assert "prompt_cache_breakpoint" not in request.content.decode()
|
||||
assert json.loads(request.content)["messages"] == [
|
||||
{"role": "user", "content": [{"type": "text", "text": "weather?"}]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_use",
|
||||
"id": "call_1",
|
||||
"name": "get_current_weather",
|
||||
"input": {"location": "Boston"},
|
||||
"cache_control": EPHEMERAL,
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": "call_1",
|
||||
"content": [{"type": "text", "text": "sunny", "cache_control": EPHEMERAL}],
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": [{"type": "text", "text": "It is sunny."}]},
|
||||
{"role": "user", "content": [{"type": "text", "text": "thanks", "cache_control": EPHEMERAL}]},
|
||||
]
|
||||
|
||||
|
||||
async def test_a_cache_control_the_client_wrote_wins_over_the_breakpoint_on_the_same_block(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request: Final = await _send_prompt_caching_request(
|
||||
respx_mock,
|
||||
monkeypatch,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "hi",
|
||||
"cache_control": {"type": "ephemeral", "ttl": "1h"},
|
||||
"prompt_cache_breakpoint": EXPLICIT_30M,
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
assert "prompt_cache_breakpoint" not in request.content.decode()
|
||||
assert json.loads(request.content)["messages"] == [
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}]}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("drop_params_source", ["kwarg", "global"])
|
||||
async def test_five_breakpoints_keep_the_latest_four_under_drop_params(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, drop_params_source: str
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", drop_params_source == "global")
|
||||
request: Final = await _send_prompt_caching_request(
|
||||
respx_mock,
|
||||
monkeypatch,
|
||||
messages=[{"role": "user", "content": _five_marked_parts()}],
|
||||
**({"drop_params": True} if drop_params_source == "kwarg" else {}),
|
||||
)
|
||||
|
||||
assert json.loads(request.content)["messages"] == [{"role": "user", "content": _latest_four_marked()}]
|
||||
|
||||
|
||||
async def test_five_breakpoints_are_a_bad_request_before_any_provider_call_without_drop_params(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
|
||||
return_value=httpx.Response(200, json=_ANTHROPIC_REPLY)
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match=r"Claude accepts at most 4 cache breakpoints, found 5") as caught:
|
||||
await litellm.acompletion(
|
||||
api_key="mock_api_key",
|
||||
model=f"anthropic/{PROMPT_CACHING_MODEL}",
|
||||
messages=[{"role": "user", "content": _five_marked_parts()}],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
assert "drop_params=True" in str(caught.value)
|
||||
assert route.call_count == 0
|
||||
|
||||
|
||||
async def test_configured_injection_points_stand_down_when_five_breakpoints_are_cut_to_four(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request: Final = await _send_prompt_caching_request(
|
||||
respx_mock,
|
||||
monkeypatch,
|
||||
messages=[{"role": "system", "content": "rules"}, {"role": "user", "content": _five_marked_parts()}],
|
||||
drop_params=True,
|
||||
cache_control_injection_points=[
|
||||
{"location": "message", "role": "system"},
|
||||
{"location": "message", "index": -1},
|
||||
],
|
||||
)
|
||||
|
||||
body: Final = json.loads(request.content)
|
||||
assert body["system"] == [{"type": "text", "text": "rules"}]
|
||||
assert body["messages"] == [{"role": "user", "content": _latest_four_marked()}]
|
||||
|
||||
|
||||
async def test_prompt_cache_options_are_rejected_before_any_provider_call_without_drop_params(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
|
||||
return_value=httpx.Response(200, json=_ANTHROPIC_REPLY)
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match="drop_params"):
|
||||
await litellm.acompletion(
|
||||
api_key="mock_api_key",
|
||||
model=f"anthropic/{PROMPT_CACHING_MODEL}",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
prompt_cache_options={"mode": "explicit", "ttl": "30m"},
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
assert route.call_count == 0
|
||||
|
||||
|
||||
async def test_prompt_cache_options_are_dropped_under_drop_params(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
request: Final = await _send_prompt_caching_request(
|
||||
respx_mock,
|
||||
monkeypatch,
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
prompt_cache_options={"mode": "explicit", "ttl": "30m"},
|
||||
drop_params=True,
|
||||
)
|
||||
|
||||
assert json.loads(request.content) == {
|
||||
"model": PROMPT_CACHING_MODEL,
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}],
|
||||
"max_tokens": litellm.get_max_tokens(PROMPT_CACHING_MODEL),
|
||||
}
|
||||
|
||||
|
||||
async def test_five_native_cache_control_marks_are_forwarded_and_the_provider_answer_reaches_the_caller(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", True)
|
||||
monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True")
|
||||
marked: Final = [{"type": "text", "text": f"part {index}", "cache_control": EPHEMERAL} for index in range(5)]
|
||||
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
|
||||
return_value=httpx.Response(
|
||||
400,
|
||||
json={
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": "invalid_request_error",
|
||||
"message": "A maximum of 4 blocks with cache_control may be provided. Found 5.",
|
||||
},
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match=r"Found 5"):
|
||||
await litellm.acompletion(
|
||||
api_key="mock_api_key",
|
||||
model=f"anthropic/{PROMPT_CACHING_MODEL}",
|
||||
messages=[{"role": "user", "content": marked}],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
assert route.call_count == 1
|
||||
assert json.loads(route.calls.last.request.content)["messages"] == [{"role": "user", "content": marked}]
|
||||
|
|
|
|||
175
tests/unit/llms/anthropic/chat/test_prompt_cache_breakpoints.py
Normal file
175
tests/unit/llms/anthropic/chat/test_prompt_cache_breakpoints.py
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.llms.anthropic.chat.prompt_cache_breakpoints import (
|
||||
CacheBreakpointCapExceeded,
|
||||
ClaudeCacheMarks,
|
||||
claude_cache_marks,
|
||||
reject_prompt_cache_options,
|
||||
resolved_claude_cache_marks,
|
||||
tools_carry_prompt_cache_breakpoint,
|
||||
)
|
||||
|
||||
EPHEMERAL: Final = {"type": "ephemeral"}
|
||||
EXPLICIT: Final = {"mode": "explicit"}
|
||||
EXPLICIT_30M: Final = {"mode": "explicit", "ttl": "30m"}
|
||||
MODEL: Final = "claude-opus-5-5"
|
||||
|
||||
|
||||
def _text(text: str, **fields: object) -> dict[str, object]:
|
||||
return {"type": "text", "text": text, **fields}
|
||||
|
||||
|
||||
def _tool(**fields: object) -> dict[str, object]:
|
||||
return {"type": "function", "function": {"name": "lookup", "parameters": {}}, **fields}
|
||||
|
||||
|
||||
def _call(identity: str, **fields: object) -> dict[str, object]:
|
||||
return {"id": identity, "type": "function", "function": {"name": "lookup", "arguments": "{}"}, **fields}
|
||||
|
||||
|
||||
def _six_breakpoints() -> tuple[list[dict[str, object]], list[dict[str, object]]]:
|
||||
"""A tool definition, a system part, a user part, a tool call, a tool result part and a string message, every
|
||||
one marked, in the order Claude reads them"""
|
||||
messages: Final = [
|
||||
{"role": "system", "content": [_text("rules", prompt_cache_breakpoint=EXPLICIT)]},
|
||||
{"role": "user", "content": [_text("hi", prompt_cache_breakpoint=EXPLICIT_30M)]},
|
||||
{"role": "assistant", "content": "", "tool_calls": [_call("call_1", prompt_cache_breakpoint=EXPLICIT)]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": [_text("sunny", prompt_cache_breakpoint=EXPLICIT)]},
|
||||
{"role": "user", "content": "done", "prompt_cache_breakpoint": EXPLICIT},
|
||||
]
|
||||
return messages, [_tool(prompt_cache_breakpoint=EXPLICIT)]
|
||||
|
||||
|
||||
def test_drop_params_keeps_the_latest_four_marks_in_wire_order() -> None:
|
||||
messages, tools = _six_breakpoints()
|
||||
|
||||
marks: Final = claude_cache_marks(messages, tools, automatic_breakpoint=False, drop_params=True)
|
||||
|
||||
assert marks == ClaudeCacheMarks(
|
||||
messages=(
|
||||
{"role": "system", "content": [_text("rules")]},
|
||||
{"role": "user", "content": [_text("hi", cache_control=EPHEMERAL)]},
|
||||
{"role": "assistant", "content": "", "tool_calls": [_call("call_1", cache_control=EPHEMERAL)]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": [_text("sunny", cache_control=EPHEMERAL)]},
|
||||
{"role": "user", "content": "done", "cache_control": EPHEMERAL},
|
||||
),
|
||||
tools=(_tool(),),
|
||||
)
|
||||
|
||||
|
||||
def test_four_breakpoints_are_renamed_without_dropping_anything() -> None:
|
||||
messages, tools = _six_breakpoints()
|
||||
|
||||
marks: Final = claude_cache_marks(messages[2:], tools, automatic_breakpoint=False, drop_params=False)
|
||||
|
||||
assert marks == ClaudeCacheMarks(
|
||||
messages=(
|
||||
{"role": "assistant", "content": "", "tool_calls": [_call("call_1", cache_control=EPHEMERAL)]},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": [_text("sunny", cache_control=EPHEMERAL)]},
|
||||
{"role": "user", "content": "done", "cache_control": EPHEMERAL},
|
||||
),
|
||||
tools=(_tool(cache_control=EPHEMERAL),),
|
||||
)
|
||||
|
||||
|
||||
def test_the_automatic_breakpoint_takes_a_slot_and_is_never_dropped() -> None:
|
||||
parts: Final = [_text(f"part {index}", prompt_cache_breakpoint=EXPLICIT) for index in range(4)]
|
||||
|
||||
marks: Final = claude_cache_marks(
|
||||
[{"role": "user", "content": parts}], None, automatic_breakpoint=True, drop_params=True
|
||||
)
|
||||
|
||||
assert marks == ClaudeCacheMarks(
|
||||
messages=(
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
_text("part 0"),
|
||||
*(_text(f"part {index}", cache_control=EPHEMERAL) for index in range(1, 4)),
|
||||
],
|
||||
},
|
||||
),
|
||||
tools=(),
|
||||
)
|
||||
|
||||
|
||||
def test_without_drop_params_the_excess_is_reported_with_the_count() -> None:
|
||||
messages, tools = _six_breakpoints()
|
||||
|
||||
assert claude_cache_marks(messages, tools, automatic_breakpoint=False, drop_params=False) == (
|
||||
CacheBreakpointCapExceeded(found=6)
|
||||
)
|
||||
|
||||
|
||||
def test_a_request_without_breakpoints_is_forwarded_as_written_even_past_the_cap() -> None:
|
||||
messages: Final = [
|
||||
{"role": "user", "content": [_text(f"part {index}", cache_control=EPHEMERAL) for index in range(5)]}
|
||||
]
|
||||
tools: Final = [_tool(cache_control=EPHEMERAL)]
|
||||
|
||||
marks: Final = claude_cache_marks(messages, tools, automatic_breakpoint=True, drop_params=True)
|
||||
|
||||
assert marks == ClaudeCacheMarks(messages=(messages[0],), tools=(tools[0],))
|
||||
assert marks.messages[0] is messages[0]
|
||||
|
||||
|
||||
def test_tools_already_mapped_to_the_anthropic_shape_still_arm_the_cap() -> None:
|
||||
mapped_tools: Final = [
|
||||
{"name": f"tool_{index}", "input_schema": {}, "cache_control": EPHEMERAL} for index in range(4)
|
||||
]
|
||||
messages: Final = [{"role": "user", "content": [_text("hi", prompt_cache_breakpoint=EXPLICIT)]}]
|
||||
|
||||
assert claude_cache_marks(
|
||||
messages, mapped_tools, automatic_breakpoint=False, drop_params=False, tools_carried_breakpoint=True
|
||||
) == CacheBreakpointCapExceeded(found=5)
|
||||
assert claude_cache_marks(
|
||||
messages, mapped_tools, automatic_breakpoint=False, drop_params=True, tools_carried_breakpoint=True
|
||||
) == ClaudeCacheMarks(
|
||||
messages=({"role": "user", "content": [_text("hi", cache_control=EPHEMERAL)]},),
|
||||
tools=({"name": "tool_0", "input_schema": {}}, *mapped_tools[1:]),
|
||||
)
|
||||
|
||||
|
||||
def test_the_cap_error_is_a_public_bad_request_naming_drop_params() -> None:
|
||||
with pytest.raises(litellm.BadRequestError) as caught:
|
||||
resolved_claude_cache_marks(CacheBreakpointCapExceeded(found=6), model=MODEL, llm_provider="anthropic")
|
||||
|
||||
assert "Claude accepts at most 4 cache breakpoints, found 6" in str(caught.value)
|
||||
assert "drop_params=True" in str(caught.value)
|
||||
assert caught.value.status_code == 400
|
||||
assert caught.value.llm_provider == "anthropic"
|
||||
|
||||
|
||||
def test_marks_within_the_cap_resolve_to_themselves() -> None:
|
||||
marks: Final = ClaudeCacheMarks(messages=(), tools=())
|
||||
|
||||
assert resolved_claude_cache_marks(marks, model=MODEL, llm_provider="anthropic") is marks
|
||||
|
||||
|
||||
def test_prompt_cache_options_are_rejected_unless_drop_params() -> None:
|
||||
with pytest.raises(litellm.UnsupportedParamsError) as caught:
|
||||
reject_prompt_cache_options(model=MODEL, options={"mode": "explicit"}, drop_params=False)
|
||||
|
||||
assert "prompt_cache_options" in str(caught.value)
|
||||
assert "drop_params" in str(caught.value)
|
||||
assert caught.value.status_code == 400
|
||||
assert reject_prompt_cache_options(model=MODEL, options={"mode": "explicit"}, drop_params=True) is None
|
||||
assert reject_prompt_cache_options(model=MODEL, options=None, drop_params=False) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tools", "expected"),
|
||||
[
|
||||
([_tool(prompt_cache_breakpoint=EXPLICIT)], True),
|
||||
([{"type": "function", "function": {"name": "lookup", "prompt_cache_breakpoint": EXPLICIT}}], True),
|
||||
([_tool(cache_control=EPHEMERAL), _tool()], False),
|
||||
(None, False),
|
||||
("not a list", False),
|
||||
],
|
||||
ids=["top-level", "under-function", "cache-control-only", "none", "opaque"],
|
||||
)
|
||||
def test_tools_carry_prompt_cache_breakpoint_reads_both_tool_shapes(tools: object, expected: bool) -> None:
|
||||
assert tools_carry_prompt_cache_breakpoint(tools) is expected
|
||||
|
|
@ -4267,3 +4267,18 @@ def test_anthropic_error_frame_exception_classes_an_overloaded_frame_as_internal
|
|||
error: Final = anthropic_error_frame_exception("overloaded_error", "Overloaded", 503, "claude-sonnet-4-5")
|
||||
|
||||
assert type(error) is litellm.InternalServerError
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"messages",
|
||||
[
|
||||
[{"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": {"mode": "explicit"}}]}],
|
||||
[{"role": "user", "content": "hi", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}],
|
||||
[{"role": "system", "content": [{"type": "text", "text": "rules", "prompt_cache_breakpoint": {"mode": "explicit"}}]}],
|
||||
],
|
||||
ids=["content-part", "message-level-on-string", "system-part"],
|
||||
)
|
||||
def test_is_cache_control_set_counts_a_prompt_cache_breakpoint_like_a_cache_control(messages: list) -> None:
|
||||
from litellm.llms.anthropic.common_utils import AnthropicModelInfo
|
||||
|
||||
assert AnthropicModelInfo().is_cache_control_set(messages) is True
|
||||
|
|
|
|||
|
|
@ -7,6 +7,8 @@ sys.path.insert(
|
|||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.llms.azure_ai.anthropic.transformation import AzureAnthropicConfig
|
||||
|
|
@ -484,3 +486,23 @@ def test_chat_flagged_model_keeps_mid_conversation_system_role_in_place(local_mo
|
|||
"role": "system",
|
||||
"content": [{"type": "text", "text": "<system-reminder>Answer with exactly one word.</system-reminder>"}],
|
||||
}
|
||||
|
||||
|
||||
def test_prompt_cache_breakpoint_becomes_cache_control_on_the_same_block():
|
||||
result = AzureAnthropicConfig().transform_request(
|
||||
model="claude-opus-5-5",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}],
|
||||
}
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert "prompt_cache_breakpoint" not in json.dumps(result)
|
||||
assert result["messages"] == [
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral"}}]}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1351,6 +1351,92 @@ class TestBedrockAnthropicPromptCachingRegression:
|
|||
has_cache_control = any(isinstance(c, dict) and "cache_control" in c for c in user_msg["content"])
|
||||
assert has_cache_control, "cache_control should be present in invoke messages"
|
||||
|
||||
@staticmethod
|
||||
def _claude_config(model_prefix: str):
|
||||
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeConfig,
|
||||
)
|
||||
|
||||
return AmazonConverseConfig() if "converse" in model_prefix else AmazonAnthropicClaudeConfig()
|
||||
|
||||
@pytest.mark.parametrize("model_prefix", ["bedrock/invoke/", "bedrock/converse/"])
|
||||
def test_prompt_cache_breakpoint_lands_on_the_same_block_for_both_routes(self, model_prefix: str) -> None:
|
||||
messages: Final = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": LARGE_DOCUMENT_FOR_CACHING,
|
||||
"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"},
|
||||
},
|
||||
{"type": "text", "text": "What are the payment terms?"},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
result: Final = self._claude_config(model_prefix).transform_request(
|
||||
model="anthropic.claude-opus-5-5", messages=messages, optional_params={}, litellm_params={}, headers={}
|
||||
)
|
||||
|
||||
assert "prompt_cache_breakpoint" not in json.dumps(result)
|
||||
if "converse" in model_prefix:
|
||||
assert result["messages"][0]["content"] == [
|
||||
{"text": LARGE_DOCUMENT_FOR_CACHING},
|
||||
{"cachePoint": {"type": "default"}},
|
||||
{"text": "What are the payment terms?"},
|
||||
]
|
||||
else:
|
||||
assert result["messages"][0]["content"] == [
|
||||
{"type": "text", "text": LARGE_DOCUMENT_FOR_CACHING, "cache_control": {"type": "ephemeral"}},
|
||||
{"type": "text", "text": "What are the payment terms?"},
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("model_prefix", ["bedrock/invoke/", "bedrock/converse/"])
|
||||
def test_five_breakpoints_are_cut_to_four_only_under_drop_params_for_both_routes(
|
||||
self, model_prefix: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
import litellm
|
||||
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
marked: Final = [
|
||||
{"type": "text", "text": f"part {index}", "prompt_cache_breakpoint": {"mode": "explicit"}}
|
||||
for index in range(5)
|
||||
]
|
||||
config: Final = self._claude_config(model_prefix)
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match=r"Claude accepts at most 4 cache breakpoints, found 5"):
|
||||
config.transform_request(
|
||||
model="anthropic.claude-opus-5-5",
|
||||
messages=[{"role": "user", "content": marked}],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
result: Final = config.transform_request(
|
||||
model="anthropic.claude-opus-5-5",
|
||||
messages=[{"role": "user", "content": marked}],
|
||||
optional_params={},
|
||||
litellm_params={"drop_params": True},
|
||||
headers={},
|
||||
)
|
||||
if "converse" in model_prefix:
|
||||
assert result["messages"][0]["content"] == [
|
||||
{"text": "part 0"},
|
||||
*(
|
||||
block
|
||||
for index in range(1, 5)
|
||||
for block in ({"text": f"part {index}"}, {"cachePoint": {"type": "default"}})
|
||||
),
|
||||
]
|
||||
else:
|
||||
assert result["messages"][0]["content"] == [
|
||||
{"type": "text", "text": "part 0"},
|
||||
*({"type": "text", "text": f"part {index}", "cache_control": {"type": "ephemeral"}} for index in range(1, 5)),
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_prefix",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -10145,3 +10145,162 @@ async def test_bedrock_user_file_element_becomes_document_block(sync_mode: bool)
|
|||
assert content[0] == {"text": "Summarize this report"}
|
||||
assert content[1]["document"]["source"] == {"bytes": base64.b64encode(pdf_bytes).decode()}
|
||||
assert content[1]["document"]["format"] == "pdf"
|
||||
|
||||
|
||||
_BREAKPOINT_MODEL: Final = "bedrock/anthropic.claude-opus-5-5"
|
||||
_EXPLICIT: Final = {"mode": "explicit"}
|
||||
_EXPLICIT_30M: Final = {"mode": "explicit", "ttl": "30m"}
|
||||
_CACHE_POINT: Final = {"cachePoint": {"type": "default"}}
|
||||
_LOOKUP_TOOL: Final = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "lookup",
|
||||
"description": "lookup",
|
||||
"parameters": {"type": "object", "properties": {}, "required": []},
|
||||
},
|
||||
}
|
||||
_LOOKUP_TOOL_SPEC: Final = {
|
||||
"toolSpec": {
|
||||
"name": "lookup",
|
||||
"description": "lookup",
|
||||
"inputSchema": {"json": {"type": "object", "properties": {}, "required": []}},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _breakpoint_conversation() -> list[dict[str, object]]:
|
||||
return [
|
||||
{"role": "system", "content": [{"type": "text", "text": "rules", "prompt_cache_breakpoint": _EXPLICIT}]},
|
||||
{"role": "user", "content": [{"type": "text", "text": "hi", "prompt_cache_breakpoint": _EXPLICIT_30M}]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": "{}"},
|
||||
"prompt_cache_breakpoint": _EXPLICIT,
|
||||
}
|
||||
],
|
||||
},
|
||||
{"role": "tool", "tool_call_id": "call_1", "content": [{"type": "text", "text": "sunny", "prompt_cache_breakpoint": _EXPLICIT}]},
|
||||
{"role": "user", "content": "done", "prompt_cache_breakpoint": _EXPLICIT},
|
||||
]
|
||||
|
||||
|
||||
def _expected_latest_four() -> dict[str, object]:
|
||||
return {
|
||||
"system": [{"text": "rules"}],
|
||||
"toolConfig": {"tools": [_LOOKUP_TOOL_SPEC]},
|
||||
"messages": [
|
||||
{"role": "user", "content": [{"text": "hi"}, _CACHE_POINT]},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [{"toolUse": {"toolUseId": "call_1", "name": "lookup", "input": {}}}, _CACHE_POINT],
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"toolResult": {"content": [{"text": "sunny"}], "toolUseId": "call_1"}},
|
||||
_CACHE_POINT,
|
||||
{"text": "done"},
|
||||
_CACHE_POINT,
|
||||
],
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
async def _breakpoint_request(
|
||||
sync_mode: bool, optional_params: dict[str, object], litellm_params: dict[str, object]
|
||||
) -> dict[str, object]:
|
||||
config: Final = AmazonConverseConfig()
|
||||
if sync_mode:
|
||||
return config._transform_request(
|
||||
model=_BREAKPOINT_MODEL,
|
||||
messages=_breakpoint_conversation(),
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers={},
|
||||
)
|
||||
return await config._async_transform_request(
|
||||
model=_BREAKPOINT_MODEL,
|
||||
messages=_breakpoint_conversation(),
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
headers={},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_cache_breakpoints_become_cache_points_after_the_same_blocks(sync_mode: bool) -> None:
|
||||
conversation: Final = _breakpoint_conversation()
|
||||
config: Final = AmazonConverseConfig()
|
||||
optional_params: Final[dict[str, object]] = {"tools": [{**_LOOKUP_TOOL, "prompt_cache_breakpoint": _EXPLICIT}]}
|
||||
if sync_mode:
|
||||
result = config._transform_request(
|
||||
model=_BREAKPOINT_MODEL, messages=conversation[2:], optional_params=optional_params, litellm_params={}, headers={}
|
||||
)
|
||||
else:
|
||||
result = await config._async_transform_request(
|
||||
model=_BREAKPOINT_MODEL, messages=conversation[2:], optional_params=optional_params, litellm_params={}, headers={}
|
||||
)
|
||||
|
||||
assert "prompt_cache_breakpoint" not in json.dumps(result)
|
||||
assert result["toolConfig"]["tools"] == [_LOOKUP_TOOL_SPEC, _CACHE_POINT]
|
||||
assert result["messages"] == _expected_latest_four()["messages"][1:]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_six_breakpoints_keep_the_latest_four_cache_points_under_drop_params(sync_mode: bool) -> None:
|
||||
result: Final = await _breakpoint_request(
|
||||
sync_mode,
|
||||
{"tools": [{**_LOOKUP_TOOL, "prompt_cache_breakpoint": _EXPLICIT}]},
|
||||
{"drop_params": True},
|
||||
)
|
||||
|
||||
assert json.dumps(result).count("cachePoint") == 4
|
||||
assert {key: result[key] for key in ("system", "toolConfig", "messages")} == _expected_latest_four()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_six_breakpoints_are_a_bad_request_without_drop_params(sync_mode: bool, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
|
||||
with pytest.raises(litellm.BadRequestError, match=r"Claude accepts at most 4 cache breakpoints, found 6"):
|
||||
await _breakpoint_request(sync_mode, {"tools": [{**_LOOKUP_TOOL, "prompt_cache_breakpoint": _EXPLICIT}]}, {})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False], ids=["sync", "async"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_cache_options_are_rejected_unless_drop_params(sync_mode: bool, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "drop_params", False)
|
||||
|
||||
with pytest.raises(litellm.UnsupportedParamsError, match="drop_params"):
|
||||
await _breakpoint_request(sync_mode, {"prompt_cache_options": {"mode": "explicit"}}, {})
|
||||
|
||||
dropped: Final = await _breakpoint_request(sync_mode, {"prompt_cache_options": {"mode": "explicit"}}, {"drop_params": True})
|
||||
assert "prompt_cache_options" not in json.dumps(dropped)
|
||||
|
||||
|
||||
def test_five_native_cache_control_marks_are_forwarded_as_written_past_the_cap() -> None:
|
||||
marked: Final = [{"type": "text", "text": f"part {index}", "cache_control": {"type": "ephemeral"}} for index in range(5)]
|
||||
|
||||
result: Final = AmazonConverseConfig()._transform_request(
|
||||
model=_BREAKPOINT_MODEL,
|
||||
messages=[{"role": "user", "content": marked}],
|
||||
optional_params={},
|
||||
litellm_params={"drop_params": True},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert result["messages"] == [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [block for index in range(5) for block in ({"text": f"part {index}"}, _CACHE_POINT)],
|
||||
}
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1090,3 +1090,64 @@ def test_chat_flagged_model_replays_a_byte_identical_prefix_around_a_mid_convers
|
|||
_assert_prefix_stable(requests)
|
||||
assert [m["role"] for m in requests[1]["messages"]] == ["user", "assistant", "user", "system"]
|
||||
assert [m["role"] for m in requests[2]["messages"]] == ["user", "assistant", "user", "system", "assistant", "user"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sync_mode", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_ai_anthropic_prompt_cache_breakpoint_becomes_cache_control_on_the_same_block(
|
||||
respx_mock: respx.MockRouter,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
sync_mode: bool,
|
||||
) -> None:
|
||||
from litellm.main import vertex_partner_models_chat_completion
|
||||
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
monkeypatch.setitem(
|
||||
vertex_partner_models_chat_completion._credentials_project_mapping,
|
||||
("test-vertex-credentials", "test-project"),
|
||||
(Credentials(token="test-vertex-token"), "test-project"),
|
||||
)
|
||||
route: Final = respx_mock.post(
|
||||
"https://us-east5-aiplatform.googleapis.com/v1/projects/test-project/locations/us-east5/"
|
||||
"publishers/anthropic/models/claude-opus-5-5:rawPredict"
|
||||
).mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "msg-breakpoint",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-opus-5-5",
|
||||
"content": [{"type": "text", "text": "cached"}],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {"input_tokens": 4, "output_tokens": 1},
|
||||
},
|
||||
)
|
||||
)
|
||||
call_params: Final = {
|
||||
"model": "vertex_ai/claude-opus-5-5",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Cache this request", "prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}
|
||||
],
|
||||
}
|
||||
],
|
||||
"vertex_project": "test-project",
|
||||
"vertex_location": "us-east5",
|
||||
"vertex_credentials": "test-vertex-credentials",
|
||||
"max_retries": 0,
|
||||
}
|
||||
|
||||
if sync_mode:
|
||||
litellm.completion(**call_params)
|
||||
else:
|
||||
await litellm.acompletion(**call_params)
|
||||
|
||||
assert route.call_count == 1
|
||||
request: Final = route.calls[0].request
|
||||
assert "prompt_cache_breakpoint" not in request.content.decode()
|
||||
assert json.loads(request.content)["messages"] == [
|
||||
{"role": "user", "content": [{"type": "text", "text": "Cache this request", "cache_control": {"type": "ephemeral"}}]}
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue