diff --git a/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_chaos.py b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_chaos.py new file mode 100644 index 00000000000..e5d7208f1a7 --- /dev/null +++ b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_chaos.py @@ -0,0 +1,222 @@ +import asyncio +import socket +import threading +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from typing import Final, Literal + +import httpx +import pytest +from integration._support import bedrock_runtime_peer +from integration._support.client import Gateway, Scenario, eventually +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.prompt_cache_breakpoint import open_upstream_connections, started_worker_pids +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import ( + ASK, + ASK_LABEL, + BEDROCK_MODEL, + PROVIDER_KEY, + SYSTEM, + bedrock_labels, + final_text, + new_marker, + owned_config, +) +from pydantic import JsonValue + +Kind = Literal["chat", "chat_stream", "messages", "messages_stream"] + +_KINDS: Final[tuple[Kind, ...]] = ("chat", "chat_stream", "messages", "messages_stream") +_EXPLICIT: Final[dict[str, JsonValue]] = {"mode": "explicit"} +_CONVERSE_SYSTEM: Final = f"system:text:{SYSTEM}" +_OWNED_MODEL: Final = "breakpoint-converse" +_HOLD_SECONDS: Final = 60 + + +@dataclass(frozen=True, slots=True) +class _Call: + kind: Kind + marker: str + + +@dataclass(frozen=True, slots=True) +class _Answer: + call: _Call + status: int + text: str + + +@dataclass(frozen=True, slots=True) +class _Held: + release: threading.Event + arrived: SimpleQueue[str] + + def respond(self, request: Request) -> Reply: + self.arrived.put(bedrock_runtime_peer.marker_of(request)) + self.release.wait(timeout=_HOLD_SECONDS) + return bedrock_runtime_peer.respond(request) + + +def _calls(count: int) -> tuple[_Call, ...]: + return tuple(_Call(_KINDS[index % len(_KINDS)], new_marker()) for index in range(count)) + + +def _converse_params(url: str) -> dict[str, JsonValue]: + return { + "model": f"bedrock/converse/{BEDROCK_MODEL}", + "api_key": PROVIDER_KEY, + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": url, + } + + +def _converse(scenario: Scenario, url: str) -> str: + return scenario.model(**_converse_params(url)) + + +def _request(call: _Call, model: str) -> tuple[str, dict[str, JsonValue]]: + stream: Final = call.kind.endswith("_stream") + if call.kind.startswith("chat"): + return "/v1/chat/completions", { + "model": model, + "max_tokens": 64, + "stream": stream, + "num_retries": 0, + "messages": [ + {"role": "system", "content": [{"type": "text", "text": SYSTEM, "prompt_cache_breakpoint": _EXPLICIT}]}, + {"role": "user", "content": final_text(call.marker)}, + ], + } + return "/v1/messages", { + "model": model, + "max_tokens": 64, + "stream": stream, + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": ASK, "prompt_cache_breakpoint": _EXPLICIT}, + {"type": "text", "text": final_text(call.marker)}, + ], + } + ], + } + + +def _expected_label(call: _Call) -> str: + return _CONVERSE_SYSTEM if call.kind.startswith("chat") else ASK_LABEL + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Answer: + path, body = _request(call, model) + async with client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {key}"}) as response: + raw: Final = (await response.aread()).decode() + return _Answer(call, response.status_code, raw) + + +async def _burst( + base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Answer, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=90, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Answer)) + + +def _posts(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if request.method == "POST") + + +def _assert_answered(answers: tuple[_Answer, ...]) -> None: + for answer in answers: + assert answer.status == 200, answer + assert f"marker-{answer.call.marker}" in answer.text, answer + + +def _assert_marked(received: tuple[Request, ...], calls: tuple[_Call, ...]) -> None: + by_marker: Final = {bedrock_runtime_peer.marker_of(request): request for request in received} + assert len(received) == len(by_marker) == len(calls), [request.target for request in received] + assert set(by_marker) == {call.marker for call in calls}, by_marker.keys() + for call in calls: + assert bedrock_labels(by_marker[call.marker]) == [_expected_label(call)], call + + +def _reserved_port() -> int: + with socket.socket() as reserve: + reserve.bind(("127.0.0.1", 0)) + return reserve.getsockname()[1] + + +async def test_a_mixed_burst_marks_every_request_once(gateway: Gateway) -> None: + calls: Final = _calls(32) + with wire_server(bedrock_runtime_peer.respond) as wire, gateway.scenario() as scenario: + model: Final = _converse(scenario, wire.url) + answers: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + received: Final = _posts(wire) + assert len(answers) == len(calls) + _assert_answered(answers) + _assert_marked(received, calls) + + +async def test_an_upstream_outage_mid_burst_reaches_the_caller_and_recovered_requests_are_marked( + gateway: Gateway, +) -> None: + port: Final = _reserved_port() + url: Final = str(gateway.client.base_url) + before_calls, during_calls, after_calls = _calls(8), _calls(8), _calls(8) + with gateway.scenario() as scenario: + model: Final = _converse(scenario, f"http://127.0.0.1:{port}") + with wire_server(bedrock_runtime_peer.respond, port=port) as first: + before: Final = await _burst(url, gateway.key, model, before_calls) + first_received: Final = _posts(first) + during: Final = await _burst(url, gateway.key, model, during_calls) + liveliness: Final = gateway.request("GET", "/health/liveliness") + with wire_server(bedrock_runtime_peer.respond, port=port) as second: + after: Final = await _burst(url, gateway.key, model, after_calls) + second_received: Final = _posts(second) + _assert_answered(before) + _assert_marked(first_received, before_calls) + assert len(during) == len(during_calls), during + assert all(answer.status >= 500 and "error" in answer.text for answer in during), during + assert liveliness.status_code == 200, liveliness.text + _assert_answered(after) + _assert_marked(second_received, after_calls) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +async def test_a_worker_killed_mid_burst_leaves_the_survivor_answering_with_marks( + gateway: Gateway, tmp_path: Path +) -> None: + held: Final = _Held(threading.Event(), SimpleQueue()) + calls: Final = _calls(12) + follow_up_calls: Final = _calls(8) + with wire_server(held.respond) as wire: + config: Final = owned_config( + tmp_path, [{"model_name": _OWNED_MODEL, "litellm_params": _converse_params(wire.url)}] + ) + with owned_proxy_process(gateway, tmp_path, {}, config=config, workers=2) as owned: + url: Final = str(owned.gateway.client.base_url) + workers: Final = eventually(lambda: started_worker_pids(owned.log), lambda pids: len(pids) == 2, seconds=60) + burst: Final = asyncio.create_task( + _burst(url, owned.gateway.key, _OWNED_MODEL, calls, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, held.arrived.qsize, lambda size: size == len(calls), 60) + held_by: Final = {pid: open_upstream_connections(pid, wire.url) for pid in workers} + assert sum(held_by.values()) == len(calls), held_by + victim, survivor = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + held.release.set() + served: Final = await burst + follow_up: Final = await _burst(url, owned.gateway.key, _OWNED_MODEL, follow_up_calls) + received: Final = _posts(wire) + assert len(served) == held_by[survivor], (held_by, served) + _assert_answered(served) + assert len(follow_up) == len(follow_up_calls), follow_up + _assert_answered(follow_up) + _assert_marked(received, (*calls, *follow_up_calls)) diff --git a/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_routes.py b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_routes.py new file mode 100644 index 00000000000..378f6f5f748 --- /dev/null +++ b/tests/integration/providers/test_claude_chat_prompt_cache_breakpoint_routes.py @@ -0,0 +1,434 @@ +import json +from collections.abc import Iterable +from typing import Final, cast + +import anthropic +import openai +import pytest +from anthropic.types import MessageParam +from integration._support import bedrock_runtime_peer +from integration._support.client import Gateway, Scenario, eventually +from integration._support.database import read_rows +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, + PROVIDER_KEY, + SYSTEM, + SYSTEM_LABEL, + Mark, + anthropic_labels, + anthropic_marks, + anthropic_peer, + bedrock_labels, + call_id, + chat_body, + final_text, + gateway_injected, + marker_of, + new_marker, + post_chat, + tool_call, + tool_use_label, +) +from openai.types.chat import ChatCompletionMessageParam +from openai.types.responses import ResponseInputParam +from pydantic import JsonValue, TypeAdapter + +_EXPLICIT: Final[dict[str, JsonValue]] = {"mode": "explicit"} +_THIRTY_MINUTES: Final[dict[str, JsonValue]] = {"mode": "explicit", "ttl": "30m"} +_CONVERSE_SYSTEM: Final = f"system:text:{SYSTEM}" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MALFORMED: Final[tuple[JsonValue, ...]] = ( + 5, + [], + "", + "x" * 5000, + {}, + {"mode": "explicit", "ttl": 5}, + {"mode": "explicit", "ttl": ""}, + {"mode": ["explicit"]}, +) +_MALFORMED_IDS: Final = ("int", "list", "empty_string", "5kb_string", "empty_object", "ttl_int", "ttl_empty", "mode_list") + + +def _anthropic(scenario: Scenario, wire: Wire, **fields: JsonValue) -> str: + return scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=PROVIDER_KEY, model_info=None, **fields + ) + + +def _converse(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"bedrock/converse/{BEDROCK_MODEL}", + api_key=PROVIDER_KEY, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + ) + + +def _turns(marker: str, **system_fields: JsonValue) -> list[JsonValue]: + return [ + {"role": "system", "content": [{"type": "text", "text": SYSTEM, **system_fields}]}, + {"role": "user", "content": final_text(marker)}, + ] + + +def _messages_turns(marker: str, breakpoint: JsonValue) -> list[JsonValue]: + return [ + { + "role": "user", + "content": [ + {"type": "text", "text": ASK, "prompt_cache_breakpoint": breakpoint}, + {"type": "text", "text": final_text(marker)}, + ], + } + ] + + +def _posts(wire: Wire) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if request.method == "POST") + + +def _only_request(wire: Wire) -> Request: + received: Final = _posts(wire) + assert len(received) == 1, [request.target for request in received] + return received[0] + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI(base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0) + + +def _anthropic_client(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + + +def _chat_messages(turns: list[JsonValue]) -> Iterable[ChatCompletionMessageParam]: + return cast(Iterable[ChatCompletionMessageParam], turns) + + +def _sdk_messages(turns: list[JsonValue]) -> Iterable[MessageParam]: + return cast(Iterable[MessageParam], turns) + + +def _messages_text(client: anthropic.Anthropic, model: str, turns: list[JsonValue], *, stream: bool) -> str: + if not stream: + message: Final = client.messages.create(model=model, max_tokens=64, messages=_sdk_messages(turns)) + return "".join(block.text for block in message.content if block.type == "text") + with client.messages.stream(model=model, max_tokens=64, messages=_sdk_messages(turns)) as events: + return "".join(events.text_stream) + + +def _spend_row_lands(response_id: str) -> None: + eventually( + lambda: read_rows('SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (response_id,)), + lambda rows: len(rows) == 1, + seconds=70, + ) + + +async def test_async_openai_sdk_stream_to_claude_carries_the_system_breakpoint(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + async with openai.AsyncOpenAI( + base_url=f"{gateway.client.base_url}/v1", api_key=gateway.key, max_retries=0 + ) as client: + stream: Final = await client.chat.completions.create( + model=model, + messages=_chat_messages(_turns(marker, prompt_cache_breakpoint=_EXPLICIT)), + max_tokens=64, + stream=True, + ) + chunks: Final = [chunk async for chunk in stream] + request: Final = _only_request(wire) + assert len({chunk.id for chunk in chunks}) == 1, chunks + assert marker_of(request) == marker + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "sunny" + assert _JSON_OBJECT.validate_json(request.body)["stream"] is True + assert b"prompt_cache_breakpoint" not in request.body, request.body + assert anthropic_marks(request) == (Mark(SYSTEM_LABEL, None),) + + +def test_sync_openai_sdk_chat_to_converse_carries_the_system_breakpoint_as_a_cache_point(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(bedrock_runtime_peer.respond) as wire, gateway.scenario() as scenario: + model: Final = _converse(scenario, wire) + with _openai_client(gateway) as client: + completion: Final = client.chat.completions.create( + model=model, messages=_chat_messages(_turns(marker, prompt_cache_breakpoint=_EXPLICIT)), max_tokens=64 + ) + request: Final = _only_request(wire) + assert completion.choices[0].message.content == bedrock_runtime_peer.answer(marker) + assert request.target.endswith("/converse"), request.target + assert b"prompt_cache_breakpoint" not in request.body, request.body + assert bedrock_labels(request) == [_CONVERSE_SYSTEM] + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +def test_anthropic_sdk_messages_to_converse_carry_the_breakpoint_as_a_cache_point( + gateway: Gateway, stream: bool +) -> None: + marker: Final = new_marker() + with wire_server(bedrock_runtime_peer.respond) as wire, gateway.scenario() as scenario: + model: Final = _converse(scenario, wire) + with _anthropic_client(gateway) as client: + text: Final = _messages_text(client, model, _messages_turns(marker, _EXPLICIT), stream=stream) + request: Final = _only_request(wire) + assert text == bedrock_runtime_peer.answer(marker) + assert request.target.endswith("/converse-stream" if stream else "/converse"), request.target + assert b"prompt_cache_breakpoint" not in request.body, request.body + assert bedrock_labels(request) == [ASK_LABEL] + + +def test_anthropic_sdk_messages_to_converse_with_a_ttl_claude_lacks_is_a_400_before_any_provider_call( + gateway: Gateway, +) -> None: + marker: Final = new_marker() + with wire_server(bedrock_runtime_peer.respond) as wire, gateway.scenario() as scenario: + model: Final = _converse(scenario, wire) + with _anthropic_client(gateway) as client, pytest.raises(anthropic.BadRequestError) as rejected: + _messages_text(client, model, _messages_turns(marker, _THIRTY_MINUTES), stream=False) + received: Final = _posts(wire) + assert rejected.value.status_code == 400 + assert "30m" in str(rejected.value) and "drop_params" in str(rejected.value), rejected.value + assert received == (), [request.target for request in received] + + +def test_messages_to_converse_with_a_ttl_claude_lacks_caches_under_request_drop_params(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(bedrock_runtime_peer.respond) as wire, gateway.scenario() as scenario: + model: Final = _converse(scenario, wire) + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": _messages_turns(marker, _THIRTY_MINUTES), + "drop_params": True, + }, + ) + request: Final = _only_request(wire) + assert response.status_code == 200, response.text + assert bedrock_runtime_peer.answer(marker) in response.text, response.text + assert bedrock_labels(request) == [ASK_LABEL] + + +def test_messages_to_a_native_claude_deployment_forward_the_content_as_sent(gateway: Gateway) -> None: + marker: Final = new_marker() + turns: Final = _messages_turns(marker, _EXPLICIT) + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + with _anthropic_client(gateway) as client: + message: Final = client.messages.create(model=model, max_tokens=64, messages=_sdk_messages(turns)) + request: Final = _only_request(wire) + assert message.id == f"msg_{marker}" + assert _JSON_OBJECT.validate_json(request.body)["messages"] == turns + + +def test_responses_bridge_to_claude_answers_a_breakpoint_with_a_ttl_claude_lacks(gateway: Gateway) -> None: + marker: Final = new_marker() + items: Final[list[JsonValue]] = [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": ASK, "prompt_cache_breakpoint": _THIRTY_MINUTES}, + {"type": "input_text", "text": final_text(marker)}, + ], + } + ] + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + with _openai_client(gateway) as client: + response: Final = client.responses.create( + model=model, input=cast(ResponseInputParam, items), max_output_tokens=64 + ) + request: Final = _only_request(wire) + assert response.output_text == "sunny" + assert f"marker-{marker}" in request.body.decode() + assert anthropic_labels(request) == [] + + +def _unauthorized(_: Request) -> Reply: + return Reply( + status=401, + body=json.dumps( + {"error": {"message": "Incorrect API key provided", "type": "invalid_request_error", "code": "invalid_api_key"}} + ).encode(), + ) + + +def test_a_fallback_to_claude_carries_the_breakpoint(gateway: Gateway) -> None: + marker: Final = new_marker() + with ( + wire_server(_unauthorized) as primary_wire, + wire_server(anthropic_peer) as claude_wire, + gateway.scenario() as scenario, + ): + primary: Final = scenario.model(api_base=f"{primary_wire.url}/v1") + claude: Final = _anthropic(scenario, claude_wire) + status, _, text = post_chat( + gateway, + chat_body(primary, _turns(marker, prompt_cache_breakpoint=_EXPLICIT), fallbacks=[claude], num_retries=0), + ) + primary_posts: Final = _posts(primary_wire) + request: Final = _only_request(claude_wire) + assert status == 200, text + assert marker_of(request) == marker + assert len(primary_posts) == 1, [post.target for post in primary_posts] + assert anthropic_marks(request) == (Mark(SYSTEM_LABEL, None),) + + +def test_a_fallback_to_claude_with_a_ttl_claude_lacks_never_reaches_claude(gateway: Gateway) -> None: + marker: Final = new_marker() + with ( + wire_server(_unauthorized) as primary_wire, + wire_server(anthropic_peer) as claude_wire, + gateway.scenario() as scenario, + ): + primary: Final = scenario.model(api_base=f"{primary_wire.url}/v1") + claude: Final = _anthropic(scenario, claude_wire) + status, _, text = post_chat( + gateway, + chat_body( + primary, _turns(marker, prompt_cache_breakpoint=_THIRTY_MINUTES), fallbacks=[claude], num_retries=0 + ), + ) + primary_posts: Final = _posts(primary_wire) + claude_posts: Final = _posts(claude_wire) + assert status == 401, text + assert "UnsupportedParamsError" in text and "30m" in text, text + assert len(primary_posts) == 1, [post.target for post in primary_posts] + assert claude_posts == (), [post.target for post in claude_posts] + + +def test_auto_caching_stands_down_for_one_breakpoint_on_a_tool_call(gateway: Gateway) -> None: + marker: Final = new_marker() + city: Final = CITIES[0] + turns: Final[list[JsonValue]] = [ + {"role": "system", "content": SYSTEM}, + {"role": "user", "content": ASK}, + {"role": "assistant", "content": "", "tool_calls": [tool_call(city, prompt_cache_breakpoint=_EXPLICIT)]}, + {"role": "tool", "tool_call_id": call_id(city), "content": "sunny."}, + {"role": "user", "content": final_text(marker)}, + ] + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + key: Final = scenario.key(metadata={"enable_prompt_caching": True}) + status, response_id, text = post_chat(gateway, chat_body(model, turns), key=key) + request: Final = _only_request(wire) + assert status == 200, text + assert marker_of(request) == marker + assert anthropic_labels(request) == [tool_use_label(city)] + assert gateway_injected(response_id) is False + + +def test_a_repeated_breakpoint_request_is_answered_from_the_response_cache(gateway: Gateway) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + body: Final = chat_body(model, _turns(marker, prompt_cache_breakpoint=_EXPLICIT)) + first: Final = post_chat(gateway, body) + _spend_row_lands(first[1]) + second: Final = post_chat(gateway, body) + request: Final = _only_request(wire) + assert first[0] == 200 and second[0] == 200, (first, second) + assert first[1] == second[1], (first, second) + assert marker_of(request) == marker + assert anthropic_marks(request) == (Mark(SYSTEM_LABEL, None),) + + +@pytest.mark.parametrize("value", _MALFORMED, ids=_MALFORMED_IDS) +def test_a_malformed_breakpoint_is_a_400_and_the_deployment_keeps_serving(gateway: Gateway, value: JsonValue) -> None: + marker: Final = new_marker() + plain: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + status, _, text = post_chat(gateway, chat_body(model, _turns(marker, prompt_cache_breakpoint=value))) + rejected: Final = _posts(wire) + plain_status, _, plain_text = post_chat(gateway, chat_body(model, _turns(plain))) + served: Final = _only_request(wire) + assert status == 400, text + assert "prompt_cache_breakpoint" in text and "drop_params" in text, text + assert rejected == (), [request.target for request in rejected] + assert plain_status == 200, plain_text + assert marker_of(served) == plain + assert anthropic_labels(served) == [] + + +@pytest.mark.parametrize( + ("first", "last", "status", "labels"), + ( + ({"mode": "implicit"}, _EXPLICIT, 200, [SYSTEM_LABEL]), + (_EXPLICIT, {"mode": "implicit"}, 400, []), + ), + ids=("last_explicit", "last_implicit"), +) +def test_a_breakpoint_sent_twice_in_one_block_takes_the_last_value( + gateway: Gateway, first: JsonValue, last: JsonValue, status: int, labels: list[str] +) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + raw: Final = json.dumps(chat_body(model, _turns(marker, duplicated=0))).replace( + '"duplicated": 0', + 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 = _posts(wire) + assert response.status_code == status, response.text + assert [anthropic_labels(request) for request in received] == ([labels] if status == 200 else []) + + +@pytest.mark.parametrize("upstream_status", (500, 401)) +def test_an_upstream_error_on_a_breakpoint_request_reaches_the_caller_after_one_attempt( + gateway: Gateway, upstream_status: int +) -> None: + marker: Final = new_marker() + failure: Final = Reply( + status=upstream_status, + body=json.dumps( + {"type": "error", "error": {"type": "api_error", "message": f"scripted {upstream_status} {marker}"}} + ).encode(), + ) + with wire_server(lambda _: failure) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + status, _, text = post_chat( + gateway, chat_body(model, _turns(marker, prompt_cache_breakpoint=_EXPLICIT), num_retries=0) + ) + request: Final = _only_request(wire) + assert status == upstream_status, text + assert f"scripted {upstream_status} {marker}" in text, text + assert anthropic_marks(request) == (Mark(SYSTEM_LABEL, None),) + + +@pytest.mark.parametrize( + ("value", "marks"), + ( + (None, ()), + ({"mode": "explicit", "ttl": None}, (Mark(SYSTEM_LABEL, None),)), + ({"mode": "explicit"}, (Mark(SYSTEM_LABEL, None),)), + ), + ids=("null", "ttl_null", "ttl_missing"), +) +def test_a_null_breakpoint_is_no_mark_and_a_null_or_missing_ttl_is_the_default_ttl( + gateway: Gateway, value: JsonValue, marks: tuple[Mark, ...] +) -> None: + marker: Final = new_marker() + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = _anthropic(scenario, wire) + status, _, text = post_chat(gateway, chat_body(model, _turns(marker, prompt_cache_breakpoint=value))) + request: Final = _only_request(wire) + assert status == 200, text + assert marker_of(request) == marker + assert b"prompt_cache_breakpoint" not in request.body, request.body + assert anthropic_marks(request) == marks 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 47ea2c11a14..8a439a2e855 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 @@ -396,3 +396,15 @@ def test_proxy_level_drop_params_keeps_the_latest_four_breakpoints(drop_params_g 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) + + +@pytest.mark.timeout(240) +def test_proxy_level_drop_params_wins_over_a_request_that_sets_it_false(drop_params_gateway: Gateway) -> None: + 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") + status, _, text = post_chat(drop_params_gateway, _body(model, marker, marks, drop_params=False)) + assert status == 200, text + 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 new file mode 100644 index 00000000000..27e7a2ef0f9 --- /dev/null +++ b/tests/integration/routing/test_prompt_caching_affinity_breakpoints.py @@ -0,0 +1,93 @@ +from pathlib import Path +from typing import Final + +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.wire import Wire, wire_server +from integration.providers._cache_control_marks_support import ( + anthropic_deployment, + anthropic_peer, + final_text, + marker_of, + new_marker, + owned_config, + post_chat, +) +from pydantic import JsonValue + +_MODEL: Final = "affinity-breakpoint-claude" +_FOLLOW_UPS: Final = 16 +_LONG_SYSTEM: Final = "Answer from the tool results below. " + "lorem ipsum " * 1500 + + +def _first_turn(session: str, marker: str) -> list[JsonValue]: + return [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": f"{_LONG_SYSTEM} session {session}", + "prompt_cache_breakpoint": {"mode": "explicit"}, + } + ], + }, + {"role": "user", "content": final_text(marker)}, + ] + + +def _follow_up(first_turn: list[JsonValue], marker: str) -> list[JsonValue]: + return [*first_turn, {"role": "assistant", "content": "sunny"}, {"role": "user", "content": final_text(marker)}] + + +def _served(wire: Wire) -> frozenset[str]: + return frozenset(marker_of(request) for request in wire.drain()) + + +def _body(messages: list[JsonValue]) -> dict[str, JsonValue]: + return {"model": _MODEL, "messages": messages, "max_tokens": 64} + + +@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: + session: Final = new_marker() + first_marker: Final = new_marker() + follow_up_markers: Final = frozenset(new_marker() for _ in range(_FOLLOW_UPS)) + first_turn: Final = _first_turn(session, first_marker) + with scratch_database() as database_url, wire_server(anthropic_peer) as left, wire_server(anthropic_peer) as right: + config: Final = owned_config( + tmp_path, + [ + {**anthropic_deployment(_MODEL, left.url), "model_info": {"id": f"breakpoint-left-{session}"}}, + {**anthropic_deployment(_MODEL, right.url), "model_info": {"id": f"breakpoint-right-{session}"}}, + ], + router_settings={"optional_pre_call_checks": ["prompt_caching"]}, + ) + with owned_proxy_process( + gateway, + tmp_path, + {"DATABASE_URL": database_url}, + config=config, + remove_environment=("DATABASE_URL_READ_REPLICA",), + workers=2, + ) as owned: + status, response_id, text = post_chat(owned.gateway, _body(first_turn)) + assert status == 200, text + eventually( + lambda: read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (response_id,), + database_url=database_url, + ), + lambda rows: len(rows) == 1, + seconds=70, + ) + follow_ups: Final = tuple( + post_chat(owned.gateway, _body(_follow_up(first_turn, marker))) for marker in follow_up_markers + ) + served: Final = (_served(left), _served(right)) + 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