mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
d5173e3d70
* test: move 76 live top-level tests to offline unit and integration coverage * test: keep the mock-param opt-in guard valid once the top-level senders are gone * test: address review on offline replacements * docs(test): drop the deleted harness test file from the harness check command * test: make the key rebind, team member delete and routes integration tests exercise the legacy paths * test: make fallback, rpm and spend integration contracts deterministic and clean up their rows, assert the rpm limit in usage-based routing * test: expect the no-deployments error at the rpm limit and cover the strategy check without pre-call checks * test: hand member cleanups to the scenario instead of growing a budget list, flatten callback kinds * test: scope the admin health check to the test's own deployment --------- Co-authored-by: yuneng <yuneng@berri.ai>
291 lines
13 KiB
Python
291 lines
13 KiB
Python
import threading
|
|
import time
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from typing import Final, Literal
|
|
|
|
import httpx
|
|
import pytest
|
|
from pydantic import JsonValue, TypeAdapter
|
|
|
|
from tests.integration._support.database import read_rows
|
|
from tests.integration._support.client import (
|
|
JSON_OBJECT,
|
|
Gateway,
|
|
Scenario,
|
|
eventually,
|
|
object_value,
|
|
string_value,
|
|
)
|
|
from tests.integration._support.openai_wire import chat_reply
|
|
from tests.integration._support.wire import Reply, Request, wire_server
|
|
|
|
UNREACHABLE_API_BASE: Final = "http://127.0.0.1:9/v1"
|
|
END_USER_REQUESTS: Final = 10
|
|
CONCURRENT_REQUESTS: Final = 25
|
|
DISTRIBUTION_REQUESTS: Final = 20
|
|
SLOW_UPSTREAM_SECONDS: Final = 3
|
|
HELD_REPLY_SECONDS: Final = 30
|
|
END_USER_ROW_SECONDS: Final = 70
|
|
CUSTOM_FALLBACK_TEXT: Final = "custom fallback prompt"
|
|
Caller = Literal["virtual-key", "master-key"]
|
|
JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]])
|
|
|
|
|
|
def _messages(text: str = "integration control") -> list[JsonValue]:
|
|
return [{"role": "user", "content": text}]
|
|
|
|
|
|
def _chat(
|
|
gateway: Gateway,
|
|
body: dict[str, JsonValue],
|
|
*,
|
|
key: str | None = None,
|
|
headers: dict[str, str] | None = None,
|
|
) -> httpx.Response:
|
|
return gateway.request("POST", "/v1/chat/completions", {"messages": _messages(), **body}, key=key, headers=headers)
|
|
|
|
|
|
def _body(response: httpx.Response) -> dict[str, JsonValue]:
|
|
return JSON_OBJECT.validate_json(response.content)
|
|
|
|
|
|
def _content(response: httpx.Response) -> str:
|
|
choices: Final = _body(response)["choices"]
|
|
assert isinstance(choices, list) and choices, response.text
|
|
return string_value(object_value(object_value(choices[0])["message"])["content"])
|
|
|
|
|
|
def _scripted_model(gateway: Gateway, scenario: Scenario, statuses: list[int], num_retries: int) -> str:
|
|
upstream_model: Final = f"integration-{uuid.uuid4().hex}"
|
|
script_url: Final = f"{gateway.upstream_url}/__scripts/{upstream_model}"
|
|
configured: Final = httpx.post(script_url, json={"statuses": statuses})
|
|
assert configured.status_code == 200, configured.text
|
|
scenario.cleanups.callback(httpx.delete, script_url)
|
|
return scenario.model(model=f"openai/{upstream_model}", num_retries=num_retries)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class UniqueModel:
|
|
name: str
|
|
upstream: str
|
|
deployment_id: str
|
|
|
|
|
|
def _unique_model(scenario: Scenario) -> UniqueModel:
|
|
upstream_model: Final = f"integration-{uuid.uuid4().hex}"
|
|
deployment_id: Final = f"integration-{uuid.uuid4().hex}"
|
|
model_info: Final[dict[str, JsonValue]] = {"id": deployment_id}
|
|
name: Final = scenario.model(model=f"openai/{upstream_model}", model_info=model_info)
|
|
return UniqueModel(name, upstream_model, deployment_id)
|
|
|
|
|
|
def _slow_reply(_: Request) -> Reply:
|
|
time.sleep(SLOW_UPSTREAM_SECONDS)
|
|
return chat_reply("chatcmpl-slow", "gpt-4o-mini", "late", stream=False)
|
|
|
|
|
|
def _fallback_reply(_: Request) -> Reply:
|
|
return chat_reply("chatcmpl-fallback", "gpt-4o-mini", "served by fallback", stream=False)
|
|
|
|
|
|
def _held_reply(release: threading.Event) -> Callable[[Request], Reply]:
|
|
def respond(_: Request) -> Reply:
|
|
assert release.wait(timeout=HELD_REPLY_SECONDS)
|
|
return chat_reply("chatcmpl-held", "gpt-4o-mini", "held", stream=False)
|
|
|
|
return respond
|
|
|
|
|
|
def _delete_auto_created_end_user(gateway: Gateway, user_id: str) -> None:
|
|
eventually(
|
|
lambda: read_rows('SELECT user_id FROM "LiteLLM_EndUserTable" WHERE user_id = %s', (user_id,)),
|
|
lambda rows: len(rows) == 1,
|
|
seconds=END_USER_ROW_SECONDS,
|
|
)
|
|
gateway.post("/end_user/delete", {"user_ids": [user_id]})
|
|
assert read_rows('SELECT user_id FROM "LiteLLM_EndUserTable" WHERE user_id = %s', (user_id,)) == []
|
|
|
|
|
|
def _status(gateway: Gateway, model: str) -> int:
|
|
return _chat(gateway, {"model": model}).status_code
|
|
|
|
|
|
def _served_model_id(gateway: Gateway, model: str) -> str:
|
|
response: Final = _chat(gateway, {"model": model})
|
|
assert response.status_code == 200, response.text
|
|
return response.headers["x-litellm-model-id"]
|
|
|
|
|
|
def _concurrent_round(gateway: Gateway, bad: str, good: str) -> tuple[list[int], list[int]]:
|
|
with ThreadPoolExecutor(max_workers=CONCURRENT_REQUESTS * 2) as pool:
|
|
bad_calls: Final = [pool.submit(_status, gateway, bad) for _ in range(CONCURRENT_REQUESTS)]
|
|
good_calls: Final = [pool.submit(_status, gateway, good) for _ in range(CONCURRENT_REQUESTS)]
|
|
return [call.result() for call in bad_calls], [call.result() for call in good_calls]
|
|
|
|
|
|
def _deployment(gateway: Gateway, scenario: Scenario, model_name: str, model: str = "openai/gpt-4o-mini") -> str:
|
|
created: Final = gateway.post(
|
|
"/model/new",
|
|
{
|
|
"model_name": model_name,
|
|
"litellm_params": {
|
|
"model": model,
|
|
"api_key": "integration-provider-key",
|
|
"api_base": f"{gateway.upstream_url}/v1",
|
|
},
|
|
},
|
|
)
|
|
deployment_id: Final = string_value(object_value(created["model_info"])["id"])
|
|
scenario.cleanups.callback(scenario.delete_model, deployment_id)
|
|
return deployment_id
|
|
|
|
|
|
@pytest.mark.parametrize("caller", ["virtual-key", "master-key"])
|
|
def test_end_user_budget_tpm_limit_rate_limits_their_requests(gateway: Gateway, caller: Caller) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model()
|
|
budget: Final = scenario.budget(tpm_limit=2)
|
|
end_user: Final = f"integration-{uuid.uuid4().hex}"
|
|
gateway.post("/end_user/new", {"user_id": end_user, "budget_id": budget})
|
|
scenario.cleanups.callback(gateway.post, "/end_user/delete", {"user_ids": [end_user]})
|
|
key: Final = scenario.key(models=[model]) if caller == "virtual-key" else gateway.key
|
|
control_user: Final = f"integration-{uuid.uuid4().hex}"
|
|
scenario.cleanups.callback(_delete_auto_created_end_user, gateway, control_user)
|
|
control: Final = _chat(gateway, {"model": model, "user": control_user}, key=key)
|
|
assert control.status_code == 200, control.text
|
|
statuses: Final = [
|
|
_chat(gateway, {"model": model, "user": end_user}, key=key).status_code for _ in range(END_USER_REQUESTS)
|
|
]
|
|
assert statuses.count(200) < 5, statuses
|
|
assert set(statuses) <= {200, 429}, statuses
|
|
|
|
|
|
def test_client_fallbacks_reach_an_allowed_model_and_name_a_denied_one(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
primary: Final = scenario.model(api_base=UNREACHABLE_API_BASE)
|
|
fallback: Final = _unique_model(scenario)
|
|
body: Final[dict[str, JsonValue]] = {"model": primary, "fallbacks": [fallback.name]}
|
|
served: Final = _chat(gateway, body, key=scenario.key(models=[primary, fallback.name]))
|
|
assert served.status_code == 200, served.text
|
|
assert _body(served)["model"] == fallback.upstream
|
|
assert served.headers["x-litellm-model-id"] == fallback.deployment_id
|
|
assert _content(served)
|
|
denied: Final = _chat(gateway, body, key=scenario.key(models=[primary]))
|
|
assert denied.status_code == 403, denied.text
|
|
assert fallback.name in denied.text
|
|
|
|
|
|
def test_client_fallback_with_custom_messages_sends_them_to_the_fallback(gateway: Gateway) -> None:
|
|
custom: Final = _messages(CUSTOM_FALLBACK_TEXT)
|
|
with gateway.scenario() as scenario, wire_server(_fallback_reply) as wire:
|
|
primary: Final = scenario.model(api_base=UNREACHABLE_API_BASE)
|
|
fallback: Final = scenario.model(api_base=wire.url)
|
|
body: Final[dict[str, JsonValue]] = {
|
|
"model": primary,
|
|
"fallbacks": [{"model": fallback, "messages": custom}],
|
|
}
|
|
served: Final = _chat(gateway, body, key=scenario.key(models=[primary, fallback]))
|
|
assert served.status_code == 200, served.text
|
|
assert _content(served) == "served by fallback"
|
|
forwarded: Final = [JSON_OBJECT.validate_json(request.body)["messages"] for request in wire.drain()]
|
|
assert forwarded == [custom]
|
|
denied: Final = _chat(gateway, body, key=scenario.key(models=[primary]))
|
|
assert denied.status_code == 403, denied.text
|
|
assert fallback in denied.text
|
|
assert wire.drain() == ()
|
|
|
|
|
|
def test_rate_limited_deployment_is_retried_and_reports_retry_counts(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = _scripted_model(gateway, scenario, [429, 200], 50)
|
|
response: Final = _chat(gateway, {"model": model})
|
|
assert response.status_code == 200, response.text
|
|
assert response.headers["x-litellm-attempted-retries"] == "1"
|
|
assert response.headers["x-litellm-max-retries"] == "50"
|
|
|
|
|
|
def test_request_fallbacks_reroute_after_a_connection_failure(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
primary: Final = scenario.model(api_base=UNREACHABLE_API_BASE)
|
|
fallback: Final = _unique_model(scenario)
|
|
response: Final = _chat(gateway, {"model": primary, "fallbacks": [fallback.name]})
|
|
assert response.status_code == 200, response.text
|
|
assert _body(response)["model"] == fallback.upstream
|
|
assert response.headers["x-litellm-model-id"] == fallback.deployment_id
|
|
assert response.headers["x-litellm-attempted-fallbacks"] == "1"
|
|
|
|
|
|
def test_model_level_timeout_is_reported_on_a_timed_out_request(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario, wire_server(_slow_reply) as wire:
|
|
response: Final = _chat(gateway, {"model": scenario.model(api_base=wire.url, timeout=1)})
|
|
assert response.status_code == 408, response.text
|
|
assert response.headers["x-litellm-timeout"] == "1.0"
|
|
|
|
|
|
def test_request_timeout_header_overrides_the_model_timeout(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario, wire_server(_slow_reply) as wire:
|
|
response: Final = _chat(
|
|
gateway,
|
|
{"model": scenario.model(api_base=wire.url, timeout=1)},
|
|
headers={"x-litellm-timeout": "0.001"},
|
|
)
|
|
assert response.status_code == 408, response.text
|
|
assert response.headers["x-litellm-timeout"] == "0.001"
|
|
|
|
|
|
def test_failing_model_traffic_does_not_starve_concurrent_good_requests(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
bad: Final = scenario.model(api_base=UNREACHABLE_API_BASE)
|
|
good: Final = scenario.model()
|
|
for _ in range(2):
|
|
bad_calls, good_calls = _concurrent_round(gateway, bad, good)
|
|
assert good_calls == [200] * CONCURRENT_REQUESTS
|
|
assert 200 not in bad_calls
|
|
|
|
|
|
def test_rpm_limited_deployment_rejects_a_second_call_while_the_first_is_in_flight(gateway: Gateway) -> None:
|
|
release: Final = threading.Event()
|
|
with gateway.scenario() as scenario, wire_server(_held_reply(release)) as wire, ThreadPoolExecutor(1) as pool:
|
|
model: Final = scenario.model(api_base=wire.url, rpm=1)
|
|
first: Final = pool.submit(_chat, gateway, {"model": model})
|
|
try:
|
|
eventually(wire.received.qsize, lambda count: count == 1)
|
|
second: Final = _chat(gateway, {"model": model})
|
|
assert second.status_code == 429, second.text
|
|
assert wire.received.qsize() == 1
|
|
finally:
|
|
release.set()
|
|
assert first.result().status_code == 200
|
|
assert len(wire.drain()) == 1
|
|
|
|
|
|
def test_model_group_with_two_deployments_serves_from_both(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
first_id: Final = f"integration-{uuid.uuid4().hex}"
|
|
model: Final = scenario.model(model_info={"id": first_id})
|
|
second_id: Final = _deployment(gateway, scenario, model)
|
|
served_by: Final = {_served_model_id(gateway, model) for _ in range(DISTRIBUTION_REQUESTS)}
|
|
assert served_by == {first_id, second_id}
|
|
|
|
|
|
def test_unlisted_provider_model_resolves_through_a_wildcard_deployment(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
prefix: Final = f"integration{uuid.uuid4().hex}"
|
|
_deployment(gateway, scenario, f"{prefix}/*", "openai/*")
|
|
response: Final = _chat(gateway, {"model": f"{prefix}/gpt-4o-mini"})
|
|
assert response.status_code == 200, response.text
|
|
assert _content(response)
|
|
|
|
|
|
def test_comma_separated_models_fan_out_to_one_response_per_model(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
first: Final = _unique_model(scenario)
|
|
second: Final = _unique_model(scenario)
|
|
response: Final = _chat(gateway, {"model": f"{first.name},{second.name}"})
|
|
assert response.status_code == 200, response.text
|
|
replies: Final = JSON_OBJECTS.validate_json(response.content)
|
|
assert len(replies) == 2, replies
|
|
assert {string_value(reply["model"]) for reply in replies} == {first.upstream, second.upstream}
|