mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
test(proxy): cover failure spend counters across budgets, endpoints and outages
Integration cells for team, user, end-user, tag and window budgets, concurrent streams, billed retries, native /v1/messages, Bedrock guardrail cost, a Redis stall on two workers and disable_spend_updates
This commit is contained in:
yucheng-berri
parent
5485c90ffc
commit
a8a18d600c
3 changed files with 688 additions and 144 deletions
165
tests/integration/spend/_partial_stream.py
Normal file
165
tests/integration/spend/_partial_stream.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterable, Iterator
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from integration._support.client import Gateway, Scenario, eventually, held
|
||||
from integration._support.wire import Reply, Request, Wire
|
||||
from openai.types.chat import ChatCompletionChunk
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
BACKEND: Final = "integration-partial-stream"
|
||||
PIECE: Final = "partial stream billed tokens "
|
||||
DELIVERED_PIECES: Final = 12
|
||||
_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"'
|
||||
ROWS: Final = 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"'
|
||||
|
||||
|
||||
class _SpendRow(BaseModel):
|
||||
spend: float
|
||||
|
||||
|
||||
class StatusRow(BaseModel):
|
||||
status: str
|
||||
spend: float
|
||||
|
||||
|
||||
class _BudgetError(BaseModel):
|
||||
type: str
|
||||
|
||||
|
||||
class _ErrorResponse(BaseModel):
|
||||
error: _BudgetError
|
||||
|
||||
|
||||
def _chunk(stream_id: str, delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes:
|
||||
body: Final = {
|
||||
"id": stream_id,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": BACKEND,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
||||
}
|
||||
return f"data: {json.dumps(body)}\n\n".encode()
|
||||
|
||||
|
||||
def completion() -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": BACKEND,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
).encode()
|
||||
|
||||
|
||||
def peer(delivered: int | None, gate: threading.Event | None = None) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method != "POST":
|
||||
return Reply(status=404)
|
||||
if JSON_OBJECT.validate_json(request.body).get("stream") is True:
|
||||
stream_id: Final = f"chatcmpl-{uuid.uuid4().hex}"
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
*(_chunk(stream_id, {"role": "assistant", "content": PIECE}) for _ in range(DELIVERED_PIECES)),
|
||||
_chunk(stream_id, {}, "stop"),
|
||||
b"data: [DONE]\n\n",
|
||||
),
|
||||
abort_after=delivered,
|
||||
pause_between_chunks=0.02,
|
||||
gate_before_first=gate,
|
||||
)
|
||||
return Reply(body=completion())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def priced_model(scenario: Scenario, wire: Wire) -> str:
|
||||
return scenario.model(
|
||||
model=f"openai/{BACKEND}",
|
||||
api_base=f"{wire.url}/v1",
|
||||
api_key="synthetic-openai-key",
|
||||
input_cost_per_token=0.001,
|
||||
output_cost_per_token=0.002,
|
||||
)
|
||||
|
||||
|
||||
def stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 5,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": "partial stream"}],
|
||||
**fields,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def broken_stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> None:
|
||||
response: Final = stream(gateway, model, key, **fields)
|
||||
assert response.status_code == 200, response.text
|
||||
events: Final = tuple(
|
||||
line.removeprefix("data: ")
|
||||
for line in response.text.splitlines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
assert "error" in JSON_OBJECT.validate_json(events[-1]), response.text
|
||||
delivered: Final = "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1]))
|
||||
assert delivered == PIECE * DELIVERED_PIECES, response.text
|
||||
|
||||
|
||||
def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]:
|
||||
for chunk in chunks:
|
||||
yield from (choice.delta.content or "" for choice in chunk.choices)
|
||||
|
||||
|
||||
def chat(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 5,
|
||||
"messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}],
|
||||
**fields,
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]:
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70)
|
||||
found: Final = held(
|
||||
lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2
|
||||
)
|
||||
return tuple(_SpendRow.model_validate(row).spend for row in found)
|
||||
|
||||
|
||||
def chats_sent(wire: Wire) -> int:
|
||||
return sum(1 for request in wire.drain() if request.method == "POST")
|
||||
|
||||
|
||||
def recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float:
|
||||
key: Final = scenario.key(models=[model], max_budget=100)
|
||||
broken_stream(gateway, model, key)
|
||||
(spend,) = failure_spend(gateway, key)
|
||||
assert spend > 0
|
||||
return spend
|
||||
|
||||
|
||||
def assert_budget_refusal(response: httpx.Response) -> None:
|
||||
assert response.status_code == 422, response.text
|
||||
assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text
|
||||
|
|
@ -0,0 +1,224 @@
|
|||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, eventually
|
||||
from integration._support.process import owned_proxy_process
|
||||
from integration._support.redis_process import owned_redis
|
||||
from integration._support.timing import assert_not_delayed_by, timed
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.spend._partial_stream import (
|
||||
DELIVERED_PIECES,
|
||||
assert_budget_refusal,
|
||||
broken_stream,
|
||||
chat,
|
||||
chats_sent,
|
||||
failure_spend,
|
||||
peer,
|
||||
priced_model,
|
||||
recovered_cost,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.constants import REDIS_SOCKET_TIMEOUT
|
||||
|
||||
_BURST: Final = 12
|
||||
_SPREAD: Final = 6
|
||||
_ATTACK_MARKER: Final = "synthetic-attack-marker"
|
||||
_POLICY_ID: Final = "synthetic-priced-policy"
|
||||
_POLICY_APPLY_PATH: Final = f"/guardrail/{_POLICY_ID}/version/DRAFT/apply"
|
||||
|
||||
|
||||
def _proxy_config(directory: Path, name: str, sections: dict[str, JsonValue]) -> Path:
|
||||
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
path: Final = directory / f"{name}.yaml"
|
||||
path.write_text(yaml.safe_dump({**base, **sections}))
|
||||
return path
|
||||
|
||||
|
||||
def _redis_counter_config(directory: Path) -> Path:
|
||||
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
return _proxy_config(
|
||||
directory,
|
||||
"redis-counters",
|
||||
{
|
||||
"router_settings": {
|
||||
**base["router_settings"],
|
||||
"redis_host": "os.environ/REDIS_HOST",
|
||||
"redis_port": "os.environ/REDIS_PORT",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _spread_requests(gateway: Gateway, model: str, key: str) -> tuple[httpx.Response, ...]:
|
||||
return tuple(
|
||||
gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"spread {uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
headers={"connection": "close"},
|
||||
)
|
||||
for _ in range(_SPREAD)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(600)
|
||||
def test_broken_streams_through_a_redis_stall_log_once_and_the_budget_holds_on_every_worker(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
release: Final = threading.Event()
|
||||
with (
|
||||
owned_redis(tmp_path) as counters,
|
||||
wire_server(peer(DELIVERED_PIECES)) as wire,
|
||||
wire_server(peer(DELIVERED_PIECES, gate=release)) as held_wire,
|
||||
owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{"REDIS_HOST": counters.host, "REDIS_PORT": str(counters.port)},
|
||||
config=_redis_counter_config(tmp_path),
|
||||
workers=2,
|
||||
) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
scenario.cleanups.callback(release.set)
|
||||
candidate: Final = owned.gateway
|
||||
model: Final = priced_model(scenario, wire)
|
||||
stalled_model: Final = priced_model(scenario, held_wire)
|
||||
cost: Final = recovered_cost(candidate, scenario, model)
|
||||
control: Final = timed(lambda: broken_stream(candidate, model, scenario.key(models=[model])))
|
||||
burst_key: Final = scenario.key(models=[stalled_model])
|
||||
chats_sent(held_wire)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=_BURST) as pool:
|
||||
pending: Final = tuple(
|
||||
pool.submit(broken_stream, candidate, stalled_model, burst_key) for _ in range(_BURST)
|
||||
)
|
||||
eventually(held_wire.received.qsize, lambda held: held >= _BURST, seconds=30)
|
||||
with counters.paused():
|
||||
release.set()
|
||||
drained: Final = timed(lambda: tuple(future.result() for future in pending))
|
||||
|
||||
assert_not_delayed_by((drained.seconds,), control=(control.seconds,), injected_delay=REDIS_SOCKET_TIMEOUT)
|
||||
assert failure_spend(candidate, burst_key, rows=_BURST) == pytest.approx((cost,) * _BURST)
|
||||
|
||||
budgeted: Final = scenario.key(models=[model], max_budget=cost / 2)
|
||||
broken_stream(candidate, model, budgeted)
|
||||
before_kill: Final = _spread_requests(candidate, model, budgeted)
|
||||
owned.kill_worker(min(owned.worker_pids()))
|
||||
after_kill: Final = _spread_requests(candidate, model, budgeted)
|
||||
|
||||
for refused in (*before_kill, *after_kill):
|
||||
assert_budget_refusal(refused)
|
||||
|
||||
|
||||
def _policy(request: Request) -> Reply:
|
||||
if request.method != "POST" or request.target != _POLICY_APPLY_PATH:
|
||||
return Reply(status=404)
|
||||
if _ATTACK_MARKER.encode() not in request.body:
|
||||
return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode())
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
"outputs": [{"text": "blocked by synthetic policy"}],
|
||||
"assessments": [
|
||||
{"contentPolicy": {"filters": [{"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}]}}
|
||||
],
|
||||
"usage": {"contentPolicyUnits": 1},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _policy_config(directory: Path, policy: Wire) -> Path:
|
||||
return _proxy_config(
|
||||
directory,
|
||||
"priced-policy",
|
||||
{
|
||||
"guardrails": [
|
||||
{
|
||||
"guardrail_name": "priced-bedrock-policy",
|
||||
"litellm_params": {
|
||||
"guardrail": "bedrock",
|
||||
"mode": "pre_call",
|
||||
"default_on": True,
|
||||
"guardrailIdentifier": _POLICY_ID,
|
||||
"guardrailVersion": "DRAFT",
|
||||
"aws_region_name": "us-east-1",
|
||||
"aws_access_key_id": "AKIASYNTHETICGUARDRAIL",
|
||||
"aws_secret_access_key": "synthetic-guardrail-secret",
|
||||
"aws_bedrock_runtime_endpoint": policy.url,
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _attack(gateway: Gateway, model: str, key: str) -> httpx.Response:
|
||||
blocked: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"{_ATTACK_MARKER} {uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
)
|
||||
assert blocked.status_code == 400, blocked.text
|
||||
return blocked
|
||||
|
||||
|
||||
@pytest.mark.timeout(480)
|
||||
def test_guardrail_cost_on_blocked_requests_counts_toward_the_budget(gateway: Gateway, tmp_path: Path) -> None:
|
||||
with (
|
||||
wire_server(_policy) as policy,
|
||||
wire_server(peer(None)) as wire,
|
||||
owned_proxy_process(gateway, tmp_path, {}, config=_policy_config(tmp_path, policy), workers=2) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
candidate: Final = owned.gateway
|
||||
model: Final = priced_model(scenario, wire)
|
||||
meter: Final = scenario.key(models=[model], max_budget=100)
|
||||
_attack(candidate, model, meter)
|
||||
(cost,) = failure_spend(candidate, meter)
|
||||
assert cost > 0
|
||||
key: Final = scenario.key(models=[model], max_budget=cost * 1.5)
|
||||
chats_sent(wire)
|
||||
|
||||
_attack(candidate, model, key)
|
||||
_attack(candidate, model, key)
|
||||
after: Final = chat(candidate, model, key)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
assert chats_sent(wire) == 0
|
||||
|
||||
|
||||
@pytest.mark.timeout(480)
|
||||
def test_broken_streams_leave_the_budget_alone_while_spend_updates_are_disabled(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
config: Final = _proxy_config(
|
||||
tmp_path, "spend-updates-off", {"general_settings": {**base["general_settings"], "disable_spend_updates": True}}
|
||||
)
|
||||
with (
|
||||
wire_server(peer(DELIVERED_PIECES)) as wire,
|
||||
gateway.scenario() as metering,
|
||||
owned_proxy_process(gateway, tmp_path, {}, config=config) as owned,
|
||||
owned.gateway.scenario() as scenario,
|
||||
):
|
||||
cost: Final = recovered_cost(gateway, metering, priced_model(metering, wire))
|
||||
candidate: Final = owned.gateway
|
||||
model: Final = priced_model(scenario, wire)
|
||||
key: Final = scenario.key(models=[model], max_budget=cost / 2)
|
||||
|
||||
broken_stream(candidate, model, key)
|
||||
after: Final = chat(candidate, model, key)
|
||||
|
||||
assert after.status_code == 200, after.text
|
||||
|
|
@ -1,195 +1,350 @@
|
|||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterable, Iterator
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually, held
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from openai.types.chat import ChatCompletionChunk
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
from integration._support.client import Gateway, eventually, held
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from integration.spend._partial_stream import (
|
||||
BACKEND,
|
||||
DELIVERED_PIECES,
|
||||
JSON_OBJECT,
|
||||
PIECE,
|
||||
ROWS,
|
||||
StatusRow,
|
||||
assert_budget_refusal,
|
||||
broken_stream,
|
||||
chat,
|
||||
chats_sent,
|
||||
completion,
|
||||
failure_spend,
|
||||
peer,
|
||||
priced_model,
|
||||
recovered_cost,
|
||||
stream,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
|
||||
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
|
||||
_BACKEND: Final = "integration-partial-stream"
|
||||
_PIECE: Final = "partial stream billed tokens "
|
||||
_DELIVERED_PIECES: Final = 12
|
||||
_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"'
|
||||
_CONCURRENT_STREAMS: Final = 6
|
||||
_MESSAGES_BACKEND: Final = "claude-integration-partial-stream"
|
||||
_UNCACHED: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}}
|
||||
|
||||
|
||||
class _SpendRow(BaseModel):
|
||||
spend: float
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None:
|
||||
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
cost: Final = recovered_cost(gateway, scenario, model)
|
||||
key: Final = scenario.key(models=[model], max_budget=cost / 2)
|
||||
chats_sent(wire)
|
||||
|
||||
broken_stream(gateway, model, key)
|
||||
after: Final = chat(gateway, model, key)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
assert failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0))
|
||||
assert chats_sent(wire) == 1
|
||||
|
||||
|
||||
class _BudgetError(BaseModel):
|
||||
type: str
|
||||
@pytest.mark.timeout(240)
|
||||
def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None:
|
||||
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
cost: Final = recovered_cost(gateway, scenario, model)
|
||||
key: Final = scenario.key(models=[model], max_budget=cost * 1.5)
|
||||
|
||||
broken_stream(gateway, model, key)
|
||||
broken_stream(gateway, model, key)
|
||||
after: Final = chat(gateway, model, key)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
assert failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0))
|
||||
|
||||
|
||||
class _ErrorResponse(BaseModel):
|
||||
error: _BudgetError
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None:
|
||||
with wire_server(peer(0)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
key: Final = scenario.key(models=[model], max_budget=0.05)
|
||||
|
||||
refused: Final = stream(gateway, model, key)
|
||||
after: Final = chat(gateway, model, key)
|
||||
|
||||
assert refused.status_code == 500, refused.text
|
||||
|
||||
assert after.status_code == 200, after.text
|
||||
assert failure_spend(gateway, key) == (0.0,)
|
||||
|
||||
|
||||
def _chunk(delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes:
|
||||
body: Final = {
|
||||
"id": "chatcmpl-partial",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1700000000,
|
||||
"model": _BACKEND,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
||||
}
|
||||
return f"data: {json.dumps(body)}\n\n".encode()
|
||||
def _broken_stream_then_refusal(
|
||||
gateway: Gateway, model: str, spender: str, other: str, **fields: JsonValue
|
||||
) -> httpx.Response:
|
||||
broken_stream(gateway, model, spender, **fields)
|
||||
return chat(gateway, model, other, **fields)
|
||||
|
||||
|
||||
def _completion() -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": _BACKEND,
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
).encode()
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_broken_stream_on_one_team_key_blocks_the_teams_other_keys(gateway: Gateway) -> None:
|
||||
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
cost: Final = recovered_cost(gateway, scenario, model)
|
||||
team: Final = scenario.team(models=[model], max_budget=cost / 2)
|
||||
|
||||
after: Final = _broken_stream_then_refusal(
|
||||
gateway, model, scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model])
|
||||
)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
|
||||
|
||||
def _peer(delivered: int) -> Callable[[Request], Reply]:
|
||||
frames: Final = (
|
||||
*(_chunk({"role": "assistant", "content": _PIECE}) for _ in range(_DELIVERED_PIECES)),
|
||||
_chunk({}, "stop"),
|
||||
b"data: [DONE]\n\n",
|
||||
)
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_broken_stream_on_one_of_a_users_keys_blocks_their_other_keys(gateway: Gateway) -> None:
|
||||
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
cost: Final = recovered_cost(gateway, scenario, model)
|
||||
user: Final = scenario.user(max_budget=cost / 2)
|
||||
|
||||
after: Final = _broken_stream_then_refusal(
|
||||
gateway, model, scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model])
|
||||
)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_broken_stream_for_an_end_user_blocks_their_next_request_on_any_key(gateway: Gateway) -> None:
|
||||
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
||||
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
cost: Final = recovered_cost(gateway, scenario, model)
|
||||
gateway.post("/customer/new", {"user_id": end_user, "max_budget": cost / 2})
|
||||
scenario.cleanups.callback(gateway.post, "/customer/delete", {"user_ids": [end_user]})
|
||||
|
||||
after: Final = _broken_stream_then_refusal(
|
||||
gateway, model, scenario.key(models=[model]), scenario.key(models=[model]), user=end_user
|
||||
)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_broken_stream_under_a_tag_blocks_the_next_request_with_that_tag(gateway: Gateway) -> None:
|
||||
tag: Final = f"integration-partial-stream-{uuid.uuid4().hex}"
|
||||
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
cost: Final = recovered_cost(gateway, scenario, model)
|
||||
gateway.post("/tag/new", {"name": tag, "max_budget": cost / 2})
|
||||
scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": tag})
|
||||
key: Final = scenario.key(models=[model])
|
||||
|
||||
after: Final = _broken_stream_then_refusal(gateway, model, key, key, metadata={"tags": [tag]})
|
||||
untagged: Final = chat(gateway, model, key)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
assert untagged.status_code == 200, untagged.text
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_broken_stream_fills_the_keys_budget_window(gateway: Gateway) -> None:
|
||||
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
cost: Final = recovered_cost(gateway, scenario, model)
|
||||
key: Final = scenario.key(models=[model], budget_limits=[{"budget_duration": "1d", "max_budget": cost / 2}])
|
||||
|
||||
after: Final = _broken_stream_then_refusal(gateway, model, key, key)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_concurrent_broken_streams_each_count_once_toward_the_budget(gateway: Gateway) -> None:
|
||||
with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
cost: Final = recovered_cost(gateway, scenario, model)
|
||||
key: Final = scenario.key(models=[model], max_budget=cost * (_CONCURRENT_STREAMS + 0.5))
|
||||
|
||||
with ThreadPoolExecutor(max_workers=_CONCURRENT_STREAMS) as pool:
|
||||
tuple(pool.map(lambda _: broken_stream(gateway, model, key), range(_CONCURRENT_STREAMS)))
|
||||
within_budget: Final = chat(gateway, model, key)
|
||||
broken_stream(gateway, model, key)
|
||||
over_budget: Final = chat(gateway, model, key)
|
||||
|
||||
assert within_budget.status_code == 200, within_budget.text
|
||||
assert_budget_refusal(over_budget)
|
||||
assert sorted(failure_spend(gateway, key, rows=_CONCURRENT_STREAMS + 2)) == pytest.approx(
|
||||
[0.0] + [cost] * (_CONCURRENT_STREAMS + 1)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_completed_streams_count_once_toward_the_budget(gateway: Gateway) -> None:
|
||||
with wire_server(peer(None)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = priced_model(scenario, wire)
|
||||
meter: Final = scenario.key(models=[model], max_budget=100)
|
||||
assert stream(gateway, model, meter, **_UNCACHED).status_code == 200
|
||||
(metered,) = eventually(
|
||||
lambda: gateway.rows(ROWS, (sha256(meter.encode()).hexdigest(),)),
|
||||
lambda found: len(found) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
cost: Final = StatusRow.model_validate(metered).spend
|
||||
key: Final = scenario.key(models=[model], max_budget=cost * 1.5)
|
||||
|
||||
first: Final = stream(gateway, model, key, **_UNCACHED)
|
||||
second: Final = stream(gateway, model, key, **_UNCACHED)
|
||||
assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text)
|
||||
assert _rows(gateway, key, 2) == (StatusRow(status="success", spend=cost),) * 2
|
||||
third: Final = stream(gateway, model, key, **_UNCACHED)
|
||||
|
||||
assert_budget_refusal(third)
|
||||
|
||||
|
||||
def _hung_then_served(release: threading.Event, hung_model: str) -> Callable[[Request], Reply]:
|
||||
def respond(request: Request) -> Reply:
|
||||
if request.method != "POST":
|
||||
return Reply(status=404)
|
||||
if _JSON_OBJECT.validate_json(request.body).get("stream") is True:
|
||||
return Reply(
|
||||
content_type="text/event-stream", chunks=frames, abort_after=delivered, pause_between_chunks=0.02
|
||||
)
|
||||
return Reply(body=_completion())
|
||||
if JSON_OBJECT.validate_json(request.body).get("model") == hung_model:
|
||||
return Reply(chunks=(completion(),), gate_before_first=release)
|
||||
return Reply(body=completion())
|
||||
|
||||
return respond
|
||||
|
||||
|
||||
def _priced_model(scenario: Scenario, wire: Wire) -> str:
|
||||
return scenario.model(
|
||||
model=f"openai/{_BACKEND}",
|
||||
api_base=f"{wire.url}/v1",
|
||||
api_key="synthetic-openai-key",
|
||||
input_cost_per_token=0.001,
|
||||
output_cost_per_token=0.002,
|
||||
def _timed_out_spend(gateway: Gateway, model: str, key: str) -> float:
|
||||
timed_out: Final = chat(gateway, model, key, messages=[{"role": "user", "content": "hung"}])
|
||||
assert timed_out.status_code in (408, 500), timed_out.text
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
landed: Final = eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= 3, seconds=70)
|
||||
rows: Final = tuple(
|
||||
StatusRow.model_validate(row)
|
||||
for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == len(landed), holding=2)
|
||||
)
|
||||
assert {row.status for row in rows} == {"failure"}, rows
|
||||
return sum(row.spend for row in rows)
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_retries_the_provider_billed_before_a_timeout_count_toward_the_budget(gateway: Gateway) -> None:
|
||||
hung_backend: Final = f"{BACKEND}-hung"
|
||||
release: Final = threading.Event()
|
||||
with wire_server(_hung_then_served(release, hung_backend)) as wire, gateway.scenario() as scenario:
|
||||
scenario.cleanups.callback(release.set)
|
||||
hung: Final = scenario.model(
|
||||
model=f"openai/{hung_backend}",
|
||||
api_base=f"{wire.url}/v1",
|
||||
api_key="synthetic-openai-key",
|
||||
input_cost_per_token=0.001,
|
||||
output_cost_per_token=0.002,
|
||||
timeout=1,
|
||||
num_retries=2,
|
||||
)
|
||||
served: Final = priced_model(scenario, wire)
|
||||
billed: Final = _timed_out_spend(gateway, hung, scenario.key(models=[hung, served], max_budget=100))
|
||||
assert billed > 0
|
||||
key: Final = scenario.key(models=[hung, served], max_budget=billed / 2)
|
||||
|
||||
assert _timed_out_spend(gateway, hung, key) == pytest.approx(billed)
|
||||
after: Final = chat(gateway, served, key)
|
||||
|
||||
assert_budget_refusal(after)
|
||||
|
||||
|
||||
def _messages_event(kind: str, fields: dict[str, JsonValue]) -> bytes:
|
||||
return f"event: {kind}\ndata: {json.dumps({'type': kind, **fields})}\n\n".encode()
|
||||
|
||||
|
||||
def _messages_peer(request: Request) -> Reply:
|
||||
if request.method != "POST" or request.target != "/v1/messages":
|
||||
return Reply(status=404)
|
||||
message: Final[dict[str, JsonValue]] = {
|
||||
"id": f"msg_{uuid.uuid4().hex}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": _MESSAGES_BACKEND,
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 10, "output_tokens": 1},
|
||||
}
|
||||
if JSON_OBJECT.validate_json(request.body).get("stream") is not True:
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{**message, "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn"}
|
||||
).encode()
|
||||
)
|
||||
return Reply(
|
||||
content_type="text/event-stream",
|
||||
chunks=(
|
||||
_messages_event("message_start", {"message": {**message, "content": []}}),
|
||||
_messages_event("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
*(
|
||||
_messages_event("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": PIECE}})
|
||||
for _ in range(DELIVERED_PIECES)
|
||||
),
|
||||
_messages_event("content_block_stop", {"index": 0}),
|
||||
_messages_event(
|
||||
"message_delta",
|
||||
{"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 60}},
|
||||
),
|
||||
_messages_event("message_stop", {}),
|
||||
),
|
||||
abort_after=DELIVERED_PIECES + 2,
|
||||
pause_between_chunks=0.02,
|
||||
)
|
||||
|
||||
|
||||
def _stream(gateway: Gateway, model: str, key: str) -> httpx.Response:
|
||||
def _messages(gateway: Gateway, model: str, key: str, *, streamed: bool) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
"/v1/messages",
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 5,
|
||||
"stream": True,
|
||||
"messages": [{"role": "user", "content": "partial stream"}],
|
||||
"stream": streamed,
|
||||
"messages": [{"role": "user", "content": f"partial stream {uuid.uuid4().hex}"}],
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _broken_stream(gateway: Gateway, model: str, key: str) -> str:
|
||||
response: Final = _stream(gateway, model, key)
|
||||
def _broken_messages_stream(gateway: Gateway, model: str, key: str) -> None:
|
||||
response: Final = _messages(gateway, model, key, streamed=True)
|
||||
assert response.status_code == 200, response.text
|
||||
events: Final = tuple(
|
||||
line.removeprefix("data: ")
|
||||
for line in response.text.splitlines()
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
assert "error" in _JSON_OBJECT.validate_json(events[-1]), response.text
|
||||
return "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1]))
|
||||
assert response.text.count(PIECE) == DELIVERED_PIECES, response.text
|
||||
assert "message_stop" not in response.text, response.text
|
||||
|
||||
|
||||
def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]:
|
||||
for chunk in chunks:
|
||||
yield from (choice.delta.content or "" for choice in chunk.choices)
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]:
|
||||
def _rows(gateway: Gateway, key: str, count: int) -> tuple[StatusRow, ...]:
|
||||
digest: Final = sha256(key.encode()).hexdigest()
|
||||
eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70)
|
||||
found: Final = held(
|
||||
lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2
|
||||
eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= count, seconds=70)
|
||||
return tuple(
|
||||
StatusRow.model_validate(row)
|
||||
for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == count, holding=2)
|
||||
)
|
||||
return tuple(_SpendRow.model_validate(row).spend for row in found)
|
||||
|
||||
|
||||
def _chats_sent(wire: Wire) -> int:
|
||||
return sum(1 for request in wire.drain() if request.method == "POST")
|
||||
|
||||
|
||||
def _recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float:
|
||||
key: Final = scenario.key(models=[model], max_budget=100)
|
||||
assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES
|
||||
(spend,) = _failure_spend(gateway, key)
|
||||
assert spend > 0
|
||||
return spend
|
||||
|
||||
|
||||
def _assert_budget_refusal(response: httpx.Response) -> None:
|
||||
assert response.status_code == 422, response.text
|
||||
assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(_DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario, wire)
|
||||
cost: Final = _recovered_cost(gateway, scenario, model)
|
||||
key: Final = scenario.key(models=[model], max_budget=cost / 2)
|
||||
_chats_sent(wire)
|
||||
def test_a_broken_messages_stream_counts_toward_the_budget(gateway: Gateway) -> None:
|
||||
with wire_server(_messages_peer) as wire, gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model=f"anthropic/{_MESSAGES_BACKEND}",
|
||||
api_base=wire.url,
|
||||
api_key="synthetic-anthropic-key",
|
||||
input_cost_per_token=0.001,
|
||||
output_cost_per_token=0.002,
|
||||
)
|
||||
meter: Final = scenario.key(models=[model], max_budget=100)
|
||||
_broken_messages_stream(gateway, model, meter)
|
||||
(metered,) = _rows(gateway, meter, 1)
|
||||
assert metered.spend > 0, metered
|
||||
key: Final = scenario.key(models=[model], max_budget=metered.spend / 2)
|
||||
|
||||
assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES
|
||||
after: Final = _chat(gateway, model, key)
|
||||
_broken_messages_stream(gateway, model, key)
|
||||
after: Final = _messages(gateway, model, key, streamed=False)
|
||||
|
||||
_assert_budget_refusal(after)
|
||||
assert _failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0))
|
||||
assert _chats_sent(wire) == 1
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(_DELIVERED_PIECES)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario, wire)
|
||||
cost: Final = _recovered_cost(gateway, scenario, model)
|
||||
key: Final = scenario.key(models=[model], max_budget=cost * 1.5)
|
||||
|
||||
assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES
|
||||
assert _broken_stream(gateway, model, key) == _PIECE * _DELIVERED_PIECES
|
||||
after: Final = _chat(gateway, model, key)
|
||||
|
||||
_assert_budget_refusal(after)
|
||||
assert _failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0))
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None:
|
||||
with wire_server(_peer(0)) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _priced_model(scenario, wire)
|
||||
key: Final = scenario.key(models=[model], max_budget=0.05)
|
||||
|
||||
refused: Final = _stream(gateway, model, key)
|
||||
after: Final = _chat(gateway, model, key)
|
||||
|
||||
assert refused.status_code == 500, refused.text
|
||||
|
||||
assert after.status_code == 200, after.text
|
||||
assert _failure_spend(gateway, key) == (0.0,)
|
||||
assert_budget_refusal(after)
|
||||
assert _rows(gateway, key, 2) == (metered, StatusRow(status="failure", spend=0.0))
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue