mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
fix(exceptions): map upstream 402 to PaymentRequiredError and cool down 402 deployments (#44879)
* fix(exceptions): keep upstream 402 status and cool down 402 deployments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(exceptions): map 402 to PaymentRequiredError subclass of BadRequestError Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(exceptions): annotate PaymentRequiredError methods Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): skip 402 cooldown on single-deployment model groups Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(exceptions): single prefix and 402 fallback response for PaymentRequiredError Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(anthropic): map billing_error to PaymentRequiredError regardless of status Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(router): honor explicit allowed-fails policy for single-deployment 402s Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * revert(anthropic): drop billing_error body mapping to PaymentRequiredError Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(exceptions): type the PaymentRequiredError constructor parameters * test(integration): cover 402 PaymentRequiredError mapping and cooldown --------- Co-authored-by: Mubashir Osmani <mubashir@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
d1cfe17518
commit
e98bbc2f8e
10 changed files with 1176 additions and 11 deletions
|
|
@ -1404,6 +1404,7 @@ from .exceptions import (
|
|||
BadGatewayError,
|
||||
OpenAIError,
|
||||
ContextWindowExceededError,
|
||||
PaymentRequiredError as PaymentRequiredError,
|
||||
ContentPolicyViolationError,
|
||||
BudgetExceededError,
|
||||
APIError,
|
||||
|
|
|
|||
|
|
@ -573,6 +573,54 @@ class ContextWindowExceededError(BadRequestError):
|
|||
return _message
|
||||
|
||||
|
||||
class PaymentRequiredError(BadRequestError):
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
model: str,
|
||||
llm_provider: str,
|
||||
response: httpx.Response | None = None,
|
||||
litellm_debug_info: str | None = None,
|
||||
) -> None:
|
||||
response_is_valid: Final = (
|
||||
response is not None
|
||||
and isinstance(response, httpx.Response)
|
||||
and hasattr(response, "_request")
|
||||
and getattr(response, "_request", None) is not None
|
||||
)
|
||||
response_for_parent: Final = (
|
||||
response
|
||||
if response_is_valid
|
||||
else httpx.Response(
|
||||
status_code=402,
|
||||
request=httpx.Request(method="GET", url="https://litellm.ai"),
|
||||
)
|
||||
)
|
||||
super().__init__(
|
||||
message=message,
|
||||
model=model,
|
||||
llm_provider=llm_provider,
|
||||
response=response_for_parent,
|
||||
litellm_debug_info=litellm_debug_info,
|
||||
)
|
||||
self.status_code = 402
|
||||
self.message = f"litellm.PaymentRequiredError: {message}"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
self.message
|
||||
+ (f" LiteLLM Retried: {self.num_retries} times" if self.num_retries else "")
|
||||
+ (f", LiteLLM Max Retries: {self.max_retries}" if self.max_retries else "")
|
||||
)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return (
|
||||
self.message
|
||||
+ (f" LiteLLM Retried: {self.num_retries} times" if self.num_retries else "")
|
||||
+ (f", LiteLLM Max Retries: {self.max_retries}" if self.max_retries else "")
|
||||
)
|
||||
|
||||
|
||||
# sub class of bad request error - meant to help us catch guardrails-related errors on proxy.
|
||||
class RejectedRequestError(BadRequestError):
|
||||
def __init__(
|
||||
|
|
@ -977,6 +1025,7 @@ LITELLM_EXCEPTION_TYPES: Final = [
|
|||
PermissionDeniedError,
|
||||
RateLimitError,
|
||||
ContextWindowExceededError,
|
||||
PaymentRequiredError,
|
||||
RejectedRequestError,
|
||||
ContentPolicyViolationError,
|
||||
InternalServerError,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ from ..exceptions import (
|
|||
ContextWindowExceededError,
|
||||
InternalServerError,
|
||||
NotFoundError,
|
||||
PaymentRequiredError,
|
||||
PermissionDeniedError,
|
||||
RateLimitError,
|
||||
ServiceUnavailableError,
|
||||
|
|
@ -2327,6 +2328,14 @@ def _map_exception_by_status(
|
|||
litellm_debug_info=extra_information,
|
||||
exception_status_code=status_code,
|
||||
)
|
||||
case 402:
|
||||
raise PaymentRequiredError(
|
||||
message=message,
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case _ if status_code < 500:
|
||||
raise BadRequestError(
|
||||
message=message,
|
||||
|
|
|
|||
|
|
@ -250,11 +250,7 @@ def _is_cooldown_required(
|
|||
# Cool down 429 Rate Limit Errors
|
||||
return True
|
||||
|
||||
elif exception_status == 401:
|
||||
# Cool down 401 Auth Errors
|
||||
return True
|
||||
|
||||
elif exception_status == 408 or exception_status == 404:
|
||||
elif exception_status in (401, 402, 408, 404):
|
||||
return True
|
||||
|
||||
else:
|
||||
|
|
@ -362,8 +358,17 @@ def _should_cooldown_deployment(
|
|||
or litellm_router_instance.team_model_has_alternatives(deployment)
|
||||
)
|
||||
|
||||
## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level)
|
||||
exception_status_int: Final = cast_exception_status_to_int(exception_status)
|
||||
dep_policy, dep_allowed_fails = _get_deployment_cooldown_policy(litellm_router_instance, deployment)
|
||||
if (
|
||||
is_single_deployment_model_group
|
||||
and exception_status_int == 402
|
||||
and _resolve_allowed_fails_from_policy(dep_policy, original_exception) is None
|
||||
and litellm_router_instance.get_allowed_fails_from_policy(original_exception) is None
|
||||
):
|
||||
return False
|
||||
|
||||
## CHECK DEPLOYMENT-LEVEL POLICY FIRST (overrides router-level)
|
||||
if dep_policy is not None or dep_allowed_fails is not None:
|
||||
return _should_cooldown_based_on_deployment_policy(
|
||||
litellm_router_instance,
|
||||
|
|
@ -398,7 +403,6 @@ def _should_cooldown_deployment(
|
|||
num_fails_this_minute,
|
||||
)
|
||||
|
||||
exception_status_int: Final = cast_exception_status_to_int(exception_status)
|
||||
if exception_status_int == 429 and not is_single_deployment_model_group:
|
||||
return True
|
||||
elif percent_fails == 1.0 and total_requests_this_minute >= SINGLE_DEPLOYMENT_TRAFFIC_FAILURE_THRESHOLD:
|
||||
|
|
@ -412,7 +416,7 @@ def _should_cooldown_deployment(
|
|||
# Only apply error rate cooldown when we have enough requests to make the percentage meaningful
|
||||
return True
|
||||
|
||||
elif litellm.should_retry(status_code=cast_exception_status_to_int(exception_status)) is False:
|
||||
elif litellm.should_retry(status_code=exception_status_int) is False:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
|
|
|||
564
tests/integration/routing/test_payment_required_cooldown.py
Normal file
564
tests/integration/routing/test_payment_required_cooldown.py
Normal file
|
|
@ -0,0 +1,564 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import signal
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import ExitStack
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import psutil
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import GATEWAY_LIMITS, Gateway, Scenario, eventually, gateway_from_environment
|
||||
from integration._support.openai_wire import chat_reply, responses_reply
|
||||
from integration._support.process import graceful_stop_seconds, owned_proxy_process
|
||||
from integration._support.prometheus_series import PROXY_FAILURES, Sample, scrape
|
||||
from integration._support.redis_process import OwnedRedis, owned_redis
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from integration.routing.test_payment_required_mapping import (
|
||||
OPENAI_MODEL,
|
||||
PAYMENT_REQUIRED,
|
||||
anthropic_client,
|
||||
anthropic_error,
|
||||
anthropic_params,
|
||||
assert_payment_required_row,
|
||||
chat_body,
|
||||
deployment,
|
||||
is_model_info_probe,
|
||||
model_list_reply,
|
||||
new_group,
|
||||
new_marker,
|
||||
openai_client,
|
||||
openai_params,
|
||||
openai_refusal,
|
||||
post,
|
||||
provider_calls,
|
||||
)
|
||||
from pydantic import JsonValue
|
||||
from redis import Redis
|
||||
|
||||
pytestmark: Final = pytest.mark.timeout(300)
|
||||
|
||||
COOLDOWN_SECONDS: Final = 300
|
||||
RELOAD_SECONDS: Final = 2
|
||||
NO_DEPLOYMENTS: Final = "No deployments available"
|
||||
COOLED_DOWN: Final = "litellm_deployment_cooled_down_total"
|
||||
PRIMARY_GROUP: Final = new_group()
|
||||
FALLBACK_GROUP: Final = new_group()
|
||||
MARKER: Final = re.compile(rb"pr402-[0-9a-f]{32}(?:-\d+)?")
|
||||
WORKER_PID: Final = re.compile(r"Started server process \[(\d+)\]")
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Outcome:
|
||||
marker: str
|
||||
status: int
|
||||
detail: str
|
||||
call_id: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Call:
|
||||
marker: str
|
||||
path: str
|
||||
body: Mapping[str, JsonValue]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CooldownRig:
|
||||
gateway: Gateway
|
||||
cache: OwnedRedis
|
||||
log: Path
|
||||
|
||||
def cooled(self, deployment_id: str) -> bool:
|
||||
return cooled(self.cache, deployment_id)
|
||||
|
||||
|
||||
def cooled(cache: OwnedRedis, deployment_id: str) -> bool:
|
||||
with Redis(host=cache.host, port=cache.port) as client:
|
||||
return int(client.exists(f"deployment:{deployment_id}:cooldown")) == 1
|
||||
|
||||
|
||||
def marker_in(request: Request) -> str:
|
||||
found: Final = MARKER.search(request.body)
|
||||
assert found is not None, (request.method, request.target, request.headers, request.body[:300])
|
||||
return found.group(0).decode()
|
||||
|
||||
|
||||
def markers_received(wire: Wire) -> tuple[str, ...]:
|
||||
return tuple(marker_in(request) for request in provider_calls(wire))
|
||||
|
||||
|
||||
def refusing_reply(request: Request) -> Reply:
|
||||
return anthropic_error(402, f"scripted 402 {marker_in(request)}")
|
||||
|
||||
|
||||
def healthy_reply(request: Request) -> Reply:
|
||||
if is_model_info_probe(request):
|
||||
return model_list_reply()
|
||||
marker: Final = marker_in(request)
|
||||
body: Final = json.loads(request.body)
|
||||
stream: Final = isinstance(body, dict) and body.get("stream") is True
|
||||
if request.target.endswith("/responses"):
|
||||
return responses_reply(f"resp_{marker}", OPENAI_MODEL, f"served {marker}", stream=stream)
|
||||
return chat_reply(f"chatcmpl-{marker}", OPENAI_MODEL, f"served {marker}", stream=stream)
|
||||
|
||||
|
||||
def router_settings(cache: OwnedRedis, **extra: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"num_retries": 0,
|
||||
"cooldown_time": COOLDOWN_SECONDS,
|
||||
"redis_host": cache.host,
|
||||
"redis_port": cache.port,
|
||||
**extra,
|
||||
}
|
||||
|
||||
|
||||
def cooldown_config(
|
||||
directory: Path,
|
||||
settings: Mapping[str, JsonValue],
|
||||
*,
|
||||
model_list: Sequence[Mapping[str, JsonValue]] = (),
|
||||
prometheus: bool = True,
|
||||
) -> Path:
|
||||
base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
assert isinstance(base, dict)
|
||||
litellm_settings: Final = (
|
||||
{**base["litellm_settings"], "callbacks": ["prometheus"]} if prometheus else base["litellm_settings"]
|
||||
)
|
||||
config: Final = {
|
||||
**base,
|
||||
"litellm_settings": litellm_settings,
|
||||
"router_settings": dict(settings),
|
||||
**({"model_list": [dict(entry) for entry in model_list]} if model_list else {}),
|
||||
}
|
||||
path: Final = directory / f"cooldown-{uuid.uuid4().hex[:8]}.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[CooldownRig]:
|
||||
directory: Final = tmp_path_factory.mktemp("payment-required-cooldown")
|
||||
with ExitStack() as stack:
|
||||
upstream: Final = stack.enter_context(gateway_from_environment())
|
||||
cache: Final = stack.enter_context(owned_redis(directory))
|
||||
settings: Final = router_settings(cache, fallbacks=[{PRIMARY_GROUP: [FALLBACK_GROUP]}])
|
||||
owned: Final = stack.enter_context(
|
||||
owned_proxy_process(
|
||||
upstream,
|
||||
directory,
|
||||
{"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": str(RELOAD_SECONDS)},
|
||||
config=cooldown_config(directory, settings),
|
||||
remove_environment=("PROMETHEUS_MULTIPROC_DIR",),
|
||||
)
|
||||
)
|
||||
yield CooldownRig(owned.gateway, cache, owned.log)
|
||||
|
||||
|
||||
def two_deployments(scenario: Scenario, group: str, refusing: Wire, healthy: Wire) -> str:
|
||||
refused_id: Final = deployment(scenario, group, anthropic_params(refusing))
|
||||
deployment(scenario, group, openai_params(healthy))
|
||||
return refused_id
|
||||
|
||||
|
||||
def raw_chat(gateway: Gateway, group: str, marker: str, **extra: JsonValue) -> Outcome:
|
||||
response: Final = post(gateway, "/v1/chat/completions", chat_body(group, marker, **extra))
|
||||
return Outcome(marker, response.status_code, response.text, response.headers["x-litellm-call-id"])
|
||||
|
||||
|
||||
def raw_call(gateway: Gateway, call: Call) -> Outcome:
|
||||
response: Final = post(gateway, call.path, call.body)
|
||||
return Outcome(call.marker, response.status_code, response.text, response.headers["x-litellm-call-id"])
|
||||
|
||||
|
||||
def openai_chat(client: openai.OpenAI, group: str, marker: str) -> Outcome:
|
||||
try:
|
||||
raw: Final = client.chat.completions.with_raw_response.create(
|
||||
model=group, messages=[{"role": "user", "content": marker}]
|
||||
)
|
||||
except openai.APIStatusError as error:
|
||||
return Outcome(marker, error.status_code, error.message, error.response.headers["x-litellm-call-id"])
|
||||
return Outcome(marker, raw.status_code, raw.parse().id, raw.headers["x-litellm-call-id"])
|
||||
|
||||
|
||||
def anthropic_message(client: anthropic.Anthropic, group: str, marker: str) -> Outcome:
|
||||
try:
|
||||
raw: Final = client.messages.with_raw_response.create(
|
||||
model=group, max_tokens=32, messages=[{"role": "user", "content": marker}]
|
||||
)
|
||||
except anthropic.APIStatusError as error:
|
||||
return Outcome(marker, error.status_code, error.message, error.response.headers["x-litellm-call-id"])
|
||||
block: Final = raw.parse().content[0]
|
||||
assert isinstance(block, anthropic.types.TextBlock), block
|
||||
return Outcome(marker, raw.status_code, block.text, raw.headers["x-litellm-call-id"])
|
||||
|
||||
|
||||
def until_refused(attempt: Callable[[str], Outcome], tries: int = 40) -> tuple[Outcome, ...]:
|
||||
def walk() -> Iterator[Outcome]:
|
||||
for _ in range(tries):
|
||||
outcome: Final = attempt(new_marker())
|
||||
yield outcome
|
||||
if outcome.status != 200:
|
||||
return
|
||||
|
||||
history: Final = tuple(walk())
|
||||
assert history[-1].status != 200, "the refusing deployment was never picked"
|
||||
return history
|
||||
|
||||
|
||||
def assert_refused(outcome: Outcome) -> None:
|
||||
assert outcome.status == 402, (outcome.status, outcome.detail)
|
||||
assert PAYMENT_REQUIRED in outcome.detail, outcome.detail
|
||||
|
||||
|
||||
def assert_served(outcomes: Sequence[Outcome], identity: Callable[[str], str]) -> None:
|
||||
assert all(outcome.status == 200 for outcome in outcomes), [(o.status, o.detail[:200]) for o in outcomes]
|
||||
assert all(identity(outcome.marker) in outcome.detail for outcome in outcomes), [o.detail[:200] for o in outcomes]
|
||||
|
||||
|
||||
def chat_identity(marker: str) -> str:
|
||||
return f"chatcmpl-{marker}"
|
||||
|
||||
|
||||
def served_text(marker: str) -> str:
|
||||
return f"served {marker}"
|
||||
|
||||
|
||||
def response_identity(marker: str) -> str:
|
||||
return f"resp_{marker}"
|
||||
|
||||
|
||||
def markers(outcomes: Sequence[Outcome]) -> tuple[str, ...]:
|
||||
return tuple(outcome.marker for outcome in outcomes)
|
||||
|
||||
|
||||
def failure_samples(gateway: Gateway, group: str) -> tuple[Sample, ...]:
|
||||
return tuple(
|
||||
sample
|
||||
for sample in scrape(gateway)
|
||||
if sample.name == PROXY_FAILURES and sample.labels.get("requested_model") == group
|
||||
)
|
||||
|
||||
|
||||
def cooled_down_samples(gateway: Gateway, deployment_id: str) -> tuple[Sample, ...]:
|
||||
return tuple(
|
||||
sample
|
||||
for sample in scrape(gateway)
|
||||
if sample.name == COOLED_DOWN and sample.labels.get("model_id") == deployment_id
|
||||
)
|
||||
|
||||
|
||||
def assert_payment_required_metrics(gateway: Gateway, group: str, deployment_id: str) -> None:
|
||||
failures: Final = eventually(lambda: failure_samples(gateway, group), lambda found: len(found) >= 1, seconds=10)
|
||||
assert [(sample.labels["exception_class"], sample.labels["exception_status"]) for sample in failures] == [
|
||||
("Anthropic.PaymentRequiredError", "402")
|
||||
], failures
|
||||
cooldowns: Final = eventually(
|
||||
lambda: cooled_down_samples(gateway, deployment_id), lambda found: len(found) >= 1, seconds=10
|
||||
)
|
||||
assert [sample.labels["exception_status"] for sample in cooldowns] == ["402"], cooldowns
|
||||
|
||||
|
||||
def test_chat_openai_sdk_402_cools_the_deployment_and_the_healthy_one_serves_after(rig: CooldownRig) -> None:
|
||||
group: Final = new_group()
|
||||
with (
|
||||
wire_server(refusing_reply) as refusing,
|
||||
wire_server(healthy_reply) as healthy,
|
||||
rig.gateway.scenario() as scenario,
|
||||
):
|
||||
refused_id: Final = two_deployments(scenario, group, refusing, healthy)
|
||||
client: Final = openai_client(rig.gateway)
|
||||
history: Final = until_refused(lambda marker: openai_chat(client, group, marker))
|
||||
assert_refused(history[-1])
|
||||
assert_served(history[:-1], chat_identity)
|
||||
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
||||
assert markers_received(refusing) == (history[-1].marker,)
|
||||
assert markers_received(healthy) == markers(history[:-1])
|
||||
served: Final = tuple(openai_chat(client, group, new_marker()) for _ in range(20))
|
||||
assert_served(served, chat_identity)
|
||||
assert markers_received(healthy) == markers(served)
|
||||
assert markers_received(refusing) == ()
|
||||
assert_payment_required_row(history[-1].call_id, refused_id)
|
||||
assert_payment_required_metrics(rig.gateway, group, refused_id)
|
||||
|
||||
|
||||
def test_messages_anthropic_sdk_402_cools_the_deployment_and_the_healthy_one_serves_after(rig: CooldownRig) -> None:
|
||||
group: Final = new_group()
|
||||
with (
|
||||
wire_server(refusing_reply) as refusing,
|
||||
wire_server(healthy_reply) as healthy,
|
||||
rig.gateway.scenario() as scenario,
|
||||
):
|
||||
refused_id: Final = two_deployments(scenario, group, refusing, healthy)
|
||||
client: Final = anthropic_client(rig.gateway)
|
||||
history: Final = until_refused(lambda marker: anthropic_message(client, group, marker))
|
||||
assert_refused(history[-1])
|
||||
assert_served(history[:-1], served_text)
|
||||
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
||||
assert markers_received(refusing) == (history[-1].marker,)
|
||||
assert markers_received(healthy) == markers(history[:-1])
|
||||
served: Final = tuple(anthropic_message(client, group, new_marker()) for _ in range(20))
|
||||
assert_served(served, served_text)
|
||||
assert markers_received(healthy) == markers(served)
|
||||
assert markers_received(refusing) == ()
|
||||
assert_payment_required_row(history[-1].call_id, refused_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
|
||||
def test_responses_httpx_402_cools_the_deployment_and_the_healthy_one_serves_after(
|
||||
rig: CooldownRig, stream: bool
|
||||
) -> None:
|
||||
group: Final = new_group()
|
||||
with (
|
||||
wire_server(refusing_reply) as refusing,
|
||||
wire_server(healthy_reply) as healthy,
|
||||
rig.gateway.scenario() as scenario,
|
||||
):
|
||||
refused_id: Final = two_deployments(scenario, group, refusing, healthy)
|
||||
|
||||
def attempt(marker: str) -> Outcome:
|
||||
return raw_call(
|
||||
rig.gateway, Call(marker, "/v1/responses", {"model": group, "input": marker, "stream": stream})
|
||||
)
|
||||
|
||||
history: Final = until_refused(attempt)
|
||||
assert_refused(history[-1])
|
||||
assert_served(history[:-1], response_identity)
|
||||
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
||||
assert markers_received(refusing) == (history[-1].marker,)
|
||||
assert markers_received(healthy) == markers(history[:-1])
|
||||
served: Final = tuple(attempt(new_marker()) for _ in range(20))
|
||||
assert_served(served, response_identity)
|
||||
assert markers_received(healthy) == markers(served)
|
||||
assert markers_received(refusing) == ()
|
||||
assert_payment_required_row(history[-1].call_id, refused_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_info",
|
||||
({}, {"allowed_fails_policy": None}, {"allowed_fails_policy": {}}),
|
||||
ids=("missing", "null", "empty"),
|
||||
)
|
||||
def test_a_lone_402_deployment_without_a_policy_is_never_cooled(
|
||||
rig: CooldownRig, model_info: dict[str, JsonValue]
|
||||
) -> None:
|
||||
group: Final = new_group()
|
||||
with wire_server(refusing_reply) as refusing, rig.gateway.scenario() as scenario:
|
||||
refused_id: Final = deployment(scenario, group, anthropic_params(refusing), model_info=model_info)
|
||||
refused: Final = tuple(raw_chat(rig.gateway, group, new_marker()) for _ in range(5))
|
||||
for outcome in refused:
|
||||
assert_refused(outcome)
|
||||
assert_payment_required_row(outcome.call_id, refused_id)
|
||||
assert not rig.cooled(refused_id)
|
||||
sixth: Final = raw_chat(rig.gateway, group, new_marker())
|
||||
assert_refused(sixth)
|
||||
assert markers_received(refusing) == (*markers(refused), sixth.marker)
|
||||
|
||||
|
||||
@pytest.mark.timeout(900)
|
||||
def test_a_router_allowed_fails_policy_cools_a_lone_402_deployment_past_its_count(tmp_path: Path) -> None:
|
||||
group: Final = new_group()
|
||||
with ExitStack() as stack:
|
||||
upstream: Final = stack.enter_context(gateway_from_environment())
|
||||
cache: Final = stack.enter_context(owned_redis(tmp_path))
|
||||
settings: Final = router_settings(cache, allowed_fails_policy={"BadRequestErrorAllowedFails": 1})
|
||||
owned: Final = stack.enter_context(
|
||||
owned_proxy_process(
|
||||
upstream,
|
||||
tmp_path,
|
||||
{},
|
||||
config=cooldown_config(tmp_path, settings),
|
||||
remove_environment=("PROMETHEUS_MULTIPROC_DIR",),
|
||||
)
|
||||
)
|
||||
refusing: Final = stack.enter_context(wire_server(refusing_reply))
|
||||
scenario: Final = stack.enter_context(owned.gateway.scenario())
|
||||
refused_id: Final = deployment(scenario, group, anthropic_params(refusing))
|
||||
first: Final = raw_chat(owned.gateway, group, new_marker())
|
||||
second: Final = raw_chat(owned.gateway, group, new_marker())
|
||||
assert_refused(first)
|
||||
assert_refused(second)
|
||||
eventually(lambda: cooled(cache, refused_id), lambda seen: seen, seconds=10)
|
||||
third: Final = raw_chat(owned.gateway, group, new_marker())
|
||||
assert third.status == 429 and NO_DEPLOYMENTS in third.detail, (third.status, third.detail)
|
||||
assert markers_received(refusing) == (first.marker, second.marker)
|
||||
|
||||
|
||||
def test_a_deployment_allowed_fails_policy_cools_a_lone_402_deployment(rig: CooldownRig) -> None:
|
||||
group: Final = new_group()
|
||||
policy: Final = {"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}
|
||||
with wire_server(refusing_reply) as refusing, rig.gateway.scenario() as scenario:
|
||||
refused_id: Final = deployment(scenario, group, anthropic_params(refusing), model_info=policy)
|
||||
first: Final = raw_chat(rig.gateway, group, new_marker())
|
||||
assert_refused(first)
|
||||
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
||||
second: Final = raw_chat(rig.gateway, group, new_marker())
|
||||
assert second.status == 429 and NO_DEPLOYMENTS in second.detail, (second.status, second.detail)
|
||||
assert markers_received(refusing) == (first.marker,)
|
||||
|
||||
|
||||
def test_a_plain_allowed_fails_count_leaves_a_lone_402_deployment_warm(rig: CooldownRig) -> None:
|
||||
group: Final = new_group()
|
||||
with wire_server(refusing_reply) as refusing, rig.gateway.scenario() as scenario:
|
||||
refused_id: Final = deployment(scenario, group, anthropic_params(refusing), model_info={"allowed_fails": 0})
|
||||
refused: Final = tuple(raw_chat(rig.gateway, group, new_marker()) for _ in range(5))
|
||||
for outcome in refused:
|
||||
assert_refused(outcome)
|
||||
assert not rig.cooled(refused_id)
|
||||
sixth: Final = raw_chat(rig.gateway, group, new_marker())
|
||||
assert_refused(sixth)
|
||||
assert markers_received(refusing) == (*markers(refused), sixth.marker)
|
||||
|
||||
|
||||
def test_an_openai_compatible_402_cools_the_deployment_too(rig: CooldownRig) -> None:
|
||||
group: Final = new_group()
|
||||
with (
|
||||
wire_server(openai_refusal) as refusing,
|
||||
wire_server(healthy_reply) as healthy,
|
||||
rig.gateway.scenario() as scenario,
|
||||
):
|
||||
refused_id: Final = deployment(scenario, group, openai_params(refusing))
|
||||
deployment(scenario, group, openai_params(healthy))
|
||||
history: Final = until_refused(lambda marker: raw_chat(rig.gateway, group, marker))
|
||||
assert history[-1].status == 402 and "scripted 402" in history[-1].detail, history[-1]
|
||||
assert_served(history[:-1], chat_identity)
|
||||
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
||||
assert markers_received(refusing) == (history[-1].marker,)
|
||||
served: Final = tuple(raw_chat(rig.gateway, group, new_marker()) for _ in range(20))
|
||||
assert_served(served, chat_identity)
|
||||
assert markers_received(refusing) == ()
|
||||
assert markers_received(healthy) == (*markers(history[:-1]), *markers(served))
|
||||
|
||||
|
||||
def test_a_lone_402_primary_with_a_fallback_group_stays_warm_and_the_fallback_answers(rig: CooldownRig) -> None:
|
||||
with (
|
||||
wire_server(refusing_reply) as refusing,
|
||||
wire_server(healthy_reply) as healthy,
|
||||
rig.gateway.scenario() as scenario,
|
||||
):
|
||||
primary_id: Final = deployment(scenario, PRIMARY_GROUP, anthropic_params(refusing))
|
||||
deployment(scenario, FALLBACK_GROUP, openai_params(healthy))
|
||||
served: Final = tuple(raw_chat(rig.gateway, PRIMARY_GROUP, new_marker()) for _ in range(3))
|
||||
assert_served(served, chat_identity)
|
||||
assert markers_received(refusing) == markers(served)
|
||||
assert markers_received(healthy) == markers(served)
|
||||
assert not rig.cooled(primary_id)
|
||||
|
||||
|
||||
def test_a_deployment_policy_added_while_402_traffic_flows_starts_the_cooldown(rig: CooldownRig) -> None:
|
||||
group: Final = new_group()
|
||||
with wire_server(refusing_reply) as refusing, rig.gateway.scenario() as scenario:
|
||||
refused_id: Final = deployment(scenario, group, anthropic_params(refusing))
|
||||
before: Final = tuple(raw_chat(rig.gateway, group, new_marker()) for _ in range(3))
|
||||
for outcome in before:
|
||||
assert_refused(outcome)
|
||||
assert not rig.cooled(refused_id)
|
||||
policy: Final = {"id": refused_id, "allowed_fails_policy": {"BadRequestErrorAllowedFails": 0}}
|
||||
patched: Final = rig.gateway.request("PATCH", f"/model/{refused_id}/update", {"model_info": policy})
|
||||
assert patched.status_code == 200, patched.text
|
||||
|
||||
def probe() -> tuple[int, bool]:
|
||||
return raw_chat(rig.gateway, group, new_marker()).status, rig.cooled(refused_id)
|
||||
|
||||
eventually(probe, lambda seen: seen[1], seconds=RELOAD_SECONDS * 10)
|
||||
after: Final = raw_chat(rig.gateway, group, new_marker())
|
||||
assert after.status == 429 and NO_DEPLOYMENTS in after.detail, (after.status, after.detail)
|
||||
assert len(markers_received(refusing)) >= 4
|
||||
|
||||
|
||||
def test_a_cooled_402_deployment_comes_back_after_its_own_cooldown_time(rig: CooldownRig) -> None:
|
||||
group: Final = new_group()
|
||||
with (
|
||||
wire_server(refusing_reply) as refusing,
|
||||
wire_server(healthy_reply) as healthy,
|
||||
rig.gateway.scenario() as scenario,
|
||||
):
|
||||
refused_id: Final = deployment(scenario, group, {**anthropic_params(refusing), "cooldown_time": 2})
|
||||
deployment(scenario, group, openai_params(healthy))
|
||||
history: Final = until_refused(lambda marker: raw_chat(rig.gateway, group, marker))
|
||||
assert_refused(history[-1])
|
||||
eventually(lambda: rig.cooled(refused_id), lambda seen: seen, seconds=10)
|
||||
eventually(lambda: rig.cooled(refused_id), lambda seen: not seen, seconds=10)
|
||||
again: Final = until_refused(lambda marker: raw_chat(rig.gateway, group, marker))
|
||||
assert_refused(again[-1])
|
||||
assert markers_received(refusing) == (history[-1].marker, again[-1].marker)
|
||||
|
||||
|
||||
def burst_call(group: str, index: int) -> Call:
|
||||
marker: Final = f"{new_marker()}-{index}"
|
||||
stream: Final = index % 2 == 1
|
||||
match index % 3:
|
||||
case 0:
|
||||
return Call(marker, "/v1/chat/completions", chat_body(group, marker, stream=stream))
|
||||
case 1:
|
||||
body: Final = {"model": group, "max_tokens": 32, "messages": [{"role": "user", "content": marker}]}
|
||||
return Call(marker, "/v1/messages", {**body, "stream": stream})
|
||||
case _:
|
||||
return Call(marker, "/v1/responses", {"model": group, "input": marker, "stream": stream})
|
||||
|
||||
|
||||
def worker_pids(log: Path) -> tuple[int, ...]:
|
||||
return tuple(int(pid) for pid in WORKER_PID.findall(log.read_text()))
|
||||
|
||||
|
||||
def ready_workers(log: Path) -> int:
|
||||
return log.read_text().count("Application startup complete.")
|
||||
|
||||
|
||||
@pytest.mark.timeout(900)
|
||||
def test_chaos_burst_cooldown_holds_across_both_workers_and_a_worker_kill(tmp_path: Path) -> None:
|
||||
group: Final = new_group()
|
||||
refused_id: Final = f"refusing-{uuid.uuid4().hex[:12]}"
|
||||
healthy_id: Final = f"healthy-{uuid.uuid4().hex[:12]}"
|
||||
with ExitStack() as stack:
|
||||
upstream: Final = stack.enter_context(gateway_from_environment())
|
||||
cache: Final = stack.enter_context(owned_redis(tmp_path))
|
||||
refusing: Final = stack.enter_context(wire_server(refusing_reply))
|
||||
healthy: Final = stack.enter_context(wire_server(healthy_reply))
|
||||
model_list: Final = (
|
||||
{"model_name": group, "litellm_params": anthropic_params(refusing), "model_info": {"id": refused_id}},
|
||||
{"model_name": group, "litellm_params": openai_params(healthy), "model_info": {"id": healthy_id}},
|
||||
)
|
||||
config: Final = cooldown_config(tmp_path, router_settings(cache), model_list=model_list, prometheus=False)
|
||||
owned: Final = stack.enter_context(owned_proxy_process(upstream, tmp_path, {}, config=config, workers=2))
|
||||
eventually(lambda: ready_workers(owned.log), lambda ready: ready == 2, seconds=graceful_stop_seconds())
|
||||
client: Final = stack.enter_context(
|
||||
httpx.Client(base_url=owned.gateway.client.base_url, timeout=60, trust_env=False, limits=GATEWAY_LIMITS)
|
||||
)
|
||||
patient: Final = Gateway(client, owned.gateway.key, owned.gateway.upstream_url)
|
||||
calls: Final = tuple(burst_call(group, index) for index in range(30))
|
||||
with ThreadPoolExecutor(max_workers=30) as pool:
|
||||
outcomes: Final = tuple(pool.map(lambda call: raw_call(patient, call), calls))
|
||||
assert {outcome.status for outcome in outcomes} <= {200, 402}, [(o.status, o.detail[:200]) for o in outcomes]
|
||||
refused: Final = tuple(outcome for outcome in outcomes if outcome.status == 402)
|
||||
served: Final = tuple(outcome for outcome in outcomes if outcome.status == 200)
|
||||
assert refused, "the burst never reached the refusing deployment"
|
||||
for outcome in refused:
|
||||
assert_refused(outcome)
|
||||
assert all(outcome.marker in outcome.detail for outcome in served), [o.detail[:200] for o in served]
|
||||
assert sorted(markers_received(refusing)) == sorted(markers(refused))
|
||||
assert sorted(markers_received(healthy)) == sorted(markers(served))
|
||||
eventually(lambda: cooled(cache, refused_id), lambda seen: seen, seconds=10)
|
||||
warm: Final = tuple(raw_chat(patient, group, new_marker()) for _ in range(20))
|
||||
assert_served(warm, chat_identity)
|
||||
assert markers_received(refusing) == ()
|
||||
assert markers_received(healthy) == markers(warm)
|
||||
victim: Final = next(pid for pid in sorted(worker_pids(owned.log)) if psutil.pid_exists(pid))
|
||||
ready_before: Final = ready_workers(owned.log)
|
||||
psutil.Process(victim).send_signal(signal.SIGKILL)
|
||||
eventually(
|
||||
lambda: ready_workers(owned.log), lambda ready: ready > ready_before, seconds=graceful_stop_seconds()
|
||||
)
|
||||
after_kill: Final = tuple(raw_chat(patient, group, new_marker()) for _ in range(20))
|
||||
assert_served(after_kill, chat_identity)
|
||||
assert markers_received(refusing) == ()
|
||||
assert markers_received(healthy) == markers(after_kill)
|
||||
387
tests/integration/routing/test_payment_required_mapping.py
Normal file
387
tests/integration/routing/test_payment_required_mapping.py
Normal file
|
|
@ -0,0 +1,387 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final, TypeVar
|
||||
|
||||
import anthropic
|
||||
import httpx
|
||||
import openai
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.openai_wire import chat_reply, openai_error
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue
|
||||
|
||||
from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
PROXY_WORKERS: Final = int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))
|
||||
WORKER_SYNC_SECONDS: Final = 0.0 if PROXY_WORKERS == 1 else PROXY_CONFIG_RELOAD_INTERVAL_SECONDS + 5.0
|
||||
FRESH_CONNECTION: Final = {"Connection": "close"}
|
||||
UNKNOWN_MODEL: Final = "Invalid model name passed in model="
|
||||
PAYMENT_REQUIRED: Final = "litellm.PaymentRequiredError: AnthropicException"
|
||||
BAD_REQUEST: Final = "litellm.BadRequestError: AnthropicException"
|
||||
ANTHROPIC_MODEL: Final = "anthropic/claude-sonnet-5"
|
||||
OPENAI_MODEL: Final = "openai/gpt-5.4"
|
||||
CREDIT_BALANCE_TEXT: Final = "Your credit balance is too low to access the Anthropic API. Please go to Plans & Billing to upgrade or purchase credits."
|
||||
FAILURE_ROW: Final = (
|
||||
"SELECT status, model_id, metadata::jsonb -> 'error_information' AS error "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
|
||||
)
|
||||
MODEL_INFO_PROBE: Final = ("GET", "/v1/models")
|
||||
|
||||
|
||||
def is_model_info_probe(request: Request) -> bool:
|
||||
return (request.method, request.target) == MODEL_INFO_PROBE
|
||||
|
||||
|
||||
def model_list_reply() -> Reply:
|
||||
return Reply(body=b'{"object": "list", "data": []}')
|
||||
|
||||
|
||||
def openai_refusal(request: Request) -> Reply:
|
||||
return model_list_reply() if is_model_info_probe(request) else openai_error(402)
|
||||
|
||||
|
||||
def provider_calls(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(request for request in wire.drain() if not is_model_info_probe(request))
|
||||
|
||||
|
||||
def new_marker() -> str:
|
||||
return f"pr402-{uuid.uuid4().hex}"
|
||||
|
||||
|
||||
def new_group() -> str:
|
||||
return f"integration-402-{uuid.uuid4().hex[:12]}"
|
||||
|
||||
|
||||
def anthropic_error(status: int, message: str, *, content_type: str = "application/json") -> Reply:
|
||||
body: Final = json.dumps({"type": "error", "error": {"type": "invalid_request_error", "message": message}})
|
||||
return Reply(status=status, body=body.encode(), content_type=content_type)
|
||||
|
||||
|
||||
def anthropic_params(wire: Wire) -> dict[str, JsonValue]:
|
||||
return {"model": ANTHROPIC_MODEL, "api_base": wire.url, "api_key": "synthetic-anthropic-key"}
|
||||
|
||||
|
||||
def openai_params(wire: Wire) -> dict[str, JsonValue]:
|
||||
return {"model": OPENAI_MODEL, "api_base": wire.url + "/v1", "api_key": "synthetic-openai-key"}
|
||||
|
||||
|
||||
def deployment(
|
||||
scenario: Scenario,
|
||||
group: str,
|
||||
litellm_params: Mapping[str, JsonValue],
|
||||
*,
|
||||
model_info: Mapping[str, JsonValue] | None = None,
|
||||
) -> str:
|
||||
created: Final = scenario.gateway.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": group,
|
||||
"litellm_params": dict(litellm_params),
|
||||
"model_info": dict(model_info) if model_info is not None else {},
|
||||
},
|
||||
)
|
||||
identity: Final = string_value(object_value(created["model_info"])["id"])
|
||||
scenario.cleanups.callback(scenario.delete_model, identity)
|
||||
return identity
|
||||
|
||||
|
||||
def settled(send: Callable[[], T], text: Callable[[T], str]) -> T:
|
||||
return eventually(send, lambda observed: UNKNOWN_MODEL not in text(observed), seconds=WORKER_SYNC_SECONDS + 10)
|
||||
|
||||
|
||||
def post(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> httpx.Response:
|
||||
return settled(
|
||||
lambda: gateway.request("POST", path, body, headers=FRESH_CONNECTION), lambda response: response.text
|
||||
)
|
||||
|
||||
|
||||
def chat_body(group: str, marker: str, **extra: JsonValue) -> dict[str, JsonValue]:
|
||||
return {"model": group, "messages": [{"role": "user", "content": marker}], **extra}
|
||||
|
||||
|
||||
def upstream_calls(wire: Wire, marker: str) -> tuple[str, ...]:
|
||||
received: Final = provider_calls(wire)
|
||||
assert all(marker.encode() in request.body for request in received), received
|
||||
return tuple(f"{request.method} {request.target}" for request in received)
|
||||
|
||||
|
||||
def failure_row(call_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(lambda: read_rows(FAILURE_ROW, (call_id,)), lambda found: len(found) == 1, seconds=30)
|
||||
return rows[0]
|
||||
|
||||
|
||||
def assert_payment_required_row(call_id: str, deployment_id: str) -> None:
|
||||
row: Final = failure_row(call_id)
|
||||
error: Final = object_value(row["error"])
|
||||
assert row["status"] == "failure", row
|
||||
assert row["model_id"] == deployment_id, row
|
||||
assert (error["error_class"], error["error_code"], error["llm_provider"]) == (
|
||||
"PaymentRequiredError",
|
||||
"402",
|
||||
"anthropic",
|
||||
), row
|
||||
|
||||
|
||||
def openai_client(gateway: Gateway) -> openai.OpenAI:
|
||||
return openai.OpenAI(
|
||||
base_url=str(gateway.client.base_url).rstrip("/") + "/v1",
|
||||
api_key=gateway.key,
|
||||
max_retries=0,
|
||||
default_headers=FRESH_CONNECTION,
|
||||
)
|
||||
|
||||
|
||||
def anthropic_client(gateway: Gateway) -> anthropic.Anthropic:
|
||||
return anthropic.Anthropic(
|
||||
base_url=str(gateway.client.base_url).rstrip("/"),
|
||||
api_key=gateway.key,
|
||||
max_retries=0,
|
||||
default_headers=FRESH_CONNECTION,
|
||||
)
|
||||
|
||||
|
||||
def test_chat_completions_openai_sdk_gets_402_payment_required(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda _: anthropic_error(402, f"scripted 402 {marker}")) as wire, gateway.scenario() as scenario:
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, anthropic_params(wire))
|
||||
client: Final = openai_client(gateway)
|
||||
|
||||
def send() -> openai.APIStatusError:
|
||||
with pytest.raises(openai.APIStatusError) as caught:
|
||||
client.chat.completions.create(model=group, messages=[{"role": "user", "content": marker}])
|
||||
return caught.value
|
||||
|
||||
error: Final = settled(send, lambda failure: failure.message)
|
||||
assert error.status_code == 402, error.message
|
||||
assert PAYMENT_REQUIRED in error.message, error.message
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/messages",)
|
||||
assert_payment_required_row(error.response.headers["x-litellm-call-id"], deployment_id)
|
||||
|
||||
|
||||
async def drain_chat_stream(client: openai.AsyncOpenAI, group: str, marker: str) -> None:
|
||||
stream: Final = await client.chat.completions.create(
|
||||
model=group, messages=[{"role": "user", "content": marker}], stream=True
|
||||
)
|
||||
async for _ in stream:
|
||||
pass
|
||||
|
||||
|
||||
async def drain_messages_stream(client: anthropic.AsyncAnthropic, group: str, marker: str) -> None:
|
||||
async with client.messages.stream(
|
||||
model=group, max_tokens=32, messages=[{"role": "user", "content": marker}]
|
||||
) as stream:
|
||||
async for _ in stream:
|
||||
pass
|
||||
|
||||
|
||||
def test_chat_completions_stream_openai_async_sdk_gets_402_payment_required(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda _: anthropic_error(402, f"scripted 402 {marker}")) as wire, gateway.scenario() as scenario:
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, anthropic_params(wire))
|
||||
|
||||
async def attempt() -> openai.APIStatusError:
|
||||
async with openai.AsyncOpenAI(
|
||||
base_url=str(gateway.client.base_url).rstrip("/") + "/v1",
|
||||
api_key=gateway.key,
|
||||
max_retries=0,
|
||||
default_headers=FRESH_CONNECTION,
|
||||
) as client:
|
||||
with pytest.raises(openai.APIStatusError) as caught:
|
||||
await drain_chat_stream(client, group, marker)
|
||||
return caught.value
|
||||
|
||||
error: Final = settled(lambda: asyncio.run(attempt()), lambda failure: failure.message)
|
||||
assert error.status_code == 402, error.message
|
||||
assert PAYMENT_REQUIRED in error.message, error.message
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/messages",)
|
||||
assert_payment_required_row(error.response.headers["x-litellm-call-id"], deployment_id)
|
||||
|
||||
|
||||
def test_messages_anthropic_sdk_gets_402_payment_required(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda _: anthropic_error(402, f"scripted 402 {marker}")) as wire, gateway.scenario() as scenario:
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, anthropic_params(wire))
|
||||
client: Final = anthropic_client(gateway)
|
||||
|
||||
def send() -> anthropic.APIStatusError:
|
||||
with pytest.raises(anthropic.APIStatusError) as caught:
|
||||
client.messages.create(model=group, max_tokens=32, messages=[{"role": "user", "content": marker}])
|
||||
return caught.value
|
||||
|
||||
error: Final = settled(send, lambda failure: failure.message)
|
||||
assert error.status_code == 402, error.message
|
||||
assert PAYMENT_REQUIRED in error.message, error.message
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/messages",)
|
||||
assert_payment_required_row(error.response.headers["x-litellm-call-id"], deployment_id)
|
||||
|
||||
|
||||
def test_messages_stream_anthropic_async_sdk_gets_402_payment_required(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda _: anthropic_error(402, f"scripted 402 {marker}")) as wire, gateway.scenario() as scenario:
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, anthropic_params(wire))
|
||||
|
||||
async def attempt() -> anthropic.APIStatusError:
|
||||
async with anthropic.AsyncAnthropic(
|
||||
base_url=str(gateway.client.base_url).rstrip("/"),
|
||||
api_key=gateway.key,
|
||||
max_retries=0,
|
||||
default_headers=FRESH_CONNECTION,
|
||||
) as client:
|
||||
with pytest.raises(anthropic.APIStatusError) as caught:
|
||||
await drain_messages_stream(client, group, marker)
|
||||
return caught.value
|
||||
|
||||
error: Final = settled(lambda: asyncio.run(attempt()), lambda failure: failure.message)
|
||||
assert error.status_code == 402, error.message
|
||||
assert PAYMENT_REQUIRED in error.message, error.message
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/messages",)
|
||||
assert_payment_required_row(error.response.headers["x-litellm-call-id"], deployment_id)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
|
||||
def test_responses_httpx_gets_402_payment_required(gateway: Gateway, stream: bool) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda _: anthropic_error(402, f"scripted 402 {marker}")) as wire, gateway.scenario() as scenario:
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, anthropic_params(wire))
|
||||
response: Final = post(gateway, "/v1/responses", {"model": group, "input": marker, "stream": stream})
|
||||
assert response.status_code == 402, response.text
|
||||
assert PAYMENT_REQUIRED in response.text, response.text
|
||||
assert response.headers["content-type"].startswith("application/json"), response.headers
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/messages",)
|
||||
assert_payment_required_row(response.headers["x-litellm-call-id"], deployment_id)
|
||||
|
||||
|
||||
def test_402_with_a_text_body_still_maps_to_payment_required(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
text: Final = f"credits exhausted {marker}"
|
||||
with (
|
||||
wire_server(lambda _: Reply(status=402, body=text.encode(), content_type="text/plain")) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, anthropic_params(wire))
|
||||
response: Final = post(gateway, "/v1/chat/completions", chat_body(group, marker))
|
||||
assert response.status_code == 402, response.text
|
||||
assert PAYMENT_REQUIRED in response.text, response.text
|
||||
assert text in response.text, response.text
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/messages",)
|
||||
assert_payment_required_row(response.headers["x-litellm-call-id"], deployment_id)
|
||||
|
||||
|
||||
def test_402_with_an_empty_body_still_maps_to_payment_required(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda _: Reply(status=402, body=b"")) as wire, gateway.scenario() as scenario:
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, anthropic_params(wire))
|
||||
response: Final = post(gateway, "/v1/chat/completions", chat_body(group, marker))
|
||||
assert response.status_code == 402, response.text
|
||||
assert PAYMENT_REQUIRED in response.text, response.text
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/messages",)
|
||||
assert_payment_required_row(response.headers["x-litellm-call-id"], deployment_id)
|
||||
|
||||
|
||||
def test_402_with_a_5kb_message_maps_and_a_healthy_deployment_keeps_serving(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
long_message: Final = f"{marker} " + "x" * 5120
|
||||
with (
|
||||
wire_server(lambda _: anthropic_error(402, long_message)) as broke,
|
||||
wire_server(
|
||||
lambda _: chat_reply(f"chatcmpl-{marker}", OPENAI_MODEL, f"served {marker}", stream=False)
|
||||
) as healthy,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
broke_group: Final = new_group()
|
||||
healthy_group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, broke_group, anthropic_params(broke))
|
||||
deployment(scenario, healthy_group, openai_params(healthy))
|
||||
refused: Final = post(gateway, "/v1/chat/completions", chat_body(broke_group, marker))
|
||||
assert refused.status_code == 402, refused.text
|
||||
assert PAYMENT_REQUIRED in refused.text, refused.text
|
||||
assert "x" * 512 in refused.text, refused.text
|
||||
assert upstream_calls(broke, marker) == ("POST /v1/messages",)
|
||||
assert_payment_required_row(refused.headers["x-litellm-call-id"], deployment_id)
|
||||
served: Final = post(gateway, "/v1/chat/completions", chat_body(healthy_group, f"{marker} healthy"))
|
||||
assert served.status_code == 200, served.text
|
||||
assert served.json()["choices"][0]["message"]["content"] == f"served {marker}", served.text
|
||||
assert upstream_calls(healthy, marker) == ("POST /v1/chat/completions",)
|
||||
|
||||
|
||||
def test_unauthenticated_request_to_a_402_deployment_is_refused_before_the_wire(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(lambda _: anthropic_error(402, f"scripted 402 {marker}")) as wire, gateway.scenario() as scenario:
|
||||
group: Final = new_group()
|
||||
deployment(scenario, group, anthropic_params(wire))
|
||||
response: Final = gateway.client.post("/v1/chat/completions", json=chat_body(group, marker))
|
||||
assert response.status_code == 401, response.text
|
||||
assert provider_calls(wire) == ()
|
||||
|
||||
|
||||
def test_anthropic_400_credit_balance_text_stays_a_bad_request(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with (
|
||||
wire_server(lambda _: anthropic_error(400, f"{CREDIT_BALANCE_TEXT} {marker}")) as wire,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, anthropic_params(wire))
|
||||
response: Final = post(gateway, "/v1/chat/completions", chat_body(group, marker))
|
||||
assert response.status_code == 400, response.text
|
||||
assert BAD_REQUEST in response.text, response.text
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/messages",)
|
||||
row: Final = failure_row(response.headers["x-litellm-call-id"])
|
||||
error: Final = object_value(row["error"])
|
||||
assert (row["status"], row["model_id"]) == ("failure", deployment_id), row
|
||||
assert (error["error_class"], error["error_code"]) == ("BadRequestError", "400"), row
|
||||
|
||||
|
||||
def test_openai_compatible_402_reaches_the_caller_as_402(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with wire_server(openai_refusal) as wire, gateway.scenario() as scenario:
|
||||
group: Final = new_group()
|
||||
deployment_id: Final = deployment(scenario, group, openai_params(wire))
|
||||
response: Final = post(gateway, "/v1/chat/completions", chat_body(group, marker))
|
||||
assert response.status_code == 402, response.text
|
||||
assert "scripted 402" in response.text, response.text
|
||||
assert upstream_calls(wire, marker) == ("POST /v1/chat/completions",)
|
||||
row: Final = failure_row(response.headers["x-litellm-call-id"])
|
||||
error: Final = object_value(row["error"])
|
||||
assert (row["status"], row["model_id"]) == ("failure", deployment_id), row
|
||||
assert (error["error_code"], error["llm_provider"]) == ("402", "openai"), row
|
||||
|
||||
|
||||
def test_disabled_cooldowns_keep_picking_the_402_deployment(gateway: Gateway) -> None:
|
||||
marker: Final = new_marker()
|
||||
with (
|
||||
wire_server(lambda _: anthropic_error(402, f"scripted 402 {marker}")) as broke,
|
||||
wire_server(
|
||||
lambda _: chat_reply(f"chatcmpl-{marker}", OPENAI_MODEL, f"served {marker}", stream=False)
|
||||
) as healthy,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
group: Final = new_group()
|
||||
deployment(scenario, group, anthropic_params(broke))
|
||||
deployment(scenario, group, openai_params(healthy))
|
||||
responses: Final = tuple(
|
||||
post(gateway, "/v1/chat/completions", chat_body(group, f"{marker}-{index}")) for index in range(40)
|
||||
)
|
||||
statuses: Final = tuple(response.status_code for response in responses)
|
||||
assert statuses.count(402) >= 2 and statuses.count(200) >= 1, statuses
|
||||
assert set(statuses) == {200, 402}, statuses
|
||||
refusals: Final = tuple(response for response in responses if response.status_code == 402)
|
||||
assert all(PAYMENT_REQUIRED in response.text for response in refusals), [r.text for r in refusals]
|
||||
assert len(upstream_calls(broke, marker)) == len(refusals)
|
||||
assert len(upstream_calls(healthy, marker)) == statuses.count(200)
|
||||
|
|
@ -773,6 +773,14 @@ class TestRouterLevelGetAllowedFailsFromPolicy:
|
|||
exc = litellm.NotFoundError("404 error", "openai", "gpt-4")
|
||||
assert router.get_allowed_fails_from_policy(exc) == 1
|
||||
|
||||
def test_payment_required_error_uses_bad_request_allowed_fails(self):
|
||||
assert (
|
||||
self._make_router(BadRequestErrorAllowedFails=6).get_allowed_fails_from_policy(
|
||||
litellm.PaymentRequiredError("402 error", "openai", "gpt-4")
|
||||
)
|
||||
== 6
|
||||
)
|
||||
|
||||
def test_unmatched_exception_returns_none(self):
|
||||
router = self._make_router(InternalServerErrorAllowedFails=5)
|
||||
exc = litellm.RateLimitError("429", "openai", "gpt-4")
|
||||
|
|
|
|||
|
|
@ -265,6 +265,114 @@ def test_should_cooldown_deployment_auth_limit_error(testing_litellm_router):
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exception_status", (401, 402))
|
||||
def test_is_cooldown_required_for_account_errors(testing_litellm_router, exception_status):
|
||||
assert (
|
||||
_is_cooldown_required(
|
||||
litellm_router_instance=testing_litellm_router,
|
||||
model_id="test_deployment",
|
||||
exception_status=exception_status,
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("allowed_fails", (None, 0))
|
||||
def test_single_deployment_402_does_not_cooldown(
|
||||
allowed_fails: int | None,
|
||||
) -> None:
|
||||
assert (
|
||||
_should_cooldown_deployment(
|
||||
Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-mini",
|
||||
"litellm_params": {"model": "gpt-5-mini"},
|
||||
"model_info": {"id": "dep-1"},
|
||||
},
|
||||
],
|
||||
allowed_fails=allowed_fails,
|
||||
),
|
||||
"dep-1",
|
||||
402,
|
||||
litellm.PaymentRequiredError(
|
||||
message="Insufficient credits",
|
||||
model="gpt-5-mini",
|
||||
llm_provider="openai",
|
||||
),
|
||||
)
|
||||
is False
|
||||
)
|
||||
|
||||
|
||||
def test_single_deployment_402_respects_router_allowed_fails_policy() -> None:
|
||||
assert (
|
||||
_should_cooldown_deployment(
|
||||
Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-mini",
|
||||
"litellm_params": {"model": "gpt-5-mini"},
|
||||
"model_info": {"id": "dep-1"},
|
||||
},
|
||||
],
|
||||
allowed_fails_policy=AllowedFailsPolicy(BadRequestErrorAllowedFails=0),
|
||||
),
|
||||
"dep-1",
|
||||
402,
|
||||
litellm.PaymentRequiredError(
|
||||
message="Insufficient credits",
|
||||
model="gpt-5-mini",
|
||||
llm_provider="openai",
|
||||
),
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_single_deployment_402_respects_deployment_allowed_fails_policy() -> None:
|
||||
assert (
|
||||
_should_cooldown_deployment(
|
||||
Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-5-mini",
|
||||
"litellm_params": {"model": "gpt-5-mini"},
|
||||
"model_info": {
|
||||
"id": "dep-1",
|
||||
"allowed_fails_policy": {"BadRequestErrorAllowedFails": 0},
|
||||
},
|
||||
},
|
||||
],
|
||||
),
|
||||
"dep-1",
|
||||
402,
|
||||
litellm.PaymentRequiredError(
|
||||
message="Insufficient credits",
|
||||
model="gpt-5-mini",
|
||||
llm_provider="openai",
|
||||
),
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def test_multi_deployment_402_cools_down(testing_litellm_router: Router) -> None:
|
||||
assert (
|
||||
_should_cooldown_deployment(
|
||||
testing_litellm_router,
|
||||
"test_deployment",
|
||||
402,
|
||||
litellm.PaymentRequiredError(
|
||||
message="Insufficient credits",
|
||||
model="gpt-5-mini",
|
||||
llm_provider="openai",
|
||||
),
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_cooldown_deployment(testing_litellm_router):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -745,11 +745,12 @@ class _UpstreamHTTPError(Exception):
|
|||
self.response = httpx.Response(status_code=status_code, request=self.request, text="upstream failure")
|
||||
|
||||
|
||||
UPSTREAM_STATUS_CODES = (400, 401, 403, 404, 408, 422, 429, 500, 503)
|
||||
UPSTREAM_STATUS_CODES = (400, 401, 402, 403, 404, 408, 422, 429, 500, 503)
|
||||
|
||||
OPENAI_SHAPED = {
|
||||
400: (litellm.BadRequestError, 400),
|
||||
401: (litellm.AuthenticationError, 401),
|
||||
402: (litellm.PaymentRequiredError, 402),
|
||||
403: (litellm.APIError, 403),
|
||||
404: (litellm.NotFoundError, 404),
|
||||
408: (litellm.Timeout, 408),
|
||||
|
|
@ -764,8 +765,10 @@ PERMISSION_DENIED = (litellm.PermissionDeniedError, 403)
|
|||
STATUS_KEYED = {**OPENAI_SHAPED, 403: PERMISSION_DENIED}
|
||||
|
||||
DEVIATIONS_FROM_THE_OPENAI_SHAPE = {
|
||||
"ai21": {402: (litellm.APIError, 402)},
|
||||
"anthropic": {403: PERMISSION_DENIED},
|
||||
"azure": {500: (litellm.APIError, 500)},
|
||||
"azure": {402: (litellm.APIError, 402), 500: (litellm.APIError, 500)},
|
||||
"azure_ai": {402: (litellm.APIError, 402)},
|
||||
"bedrock": {
|
||||
403: PERMISSION_DENIED,
|
||||
500: (litellm.ServiceUnavailableError, 503),
|
||||
|
|
@ -776,13 +779,19 @@ DEVIATIONS_FROM_THE_OPENAI_SHAPE = {
|
|||
403: PERMISSION_DENIED,
|
||||
422: (litellm.BadRequestError, 400),
|
||||
},
|
||||
"deepseek": {402: (litellm.APIError, 402)},
|
||||
"fireworks_ai": {402: (litellm.APIError, 402)},
|
||||
"gemini": {403: PERMISSION_DENIED},
|
||||
"groq": {402: (litellm.APIError, 402)},
|
||||
"huggingface": {
|
||||
402: (litellm.APIError, 402),
|
||||
404: (litellm.APIError, 404),
|
||||
422: (litellm.APIError, 422),
|
||||
500: (litellm.APIError, 500),
|
||||
},
|
||||
"mistral": {402: (litellm.APIError, 402)},
|
||||
"nlp_cloud": {
|
||||
402: (litellm.RateLimitError, 429),
|
||||
403: (litellm.AuthenticationError, 403),
|
||||
404: (litellm.APIError, 404),
|
||||
408: (litellm.APIError, 408),
|
||||
|
|
@ -790,20 +799,26 @@ DEVIATIONS_FROM_THE_OPENAI_SHAPE = {
|
|||
503: (litellm.APIError, 503),
|
||||
},
|
||||
"ollama": {403: PERMISSION_DENIED},
|
||||
"openrouter": {500: (litellm.APIError, 500)},
|
||||
"openai": {402: (litellm.APIError, 402)},
|
||||
"openrouter": {402: (litellm.APIError, 402), 500: (litellm.APIError, 500)},
|
||||
"perplexity": {402: (litellm.APIError, 402)},
|
||||
"replicate": {
|
||||
403: (litellm.APIError, 500),
|
||||
404: (litellm.APIError, 500),
|
||||
402: (litellm.APIError, 500),
|
||||
422: (litellm.UnprocessableEntityError, 422),
|
||||
500: (litellm.ServiceUnavailableError, 503),
|
||||
503: (litellm.APIError, 500),
|
||||
},
|
||||
"runwayml": {402: (litellm.APIError, 402)},
|
||||
"sagemaker": {
|
||||
403: PERMISSION_DENIED,
|
||||
500: (litellm.ServiceUnavailableError, 503),
|
||||
},
|
||||
"together_ai": {402: (litellm.APIError, 402)},
|
||||
"vertex_ai": {403: PERMISSION_DENIED},
|
||||
"vllm": {403: PERMISSION_DENIED},
|
||||
"xai": {402: (litellm.APIError, 402)},
|
||||
}
|
||||
|
||||
PROVIDERS_WITH_A_HANDLER = (
|
||||
|
|
@ -936,6 +951,16 @@ def test_a_provider_without_a_handler_maps_by_the_upstream_status(provider, stat
|
|||
assert raised.value.status_code == expected_status
|
||||
assert raised.value.llm_provider == provider
|
||||
assert raised.value.model == "test-model"
|
||||
if status_code == 402:
|
||||
assert isinstance(raised.value, litellm.BadRequestError)
|
||||
assert raised.value.response.status_code == 402
|
||||
assert str(raised.value).startswith("litellm.PaymentRequiredError: ")
|
||||
assert "litellm.BadRequestError:" not in str(raised.value)
|
||||
assert repr(raised.value) == str(raised.value)
|
||||
|
||||
|
||||
def test_payment_required_error_without_response_uses_402_response() -> None:
|
||||
assert litellm.PaymentRequiredError("x", "m", "p").response.status_code == 402
|
||||
|
||||
|
||||
def test_a_minimax_bad_key_is_an_authentication_error(quiet_exception_mapping):
|
||||
|
|
|
|||
|
|
@ -79,6 +79,16 @@ def test_subclass_falls_back_to_the_parent_field():
|
|||
)
|
||||
|
||||
|
||||
def test_payment_required_error_uses_bad_request_retries():
|
||||
assert (
|
||||
get_num_retries_from_retry_policy(
|
||||
exception=_error(litellm.PaymentRequiredError),
|
||||
retry_policy=RetryPolicy(BadRequestErrorRetries=0, DefaultRetries=3),
|
||||
)
|
||||
== 0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("exception_type", (litellm.BadGatewayError,))
|
||||
def test_default_retries_covers_exceptions_without_a_specific_field(exception_type: type[Exception]):
|
||||
exception: Final = _error(exception_type)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue