mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
test(integration): wire cells for breakpoint placement, Mantle and Claude Platform, and the router's 1h pin
This commit is contained in:
mateo-berri
parent
51af6c112e
commit
09fc34b91d
4 changed files with 795 additions and 133 deletions
|
|
@ -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))
|
||||
|
|
@ -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"),)
|
||||
|
|
@ -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),)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue