diff --git a/tests/integration/providers/_claude_prompt_cache_breakpoint_support.py b/tests/integration/providers/_claude_prompt_cache_breakpoint_support.py new file mode 100644 index 00000000000..4d66f7de955 --- /dev/null +++ b/tests/integration/providers/_claude_prompt_cache_breakpoint_support.py @@ -0,0 +1,179 @@ +from collections.abc import Iterator +from typing import Final, Literal + +from integration._support.client import Gateway, Scenario +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire +from integration.providers._cache_control_marks_support import ( + ANTHROPIC_MODEL, + BEDROCK_MODEL, + PROVIDER_KEY, + Mark, + anthropic_labels, + anthropic_peer, + bedrock_labels, + bedrock_peer, +) +from pydantic import JsonValue, TypeAdapter +from typing_extensions import assert_never + +Provider = Literal[ + "anthropic", + "bedrock_converse", + "bedrock_invoke", + "bedrock_claude_platform", + "bedrock_mantle", + "vertex_ai", + "azure_ai", +] + +PROVIDERS: Final[tuple[Provider, ...]] = ( + "anthropic", + "bedrock_converse", + "bedrock_invoke", + "bedrock_claude_platform", + "bedrock_mantle", + "vertex_ai", + "azure_ai", +) +ANTHROPIC_BODY_PROVIDERS: Final = tuple(provider for provider in PROVIDERS if provider != "bedrock_converse") +_BEDROCK_TOKEN: Final = "synthetic-bedrock-bearer" +_MANTLE_TOKEN: Final = "synthetic-mantle-bearer" +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def respond(provider: Provider, request: Request) -> Reply: + return bedrock_peer(request) if provider == "bedrock_converse" else anthropic_peer(request) + + +def labels(provider: Provider, request: Request) -> list[str]: + return bedrock_labels(request) if provider == "bedrock_converse" else anthropic_labels(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 "bedrock_claude_platform": + return { + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_base": wire.url, + "api_key": PROVIDER_KEY, + "workspace_id": "wrkspc_scripted", + "aws_region_name": "us-east-1", + } + case "bedrock_mantle": + return {"model": f"bedrock_mantle/{BEDROCK_MODEL}", "api_base": wire.url, "api_key": _MANTLE_TOKEN} + 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 only_request(wire: Wire) -> Request: + received: Final = wire.drain() + assert len(received) == 1, [request.target for request in received] + return received[0] + + +def _ttl(mark: JsonValue) -> str | None: + ttl: Final = mark.get("ttl") if isinstance(mark, dict) else None + return ttl if isinstance(ttl, str) else None + + +def _anthropic_label(block: dict[str, JsonValue]) -> str: + match block.get("type"): + case "text": + return f"text:{block.get('text')}" + case "tool_use": + return f"call:{block.get('id')}" + case "tool_result": + return f"result:{block.get('tool_use_id')}" + case _: + return f"tool:{block.get('name')}" + + +def _anthropic_marks(value: JsonValue) -> Iterator[Mark]: + if isinstance(value, list): + for item in value: + yield from _anthropic_marks(item) + return + if not isinstance(value, dict): + return + if "cache_control" in value: + yield Mark(_anthropic_label(value), _ttl(value["cache_control"])) + yield from _anthropic_marks(value.get("content")) + + +def _converse_label(block: JsonValue) -> str: + if not isinstance(block, dict): + return "start" + if isinstance(block.get("text"), str): + return f"text:{block['text']}" + for kind, label in (("toolUse", "call"), ("toolResult", "result")): + inner = block.get(kind) + if isinstance(inner, dict): + return f"{label}:{inner.get('toolUseId')}" + spec: Final = block.get("toolSpec") + return f"tool:{spec.get('name')}" if isinstance(spec, dict) else "other" + + +def _converse_marks(blocks: JsonValue) -> Iterator[Mark]: + listed: Final = blocks if isinstance(blocks, list) else [] + for previous, block in zip([None, *listed], listed): + if not isinstance(block, dict): + continue + result = block.get("toolResult") + if isinstance(result, dict): + yield from _converse_marks(result.get("content")) + if "cachePoint" in block: + yield Mark(_converse_label(previous), _ttl(block["cachePoint"])) + + +def _messages(body: dict[str, JsonValue]) -> list[JsonValue]: + messages: Final = body.get("messages") + return messages if isinstance(messages, list) else [] + + +def sent_marks(provider: Provider, request: Request) -> tuple[Mark, ...]: + """Every cache mark the provider received in wire order, nested tool result content included""" + body: Final = _JSON_OBJECT.validate_json(request.body) + contents: Final = [message.get("content") for message in _messages(body) if isinstance(message, dict)] + if provider == "bedrock_converse": + tool_config: Final = body.get("toolConfig") + return ( + *_converse_marks(tool_config.get("tools") if isinstance(tool_config, dict) else None), + *_converse_marks(body.get("system")), + *(mark for content in contents for mark in _converse_marks(content)), + ) + return (*_anthropic_marks(body.get("tools")), *_anthropic_marks(body.get("system")), *_anthropic_marks(contents)) diff --git a/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_placement_wire.py b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_placement_wire.py new file mode 100644 index 00000000000..f19166a1691 --- /dev/null +++ b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_placement_wire.py @@ -0,0 +1,488 @@ +import json +from collections.abc import Callable +from typing import Final, Literal + +import pytest +from integration._support.client import Gateway +from integration._support.wire import wire_server +from integration.providers._cache_control_marks_support import ( + ASK, + CITIES, + EPHEMERAL, + SYSTEM, + TOOL, + Mark, + call_id, + chat_body, + final_text, + new_marker, + post_chat, + tool_call, +) +from integration.providers._claude_prompt_cache_breakpoint_support import ( + ANTHROPIC_BODY_PROVIDERS, + PROVIDERS, + Provider, + deployment, + only_request, + respond, + sent_marks, +) +from pydantic import JsonValue + +Role = Literal["system", "user", "assistant"] +Shape = Literal["list_content", "developer"] + +_EXPLICIT: Final[dict[str, JsonValue]] = {"mode": "explicit"} +_ONE_HOUR_BREAKPOINT: Final[dict[str, JsonValue]] = {"mode": "explicit", "ttl": "1h"} +_THIRTY_MINUTES: Final[dict[str, JsonValue]] = {"mode": "explicit", "ttl": "30m"} +_ONE_HOUR: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} +_LONDON: Final = CITIES[0] +_PARIS: Final = CITIES[1] +_FIRST_PART: Final = "first part" +_LAST_PART: Final = "last part" +_CLOUDY: Final = "cloudy" +_SUNNY: Final = "sunny" +_CHECKING: Final = "Checking." +_NOTES: Final = ("note 0", "note 1", "note 2") +_MALFORMED: Final[tuple[JsonValue, ...]] = (5, [], "", "x" * 5000) +_MALFORMED_IDS: Final = ("int", "list", "empty_string", "5kb_string") +_SEARCH_TOOLS: Final[list[JsonValue]] = [ + {"type": "tool_search", "execution": "server"}, + {**TOOL, "defer_loading": True}, +] + +Body = Callable[[str], dict[str, JsonValue]] + + +def _text(text: str, **fields: JsonValue) -> dict[str, JsonValue]: + return {"type": "text", "text": text, **fields} + + +def _label(text: str) -> str: + return f"text:{text}" + + +def _result_label(city: str) -> str: + return f"result:{call_id(city)}" + + +def _call_label(city: str) -> str: + return f"call:{call_id(city)}" + + +def _note_labels() -> list[str]: + return [_label(note) for note in _NOTES] + + +def _final(marker: str, **fields: JsonValue) -> dict[str, JsonValue]: + if not fields: + return {"role": "user", "content": final_text(marker)} + return {"role": "user", "content": [_text(final_text(marker), **fields)]} + + +def _marks_sent(gateway: Gateway, provider: Provider, body: Body, **deployment_fields: JsonValue) -> tuple[Mark, ...]: + with wire_server(lambda request: respond(provider, request)) as wire, gateway.scenario() as scenario: + model: Final = deployment(gateway, scenario, wire, provider, **deployment_fields) + status, _, text = post_chat(gateway, body(model)) + assert status == 200, text + request: Final = only_request(wire) + assert b"prompt_cache_breakpoint" not in request.body, request.body + return sent_marks(provider, request) + + +def _rejected(gateway: Gateway, provider: Provider, body: Body) -> str: + 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)) + received: Final = wire.drain() + assert status == 400, text + assert "prompt_cache_breakpoint" in text and "drop_params" in text, text + assert received == (), [request.target for request in received] + return text + + +def _list_content_turns(role: Role, marker: str, mark: JsonValue, last: dict[str, JsonValue]) -> list[JsonValue]: + marked: Final[dict[str, JsonValue]] = { + "role": role, + "content": [_text(_FIRST_PART), {**_text(_LAST_PART), **last}], + "prompt_cache_breakpoint": mark, + } + match role: + case "system": + return [marked, _final(marker)] + case "user": + return [ + {"role": "system", "content": SYSTEM}, + marked, + {"role": "assistant", "content": "noted"}, + _final(marker), + ] + case "assistant": + return [{"role": "user", "content": ASK}, marked, _final(marker)] + + +def _developer_turns(marker: str, mark: JsonValue) -> list[JsonValue]: + return [{"role": "developer", "content": SYSTEM, "prompt_cache_breakpoint": mark}, _final(marker)] + + +def _shape_turns(shape: Shape, marker: str, mark: JsonValue) -> list[JsonValue]: + if shape == "developer": + return _developer_turns(marker, mark) + return _list_content_turns("user", marker, mark, {}) + + +def _tool_result_turns(marker: str, result: dict[str, JsonValue], *, leading: bool = False) -> list[JsonValue]: + lead: Final[list[JsonValue]] = [ + {"role": "system", "content": [_text(SYSTEM, prompt_cache_breakpoint=_EXPLICIT)]}, + {"role": "user", "content": [_text(ASK, prompt_cache_breakpoint=_EXPLICIT)]}, + ] + return [ + *(lead if leading else [{"role": "user", "content": ASK}]), + {"role": "assistant", "content": None, "tool_calls": [tool_call(_LONDON)]}, + {"role": "tool", "tool_call_id": call_id(_LONDON), **result}, + _final(marker, prompt_cache_breakpoint=_EXPLICIT) if leading else _final(marker), + ] + + +def _tool_call_turns(marker: str, **assistant_fields: JsonValue) -> list[JsonValue]: + return [ + {"role": "user", "content": ASK}, + {"role": "assistant", "tool_calls": [tool_call(_LONDON), tool_call(_PARIS)], **assistant_fields}, + {"role": "tool", "tool_call_id": call_id(_LONDON), "content": "sunny."}, + {"role": "tool", "tool_call_id": call_id(_PARIS), "content": "sunny."}, + _final(marker), + ] + + +def _noted_turns(marker: str, *history: JsonValue) -> list[JsonValue]: + return [ + {"role": "system", "content": [_text(SYSTEM, prompt_cache_breakpoint=_EXPLICIT)]}, + {"role": "user", "content": [_text(note, prompt_cache_breakpoint=_EXPLICIT) for note in _NOTES]}, + *history, + _final(marker), + ] + + +@pytest.mark.parametrize("role", ("system", "user", "assistant")) +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_message_breakpoint_on_list_content_marks_its_last_part( + gateway: Gateway, provider: Provider, role: Role +) -> None: + marker: Final = new_marker() + marks: Final = _marks_sent( + gateway, provider, lambda model: chat_body(model, _list_content_turns(role, marker, _EXPLICIT, {})) + ) + assert marks == (Mark(_label(_LAST_PART), None),) + + +@pytest.mark.parametrize("content_fields", ({"content": None}, {}), ids=("null_content", "missing_content")) +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_message_breakpoint_on_a_tool_call_turn_without_content_marks_its_last_call( + gateway: Gateway, provider: Provider, content_fields: dict[str, JsonValue] +) -> None: + marker: Final = new_marker() + turns: Final = _tool_call_turns(marker, **content_fields, prompt_cache_breakpoint=_EXPLICIT) + marks: Final = _marks_sent(gateway, provider, lambda model: chat_body(model, turns)) + assert marks == (Mark(_call_label(_PARIS), None),) + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_message_breakpoint_on_an_empty_string_tool_call_turn_lands_where_a_client_cache_control_would( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + breakpoint_turns: Final = _tool_call_turns(marker, content="", prompt_cache_breakpoint=_EXPLICIT) + cache_control_turns: Final = _tool_call_turns(marker, content="", cache_control=EPHEMERAL) + breakpoint_marks: Final = _marks_sent(gateway, provider, lambda model: chat_body(model, breakpoint_turns)) + cache_control_marks: Final = _marks_sent(gateway, provider, lambda model: chat_body(model, cache_control_turns)) + assert breakpoint_marks == cache_control_marks + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_developer_message_breakpoint_marks_its_text(gateway: Gateway, provider: Provider) -> None: + marker: Final = new_marker() + marks: Final = _marks_sent(gateway, provider, lambda model: chat_body(model, _developer_turns(marker, _EXPLICIT))) + assert marks == (Mark(_label(SYSTEM), None),) + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_tool_results_own_breakpoint_replaces_its_part_marks_and_joins_their_one_hour_ttl( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + result: Final[dict[str, JsonValue]] = { + "content": [_text(_CLOUDY), _text(_SUNNY, prompt_cache_breakpoint=_ONE_HOUR_BREAKPOINT)], + "prompt_cache_breakpoint": _EXPLICIT, + } + marks: Final = _marks_sent(gateway, provider, lambda model: chat_body(model, _tool_result_turns(marker, result))) + assert marks == (Mark(_result_label(_LONDON), "1h"),) + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_client_cache_control_on_a_tool_result_wins_over_its_part_breakpoints( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + result: Final[dict[str, JsonValue]] = { + "content": [_text(_CLOUDY), _text(_SUNNY, prompt_cache_breakpoint=_ONE_HOUR_BREAKPOINT)], + "cache_control": EPHEMERAL, + } + marks: Final = _marks_sent(gateway, provider, lambda model: chat_body(model, _tool_result_turns(marker, result))) + assert marks == (Mark(_result_label(_LONDON), None),) + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_client_cache_control_on_a_tool_result_part_wins_over_the_results_own_breakpoint( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + result: Final[dict[str, JsonValue]] = { + "content": [_text(_CLOUDY), _text(_SUNNY, cache_control=EPHEMERAL)], + "prompt_cache_breakpoint": _EXPLICIT, + } + marks: Final = _marks_sent(gateway, provider, lambda model: chat_body(model, _tool_result_turns(marker, result))) + sent_on: Final = _result_label(_LONDON) if provider == "bedrock_converse" else _label(_SUNNY) + assert marks == (Mark(sent_on, None),) + + +@pytest.mark.parametrize("client_marked", (False, True), ids=("breakpoints", "client_marked_part")) +@pytest.mark.parametrize("provider", PROVIDERS) +def test_part_marks_on_one_tool_result_take_one_slot_where_the_route_sends_one_mark_under_drop_params( + gateway: Gateway, provider: Provider, client_marked: bool +) -> None: + marker: Final = new_marker() + first: Final = ( + _text(_CLOUDY, cache_control=EPHEMERAL) if client_marked else _text(_CLOUDY, prompt_cache_breakpoint=_EXPLICIT) + ) + result: Final[dict[str, JsonValue]] = {"content": [first, _text(_SUNNY, prompt_cache_breakpoint=_EXPLICIT)]} + marks: Final = _marks_sent( + gateway, + provider, + lambda model: chat_body(model, _tool_result_turns(marker, result, leading=True)), + drop_params=True, + ) + one_mark_per_result: Final = [_label(SYSTEM), _label(ASK), _result_label(_LONDON), _label(final_text(marker))] + mark_per_part: Final = [_label(ASK), _label(_CLOUDY), _label(_SUNNY), _label(final_text(marker))] + assert [mark.label for mark in marks] == (one_mark_per_result if provider == "bedrock_converse" else mark_per_part) + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_breakpoint_on_an_unanswered_tool_search_call_takes_no_slot_under_drop_params( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + search_call: Final[dict[str, JsonValue]] = { + "id": f"srvtoolu_{marker[:24]}", + "type": "function", + "function": {"name": "tool_search_tool_regex", "arguments": json.dumps({"query": "weather"})}, + "prompt_cache_breakpoint": _EXPLICIT, + } + turns: Final = _noted_turns(marker, {"role": "assistant", "content": None, "tool_calls": [search_call]}) + marks: Final = _marks_sent( + gateway, + provider, + lambda model: {**chat_body(model, turns, drop_params=True), "tools": _SEARCH_TOOLS}, + ) + assert [mark.label for mark in marks] == [_label(SYSTEM), *_note_labels()] + + +def test_breakpoints_on_tool_history_converse_sends_as_text_take_no_slot_under_drop_params(gateway: Gateway) -> None: + marker: Final = new_marker() + turns: Final = _noted_turns( + marker, + {"role": "assistant", "content": None, "tool_calls": [tool_call(_LONDON, prompt_cache_breakpoint=_EXPLICIT)]}, + {"role": "tool", "tool_call_id": call_id(_LONDON), "content": "sunny.", "prompt_cache_breakpoint": _EXPLICIT}, + ) + marks: Final = _marks_sent( + gateway, + "bedrock_converse", + lambda model: {"model": model, "messages": turns, "max_tokens": 64, "drop_params": True}, + ) + assert [mark.label for mark in marks] == [_label(SYSTEM), *_note_labels()] + + +def _function_marked_tool_body(model: str, marker: str) -> dict[str, JsonValue]: + function: Final = TOOL["function"] + assert isinstance(function, dict), TOOL + tool: Final[dict[str, JsonValue]] = { + "type": "function", + "prompt_cache_breakpoint": _EXPLICIT, + "function": {**function, "cache_control": EPHEMERAL}, + } + return {**chat_body(model, _noted_turns(marker), drop_params=True), "tools": [tool]} + + +@pytest.mark.parametrize("provider", ANTHROPIC_BODY_PROVIDERS) +def test_a_tool_breakpoint_beside_a_function_cache_control_leaves_room_for_four_marks_under_drop_params( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + marks: Final = _marks_sent(gateway, provider, lambda model: _function_marked_tool_body(model, marker)) + labels: Final = [mark.label for mark in marks] + assert len(labels) == 4 and labels[1:] == _note_labels(), labels + + +def test_a_tool_breakpoint_beside_a_function_cache_control_keeps_the_latest_breakpoints_on_converse_under_drop_params( + gateway: Gateway, +) -> None: + marker: Final = new_marker() + marks: Final = _marks_sent(gateway, "bedrock_converse", lambda model: _function_marked_tool_body(model, marker)) + labels: Final = [mark.label for mark in marks] + assert len(labels) <= 4 and labels[-3:] == _note_labels(), labels + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_a_function_breakpoint_the_tool_cache_control_overrides_is_never_reported( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + function: Final = TOOL["function"] + assert isinstance(function, dict), TOOL + tool: Final[dict[str, JsonValue]] = { + "type": "function", + "cache_control": EPHEMERAL, + "function": {**function, "prompt_cache_breakpoint": _THIRTY_MINUTES}, + } + marks: Final = _marks_sent(gateway, provider, lambda model: {**chat_body(model, [_final(marker)]), "tools": [tool]}) + assert marks == (Mark("tool:lookup_weather", None),) + + +@pytest.mark.parametrize("provider", PROVIDERS) +def test_assistant_text_read_before_a_one_hour_tool_call_mark_is_sent_as_one_hour( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + turns: Final[list[JsonValue]] = [ + {"role": "user", "content": ASK}, + { + "role": "assistant", + "content": _CHECKING, + "prompt_cache_breakpoint": _EXPLICIT, + "tool_calls": [tool_call(_LONDON, prompt_cache_breakpoint=_ONE_HOUR_BREAKPOINT)], + }, + {"role": "tool", "tool_call_id": call_id(_LONDON), "content": "sunny."}, + _final(marker), + ] + marks: Final = _marks_sent(gateway, provider, lambda model: chat_body(model, turns)) + assert marks == (Mark(_label(_CHECKING), "1h"), Mark(_call_label(_LONDON), "1h")) + + +@pytest.mark.parametrize("value", _MALFORMED, ids=_MALFORMED_IDS) +@pytest.mark.parametrize("shape", ("list_content", "developer")) +def test_a_malformed_message_breakpoint_is_a_400_before_any_provider_call( + gateway: Gateway, shape: Shape, value: JsonValue +) -> None: + marker: Final = new_marker() + _rejected(gateway, "anthropic", lambda model: chat_body(model, _shape_turns(shape, marker, value))) + + +@pytest.mark.parametrize("value", _MALFORMED, ids=_MALFORMED_IDS) +@pytest.mark.parametrize("shape", ("list_content", "developer")) +def test_a_malformed_message_breakpoint_sends_no_mark_under_drop_params( + gateway: Gateway, shape: Shape, value: JsonValue +) -> None: + marker: Final = new_marker() + marks: Final = _marks_sent( + gateway, "anthropic", lambda model: chat_body(model, _shape_turns(shape, marker, value), drop_params=True) + ) + assert marks == () + + +@pytest.mark.parametrize( + ("provider", "result"), + ( + ( + "anthropic", + { + "content": [_text(_CLOUDY), _text(_SUNNY, prompt_cache_breakpoint=_THIRTY_MINUTES)], + "prompt_cache_breakpoint": _EXPLICIT, + }, + ), + ( + "bedrock_converse", + { + "content": [ + _text(_CLOUDY, cache_control=EPHEMERAL), + _text(_SUNNY, prompt_cache_breakpoint=_THIRTY_MINUTES), + ] + }, + ), + ), + ids=("replaced_by_the_results_own_breakpoint", "replaced_by_a_client_marked_part"), +) +def test_a_replaced_part_breakpoint_with_a_ttl_claude_lacks_is_still_a_400( + gateway: Gateway, provider: Provider, result: dict[str, JsonValue] +) -> None: + marker: Final = new_marker() + text: Final = _rejected(gateway, provider, lambda model: chat_body(model, _tool_result_turns(marker, result))) + assert "30m" in text, text + + +@pytest.mark.parametrize( + ("first", "last", "status", "labels"), + ( + ({"mode": "implicit"}, _EXPLICIT, 200, [_label(_LAST_PART)]), + (_EXPLICIT, {"mode": "implicit"}, 400, []), + ), + ids=("last_explicit", "last_implicit"), +) +def test_a_message_breakpoint_sent_twice_takes_the_last_value( + gateway: Gateway, first: JsonValue, last: JsonValue, status: int, labels: list[str] +) -> None: + marker: Final = new_marker() + with wire_server(lambda request: respond("anthropic", request)) as wire, gateway.scenario() as scenario: + model: Final = deployment(gateway, scenario, wire, "anthropic") + turns: Final = _list_content_turns("user", marker, "duplicated", {}) + raw: Final = json.dumps(chat_body(model, turns)).replace( + '"prompt_cache_breakpoint": "duplicated"', + f'"prompt_cache_breakpoint": {json.dumps(first)}, "prompt_cache_breakpoint": {json.dumps(last)}', + ) + response: Final = gateway.client.post( + "/v1/chat/completions", + content=raw, + headers={"Authorization": f"Bearer {gateway.key}", "Content-Type": "application/json"}, + ) + received: Final = wire.drain() + assert response.status_code == status, response.text + assert [[mark.label for mark in sent_marks("anthropic", request)] for request in received] == ( + [labels] if status == 200 else [] + ) + + +@pytest.mark.parametrize("provider", ("anthropic", "bedrock_converse")) +def test_a_null_message_breakpoint_on_list_content_is_no_mark(gateway: Gateway, provider: Provider) -> None: + marker: Final = new_marker() + marks: Final = _marks_sent( + gateway, provider, lambda model: chat_body(model, _list_content_turns("user", marker, None, {})) + ) + assert marks == () + + +@pytest.mark.parametrize("provider", ("anthropic", "bedrock_converse")) +def test_a_client_cache_control_on_the_last_part_wins_over_the_message_breakpoint( + gateway: Gateway, provider: Provider +) -> None: + marker: Final = new_marker() + last: Final[dict[str, JsonValue]] = {"cache_control": _ONE_HOUR} + marks: Final = _marks_sent( + gateway, provider, lambda model: chat_body(model, _list_content_turns("user", marker, _EXPLICIT, last)) + ) + assert marks == (Mark(_label(_LAST_PART), "1h"),) + + +@pytest.mark.parametrize( + ("message_mark", "part_mark"), + ((_EXPLICIT, _ONE_HOUR_BREAKPOINT), (_ONE_HOUR_BREAKPOINT, _EXPLICIT)), + ids=("one_hour_on_the_part", "one_hour_on_the_message"), +) +@pytest.mark.parametrize("provider", ("anthropic", "bedrock_converse")) +def test_a_message_breakpoint_and_its_last_parts_own_breakpoint_join_into_one_mark( + gateway: Gateway, provider: Provider, message_mark: JsonValue, part_mark: JsonValue +) -> None: + marker: Final = new_marker() + last: Final[dict[str, JsonValue]] = {"prompt_cache_breakpoint": part_mark} + marks: Final = _marks_sent( + gateway, provider, lambda model: chat_body(model, _list_content_turns("user", marker, message_mark, last)) + ) + assert marks == (Mark(_label(_LAST_PART), "1h"),) diff --git a/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_wire.py b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_wire.py index 8a439a2e855..68e5e704fe5 100644 --- a/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_wire.py +++ b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_wire.py @@ -4,10 +4,9 @@ from types import MappingProxyType from typing import Final, Literal import pytest -from integration._support.client import Gateway, Scenario, gateway_from_environment +from integration._support.client import Gateway, 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._support.wire import wire_server from integration.providers._cache_control_marks_support import ( ANTHROPIC_MODEL, ASK, @@ -16,7 +15,6 @@ from integration.providers._cache_control_marks_support import ( CITIES, EPHEMERAL, POINTS, - PROVIDER_KEY, SYSTEM, SYSTEM_LABEL, TOOL, @@ -24,8 +22,6 @@ from integration.providers._cache_control_marks_support import ( anthropic_labels, anthropic_marks, anthropic_peer, - bedrock_labels, - bedrock_peer, call_id, chat_body, final_label, @@ -37,23 +33,21 @@ from integration.providers._cache_control_marks_support import ( tool_call, tool_use_label, ) +from integration.providers._claude_prompt_cache_breakpoint_support import ( + ANTHROPIC_BODY_PROVIDERS, + PROVIDERS, + Provider, + deployment, + labels, + litellm_params, + only_request, + respond, +) from pydantic import JsonValue from typing_extensions import assert_never -Provider = Literal[ - "anthropic", "bedrock_converse", "bedrock_invoke", "bedrock_claude_platform", "vertex_ai", "azure_ai" -] Site = Literal["tool", "system", "user_part", "tool_call", "tool_result", "string_message"] -_PROVIDERS: Final[tuple[Provider, ...]] = ( - "anthropic", - "bedrock_converse", - "bedrock_invoke", - "bedrock_claude_platform", - "vertex_ai", - "azure_ai", -) -_ANTHROPIC_BODY_PROVIDERS: Final = tuple(provider for provider in _PROVIDERS if provider != "bedrock_converse") _SITES: Final[tuple[Site, ...]] = ("tool", "system", "user_part", "tool_call", "tool_result", "string_message") _FIVE_SITES: Final[tuple[Site, ...]] = ("tool", "system", "user_part", "tool_call", "string_message") _EXPLICIT: Final[dict[str, JsonValue]] = {"mode": "explicit"} @@ -61,63 +55,6 @@ _BREAKPOINT: Final[Mapping[str, JsonValue]] = MappingProxyType({"prompt_cache_br _NATIVE: Final[Mapping[str, JsonValue]] = MappingProxyType({"cache_control": EPHEMERAL}) _ONE_HOUR: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} _CITY: Final = CITIES[0] -_BEDROCK_TOKEN: Final = "synthetic-bedrock-bearer" -_VERTEX_PROJECT: Final = "scripted-project" -_VERTEX_LOCATION: Final = "us-east5" - - -def _respond(provider: Provider, request: Request) -> Reply: - return bedrock_peer(request) if provider == "bedrock_converse" else anthropic_peer(request) - - -def _labels(provider: Provider, request: Request) -> list[str]: - return bedrock_labels(request) if provider == "bedrock_converse" else anthropic_labels(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 "bedrock_claude_platform": - return { - "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", - "api_base": wire.url, - "api_key": PROVIDER_KEY, - "workspace_id": "wrkspc_scripted", - "aws_region_name": "us-east-1", - } - 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]: @@ -168,57 +105,51 @@ 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) +@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) + 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) + request: Final = only_request(wire) assert b"prompt_cache_breakpoint" not in request.body, request.body - assert _labels(provider, request) == [_expected(provider, site, marker)] + assert labels(provider, request) == [_expected(provider, site, marker)] -@pytest.mark.parametrize("provider", _PROVIDERS) +@pytest.mark.parametrize("provider", PROVIDERS) def test_breakpoint_under_a_tool_function_marks_the_tool(gateway: Gateway, provider: Provider) -> None: marker: Final = new_marker() function: Final = TOOL["function"] assert isinstance(function, dict), TOOL tool: Final[dict[str, JsonValue]] = {**TOOL, "function": {**function, "prompt_cache_breakpoint": _EXPLICIT}} - with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, provider) + 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(())), "tools": [tool]}) assert status == 200, text - request: Final = _only_request(wire) + request: Final = only_request(wire) assert b"prompt_cache_breakpoint" not in request.body, request.body - assert _labels(provider, request) == [_expected(provider, "tool", marker)] + assert labels(provider, request) == [_expected(provider, "tool", marker)] def test_breakpoint_ttl_claude_supports_is_kept(gateway: Gateway) -> None: marker: Final = new_marker() marks: Final = _marked(("user_part",), {"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "1h"}}) with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, "anthropic") + 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) + request: Final = only_request(wire) assert anthropic_marks(request) == (Mark(ASK_LABEL, "1h"),) -@pytest.mark.parametrize("provider", _PROVIDERS) +@pytest.mark.parametrize("provider", PROVIDERS) def test_breakpoint_ttl_claude_lacks_is_a_400_before_any_provider_call(gateway: Gateway, provider: Provider) -> None: marker: Final = new_marker() marks: Final = _marked(("user_part",), {"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}) - with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, provider) + 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, marks)) received: Final = wire.drain() assert status == 400, text @@ -230,10 +161,10 @@ def test_breakpoint_ttl_claude_lacks_keeps_the_mark_under_drop_params(gateway: G marker: Final = new_marker() marks: Final = _marked(("user_part",), {"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}) with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, "anthropic", drop_params=True) + model: Final = deployment(gateway, scenario, wire, "anthropic", drop_params=True) status, _, text = post_chat(gateway, _body(model, marker, marks)) assert status == 200, text - request: Final = _only_request(wire) + request: Final = only_request(wire) assert anthropic_marks(request) == (Mark(ASK_LABEL, None),) @@ -241,7 +172,7 @@ def test_malformed_breakpoint_is_a_400_before_any_provider_call(gateway: Gateway marker: Final = new_marker() marks: Final = _marked(("user_part",), {"prompt_cache_breakpoint": {"mode": "sometimes"}}) with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, "anthropic") + model: Final = deployment(gateway, scenario, wire, "anthropic") status, _, text = post_chat(gateway, _body(model, marker, marks)) received: Final = wire.drain() assert status == 400, text @@ -254,70 +185,70 @@ def test_a_client_cache_control_on_the_same_block_wins_over_the_breakpoint(gatew hour: Final[dict[str, JsonValue]] = {"type": "ephemeral", "ttl": "1h"} marks: Final = _marked(("user_part",), {"cache_control": hour, "prompt_cache_breakpoint": _EXPLICIT}) with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, "anthropic") + 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) + request: Final = only_request(wire) assert anthropic_marks(request) == (Mark(ASK_LABEL, "1h"),) @pytest.mark.parametrize("source", ("deployment", "request")) -@pytest.mark.parametrize("provider", _PROVIDERS) +@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 {} + deployment_fields: 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) + with wire_server(lambda request: respond(provider, request)) as wire, gateway.scenario() as scenario: + model: Final = deployment(gateway, scenario, wire, provider, **deployment_fields) 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) + request: Final = only_request(wire) + assert labels(provider, request) == _latest_four(provider, marker) @pytest.mark.parametrize("mark", (_BREAKPOINT, _NATIVE), ids=("breakpoint", "cache_control")) -@pytest.mark.parametrize("provider", _PROVIDERS) +@pytest.mark.parametrize("provider", PROVIDERS) def test_five_marks_without_drop_params_reach_the_provider_and_its_error_reaches_the_caller( gateway: Gateway, provider: Provider, mark: Mapping[str, JsonValue] ) -> 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) + 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, mark))) received: Final = wire.drain() assert status == 400, text assert "Found 5" in text, text - assert [len(_labels(provider, request)) for request in received] == [5], [request.target for request in received] + assert [len(labels(provider, request)) for request in received] == [5], [request.target for request in received] -@pytest.mark.parametrize("provider", _PROVIDERS) +@pytest.mark.parametrize("provider", PROVIDERS) def test_drop_params_drops_the_breakpoint_past_the_cap_and_keeps_every_mark_the_client_wrote( gateway: Gateway, provider: Provider ) -> None: marker: Final = new_marker() written: Final = _FIVE_SITES[:-1] marks: Final = MappingProxyType({**_marked(written, _NATIVE), _FIVE_SITES[-1]: _BREAKPOINT}) - with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, provider, drop_params=True) + 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, marks)) assert status == 200, text - request: Final = _only_request(wire) - assert _labels(provider, request) == [_expected(provider, site, marker) for site in written] + request: Final = only_request(wire) + assert labels(provider, request) == [_expected(provider, site, marker) for site in written] -@pytest.mark.parametrize("provider", _ANTHROPIC_BODY_PROVIDERS) +@pytest.mark.parametrize("provider", ANTHROPIC_BODY_PROVIDERS) def test_a_breakpoint_ahead_of_a_one_hour_mark_reaches_claude_as_one_hour(gateway: Gateway, provider: Provider) -> None: marker: Final = new_marker() marks: Final = MappingProxyType( {"system": _BREAKPOINT, "user_part": MappingProxyType({"cache_control": _ONE_HOUR})} ) - with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, provider) + 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, marks)) assert status == 200, text - request: Final = _only_request(wire) + request: Final = only_request(wire) assert anthropic_marks(request) == (Mark(SYSTEM_LABEL, "1h"), Mark(ASK_LABEL, "1h")) @@ -327,22 +258,22 @@ def test_a_one_hour_breakpoint_after_a_five_minute_mark_loses_its_ttl_under_drop {"system": _NATIVE, "user_part": MappingProxyType({"prompt_cache_breakpoint": {**_EXPLICIT, "ttl": "1h"}})} ) with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: - model: Final = _deployment(gateway, scenario, wire, "anthropic", drop_params=True) + model: Final = deployment(gateway, scenario, wire, "anthropic", drop_params=True) status, _, text = post_chat(gateway, _body(model, marker, marks)) assert status == 200, text - request: Final = _only_request(wire) + request: Final = only_request(wire) assert anthropic_marks(request) == (Mark(SYSTEM_LABEL, None), Mark(ASK_LABEL, None)) 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( + 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 anthropic_labels(only_request(wire)) == _latest_four("anthropic", marker) assert gateway_injected(response_id) is False @@ -363,16 +294,16 @@ def test_a_model_id_without_claude_is_marked_only_when_its_base_model_is_claude( gateway: Gateway, provider: Provider, model: str, base_model: str | None, marked: bool ) -> None: marker: Final = new_marker() - with wire_server(lambda request: _respond(provider, request)) as wire, gateway.scenario() as scenario: + with wire_server(lambda request: respond(provider, request)) as wire, gateway.scenario() as scenario: deployment: Final = scenario.model( model_info=None if base_model is None else {"base_model": base_model}, - **{**_litellm_params(gateway, wire, provider), "model": model}, + **{**litellm_params(gateway, wire, provider), "model": model}, ) status, _, text = post_chat(gateway, _body(deployment, marker, _marked(("system",)))) assert status == 200, text - request: Final = _only_request(wire) + request: Final = only_request(wire) assert b"prompt_cache_breakpoint" not in request.body, request.body - assert _labels(provider, request) == ([_expected(provider, "system", marker)] if marked else []) + assert labels(provider, request) == ([_expected(provider, "system", marker)] if marked else []) def _drop_params_config(directory: Path) -> Path: @@ -392,10 +323,10 @@ def drop_params_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Ga 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") + 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) + assert anthropic_labels(only_request(wire)) == _latest_four("anthropic", marker) @pytest.mark.timeout(240) @@ -403,8 +334,8 @@ def test_proxy_level_drop_params_wins_over_a_request_that_sets_it_false(drop_par marker: Final = new_marker() marks: Final = _marked(("user_part",), {"prompt_cache_breakpoint": {"mode": "explicit", "ttl": "30m"}}) with wire_server(anthropic_peer) as wire, drop_params_gateway.scenario() as scenario: - model: Final = _deployment(drop_params_gateway, scenario, wire, "anthropic") + model: Final = deployment(drop_params_gateway, scenario, wire, "anthropic") status, _, text = post_chat(drop_params_gateway, _body(model, marker, marks, drop_params=False)) assert status == 200, text - request: Final = _only_request(wire) + request: Final = only_request(wire) assert anthropic_marks(request) == (Mark(ASK_LABEL, None),) diff --git a/tests/integration/routing/test_prompt_caching_affinity_breakpoints.py b/tests/integration/routing/test_prompt_caching_affinity_breakpoints.py index 27e7a2ef0f9..0e81011291f 100644 --- a/tests/integration/routing/test_prompt_caching_affinity_breakpoints.py +++ b/tests/integration/routing/test_prompt_caching_affinity_breakpoints.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from pathlib import Path from typing import Final @@ -5,6 +6,7 @@ import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows, scratch_database from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.redis_process import owned_redis from integration._support.wire import Wire, wire_server from integration.providers._cache_control_marks_support import ( anthropic_deployment, @@ -15,11 +17,24 @@ from integration.providers._cache_control_marks_support import ( owned_config, post_chat, ) +from integration.routing._prompt_cache_pin_support import ( + Pin, + claude_peer, + lives_five_minutes, + lives_one_hour, + one_new_pin, + pin_config, + read_pins, + reference, +) from pydantic import JsonValue +from redis import Redis _MODEL: Final = "affinity-breakpoint-claude" _FOLLOW_UPS: Final = 16 _LONG_SYSTEM: Final = "Answer from the tool results below. " + "lorem ipsum " * 1500 +_REDIS_ENVIRONMENT: Final = ("REDIS_HOST", "REDIS_PORT", "REDIS_PASSWORD", "REDIS_URL", "REDIS_SSL") +_NO_RESPONSE_CACHE: Final[dict[str, JsonValue]] = {"cache": False, "enable_redis_auth_cache": False} def _first_turn(session: str, marker: str) -> list[JsonValue]: @@ -51,7 +66,9 @@ def _body(messages: list[JsonValue]) -> dict[str, JsonValue]: @pytest.mark.timeout(2 * graceful_stop_seconds() + 120) -def test_follow_ups_return_to_the_deployment_that_cached_the_breakpoint_prefix(gateway: Gateway, tmp_path: Path) -> None: +def test_follow_ups_return_to_the_deployment_that_cached_the_breakpoint_prefix( + gateway: Gateway, tmp_path: Path +) -> None: session: Final = new_marker() first_marker: Final = new_marker() follow_up_markers: Final = frozenset(new_marker() for _ in range(_FOLLOW_UPS)) @@ -91,3 +108,50 @@ def test_follow_ups_return_to_the_deployment_that_cached_the_breakpoint_prefix(g assert [status for status, _, _ in follow_ups] == [200] * _FOLLOW_UPS, [text for _, _, text in follow_ups] (pinned,) = tuple(side for side in served if first_marker in side) assert pinned == {first_marker, *follow_up_markers}, served + + +@pytest.mark.parametrize( + ("ttl", "lives"), + ((None, lives_five_minutes), ("1h", lives_one_hour)), + ids=("default_ttl", "one_hour"), +) +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_breakpoint_pin_lives_as_long_as_the_ttl_the_breakpoint_asks_for( + gateway: Gateway, tmp_path: Path, ttl: str | None, lives: Callable[[Pin], bool] +) -> None: + session: Final = new_marker() + marker: Final = new_marker() + identities: Final = (f"breakpoint-pin-left-{session}", f"breakpoint-pin-right-{session}") + breakpoint: Final[dict[str, JsonValue]] = {"mode": "explicit", **({"ttl": ttl} if ttl else {})} + turns: Final[list[JsonValue]] = [ + { + "role": "system", + "content": [{"type": "text", "text": reference(session), "prompt_cache_breakpoint": breakpoint}], + }, + {"role": "user", "content": final_text(marker)}, + ] + with ( + owned_redis(tmp_path) as cache, + wire_server(claude_peer) as left, + wire_server(claude_peer) as right, + Redis(host=cache.host, port=cache.port, socket_timeout=2) as redis, + ): + config: Final = pin_config( + tmp_path, + [ + {**anthropic_deployment(_MODEL, left.url), "model_info": {"id": identities[0]}}, + {**anthropic_deployment(_MODEL, right.url), "model_info": {"id": identities[1]}}, + ], + litellm_settings=_NO_RESPONSE_CACHE, + router_settings={"redis_host": cache.host, "redis_port": cache.port}, + ) + with owned_proxy_process( + gateway, tmp_path, {}, config=config, remove_environment=_REDIS_ENVIRONMENT, workers=2 + ) as owned: + before: Final = frozenset(read_pins(redis)) + status, _, text = post_chat(owned.gateway, _body(turns)) + assert status == 200, text + _, pin = one_new_pin(redis, before) + served: Final = (_served(left), _served(right)) + assert served[identities.index(pin.model_id)] == {marker}, (pin, served) + assert lives(pin), pin