mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
367 lines
17 KiB
Python
367 lines
17 KiB
Python
import uuid
|
|
from contextlib import ExitStack
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
import yaml
|
|
from integration._support.client import Gateway, eventually, object_value
|
|
from integration._support.database import read_rows
|
|
from integration._support.process import owned_proxy
|
|
from pydantic import BaseModel, JsonValue
|
|
|
|
|
|
class _BlockedCustomerInfo(BaseModel):
|
|
user_id: str
|
|
blocked: bool
|
|
|
|
|
|
class _BlockUsersResponse(BaseModel):
|
|
blocked_users: list[_BlockedCustomerInfo]
|
|
|
|
|
|
def _blocked_row(user_id: str) -> list[dict[str, JsonValue]]:
|
|
return read_rows('SELECT blocked FROM "LiteLLM_EndUserTable" WHERE user_id=%s', (user_id,))
|
|
|
|
|
|
def _observed(upstream: httpx.Client) -> list[dict[str, JsonValue]]:
|
|
response: Final = upstream.get("/__observations")
|
|
response.raise_for_status()
|
|
requests: Final = response.json()["requests"]
|
|
assert isinstance(requests, list)
|
|
return requests
|
|
|
|
|
|
def _chat(gateway: Gateway, model: str, key: str, end_user: str) -> httpx.Response:
|
|
return gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": f"end user block {uuid.uuid4().hex}"}],
|
|
"user": end_user,
|
|
},
|
|
key=key,
|
|
)
|
|
|
|
|
|
def _observation(text: str, end_user: str) -> dict[str, JsonValue]:
|
|
return {
|
|
"path": "/v1/chat/completions",
|
|
"authorization": "Bearer integration-provider-key",
|
|
"body": {
|
|
"messages": [{"role": "user", "content": text}],
|
|
"model": "gpt-4o-mini",
|
|
"user": end_user,
|
|
},
|
|
"method": "POST",
|
|
"api_key": "",
|
|
}
|
|
|
|
|
|
def _assert_blocked_refusal(response: httpx.Response, end_user: str) -> None:
|
|
assert response.status_code == 400, response.text
|
|
assert response.json() == {
|
|
"error": {
|
|
"message": f"User blocked from making LLM API Calls. User={end_user}",
|
|
"type": "invalid_request_error",
|
|
"param": None,
|
|
"code": "400",
|
|
"provider_specific_fields": {"error": f"User blocked from making LLM API Calls. User={end_user}"},
|
|
}
|
|
}, response.text
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _CallbackRig:
|
|
first: Gateway
|
|
second: Gateway
|
|
|
|
|
|
def _callback_config(directory: Path, blocked_user_list: list[str] | None = None) -> Path:
|
|
config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
|
settings: Final = {
|
|
**object_value(config["litellm_settings"]),
|
|
"callbacks": ["blocked_user_check"],
|
|
}
|
|
if blocked_user_list is not None:
|
|
settings["blocked_user_list"] = blocked_user_list
|
|
config["litellm_settings"] = settings
|
|
path: Final = directory / "config.yaml"
|
|
path.write_text(yaml.safe_dump(config))
|
|
return path
|
|
|
|
|
|
def _two_callback_proxies(gateway: Gateway, tmp_path: Path, stack: ExitStack) -> _CallbackRig:
|
|
config: Final = _callback_config(tmp_path)
|
|
first: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "a", {}, config=config))
|
|
second: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "b", {}, config=config))
|
|
return _CallbackRig(first, second)
|
|
|
|
|
|
def _one_callback_proxy(gateway: Gateway, tmp_path: Path, stack: ExitStack) -> Gateway:
|
|
config: Final = _callback_config(tmp_path)
|
|
return stack.enter_context(owned_proxy(gateway, tmp_path / "a", {}, config=config))
|
|
|
|
|
|
@pytest.mark.parametrize("block_route", ("/customer/block", "/end_user/block"))
|
|
def test_block_route_marks_the_customer_blocked_and_update_unblocks(gateway: Gateway, block_route: str) -> None:
|
|
with gateway.scenario() as scenario:
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
|
|
gateway.request("POST", block_route, {"user_ids": [end_user]})
|
|
assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15) == [
|
|
{"blocked": True}
|
|
]
|
|
info: Final = gateway.get("/customer/info", {"end_user_id": end_user})
|
|
assert _BlockedCustomerInfo.model_validate(info).blocked is True
|
|
|
|
updated: Final = gateway.request("POST", "/customer/update", {"user_id": end_user, "blocked": False})
|
|
assert updated.status_code == 200, updated.text
|
|
assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15) == [
|
|
{"blocked": False}
|
|
]
|
|
assert (
|
|
_BlockedCustomerInfo.model_validate(gateway.get("/customer/info", {"end_user_id": end_user})).blocked
|
|
is False
|
|
)
|
|
|
|
|
|
def _assert_unseen_customer_is_blocked(gateway: Gateway, block_route: str) -> None:
|
|
with gateway.scenario() as scenario:
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
|
|
blocked: Final = gateway.request("POST", block_route, {"user_ids": [end_user]})
|
|
assert blocked.status_code == 200, blocked.text
|
|
blocked_response: Final = _BlockUsersResponse.model_validate_json(blocked.content)
|
|
assert blocked_response == _BlockUsersResponse(
|
|
blocked_users=[_BlockedCustomerInfo(user_id=end_user, blocked=True)]
|
|
)
|
|
assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15) == [
|
|
{"blocked": True}
|
|
]
|
|
info: Final = gateway.request("GET", "/customer/info", params={"end_user_id": end_user})
|
|
assert info.status_code == 200, info.text
|
|
assert _BlockedCustomerInfo.model_validate(info.json()).blocked is True
|
|
|
|
updated: Final = gateway.request("POST", "/customer/update", {"user_id": end_user, "blocked": False})
|
|
assert updated.status_code == 200, updated.text
|
|
assert eventually(lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15) == [
|
|
{"blocked": False}
|
|
]
|
|
info_after_update: Final = gateway.request("GET", "/customer/info", params={"end_user_id": end_user})
|
|
assert info_after_update.status_code == 200, info_after_update.text
|
|
assert _BlockedCustomerInfo.model_validate(info_after_update.json()).blocked is False
|
|
|
|
|
|
def test_end_user_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway) -> None:
|
|
_assert_unseen_customer_is_blocked(gateway, "/end_user/block")
|
|
|
|
|
|
def test_customer_block_route_creates_and_blocks_an_unseen_customer(gateway: Gateway) -> None:
|
|
pytest.skip("BUG: POST /customer/block for an unseen customer returns 500 instead of 200 BlockUsersResponse")
|
|
_assert_unseen_customer_is_blocked(gateway, "/customer/block")
|
|
|
|
|
|
def test_listed_end_user_is_refused_and_unblock_restores_serving(gateway: Gateway, tmp_path: Path) -> None:
|
|
with (
|
|
ExitStack() as stack,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
config: Final = _callback_config(tmp_path, blocked_user_list=[end_user])
|
|
proxy: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "listed", {}, config=config))
|
|
with proxy.scenario() as scenario:
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
assert _observed(upstream) == []
|
|
|
|
refused: Final = _chat(proxy, model, key, end_user)
|
|
_assert_blocked_refusal(refused, end_user)
|
|
assert _observed(upstream) == []
|
|
|
|
unblocked: Final = proxy.request("POST", "/customer/unblock", {"user_ids": [end_user]})
|
|
assert unblocked.status_code == 200, unblocked.text
|
|
assert unblocked.json() == {"blocked_users": []}, unblocked.text
|
|
assert _blocked_row(end_user) == []
|
|
|
|
text: Final = f"end user unblocked {uuid.uuid4().hex}"
|
|
served: Final = proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": text}], "user": end_user},
|
|
key=key,
|
|
)
|
|
assert served.status_code == 200, served.text
|
|
assert _observed(upstream) == [_observation(text, end_user)]
|
|
|
|
|
|
@pytest.mark.parametrize("block_route", ("/customer/block", "/end_user/block"))
|
|
def test_blocked_end_user_is_refused_on_every_proxy_and_unblock_restores(
|
|
gateway: Gateway, tmp_path: Path, block_route: str
|
|
) -> None:
|
|
pytest.skip("BUG: blocked end-user chat requests return HTTP 200 instead of HTTP 400")
|
|
with (
|
|
ExitStack() as stack,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
rig: Final = _two_callback_proxies(gateway, tmp_path, stack)
|
|
with rig.first.scenario() as scenario:
|
|
_observed(upstream)
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
created: Final = rig.first.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(rig.first.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
served: Final = _chat(rig.first, model, key, end_user)
|
|
assert served.status_code == 200, served.text
|
|
_observed(upstream)
|
|
|
|
rig.first.request("POST", block_route, {"user_ids": [end_user]})
|
|
assert eventually(
|
|
lambda: _blocked_row(end_user), lambda value: value == [{"blocked": True}], seconds=15
|
|
) == [{"blocked": True}]
|
|
info: Final = rig.first.get("/customer/info", {"end_user_id": end_user})
|
|
assert _BlockedCustomerInfo.model_validate(info).blocked is True
|
|
|
|
denied_first: Final = _chat(rig.first, model, key, end_user)
|
|
denied_second: Final = _chat(rig.second, model, key, end_user)
|
|
for denied in (denied_first, denied_second):
|
|
_assert_blocked_refusal(denied, end_user)
|
|
upstream_calls: Final = _observed(upstream)
|
|
assert upstream_calls == [], upstream_calls
|
|
|
|
updated: Final = rig.first.request("POST", "/customer/update", {"user_id": end_user, "blocked": False})
|
|
assert updated.status_code == 200, updated.text
|
|
assert eventually(
|
|
lambda: _blocked_row(end_user), lambda value: value == [{"blocked": False}], seconds=15
|
|
) == [{"blocked": False}]
|
|
assert (
|
|
_BlockedCustomerInfo.model_validate(rig.first.get("/customer/info", {"end_user_id": end_user})).blocked
|
|
is False
|
|
)
|
|
|
|
text: Final = f"end user unblocked {uuid.uuid4().hex}"
|
|
|
|
def call(candidate: Gateway) -> httpx.Response:
|
|
return candidate.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": text}], "user": end_user},
|
|
key=key,
|
|
)
|
|
|
|
assert (
|
|
eventually(
|
|
lambda: call(rig.first), lambda response: response.status_code == 200, seconds=10
|
|
).status_code
|
|
== 200
|
|
)
|
|
assert (
|
|
eventually(
|
|
lambda: call(rig.second), lambda response: response.status_code == 200, seconds=10
|
|
).status_code
|
|
== 200
|
|
)
|
|
expected_observations: Final = [
|
|
{
|
|
"path": "/v1/chat/completions",
|
|
"authorization": "Bearer integration-provider-key",
|
|
"body": {
|
|
"messages": [{"role": "user", "content": text}],
|
|
"model": "gpt-4o-mini",
|
|
"user": end_user,
|
|
},
|
|
"method": "POST",
|
|
"api_key": "",
|
|
}
|
|
] * 2
|
|
assert _observed(upstream) == expected_observations
|
|
|
|
|
|
def test_unblock_route_requires_the_blocked_user_check_callback(gateway: Gateway, tmp_path: Path) -> None:
|
|
with ExitStack() as stack:
|
|
callback_proxy: Final = _one_callback_proxy(gateway, tmp_path, stack)
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
created: Final = callback_proxy.request(
|
|
"POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0, "blocked": True}
|
|
)
|
|
assert created.status_code == 200, created.text
|
|
stack.callback(callback_proxy.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
assert _blocked_row(end_user) == [{"blocked": True}]
|
|
|
|
shared: Final = gateway.request("POST", "/customer/unblock", {"user_ids": [end_user]})
|
|
assert shared.status_code == 400, shared.text
|
|
assert shared.json() == {"detail": {"error": "Blocked user check was never set. This call has no effect."}}, (
|
|
shared.text
|
|
)
|
|
assert _blocked_row(end_user) == [{"blocked": True}]
|
|
|
|
callback_response: Final = callback_proxy.request("POST", "/customer/unblock", {"user_ids": [end_user]})
|
|
assert callback_response.status_code == 400, callback_response.text
|
|
assert callback_response.json() == {
|
|
"detail": {"error": "Blocked user check was never set. This call has no effect."}
|
|
}, callback_response.text
|
|
assert _blocked_row(end_user) == [{"blocked": True}]
|
|
|
|
|
|
def test_block_route_returns_the_blocked_records(gateway: Gateway) -> None:
|
|
pytest.skip("BUG: POST /customer/block returns 500 instead of 200 BlockUsersResponse")
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
with gateway.scenario() as scenario:
|
|
created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
blocked: Final = gateway.request("POST", "/customer/block", {"user_ids": [end_user]})
|
|
assert blocked.status_code == 200, blocked.text
|
|
payload: Final = object_value(blocked.json())
|
|
blocked_users: Final = payload["blocked_users"]
|
|
assert isinstance(blocked_users, list) and len(blocked_users) == 1
|
|
assert object_value(blocked_users[0])["user_id"] == end_user
|
|
assert object_value(blocked_users[0])["blocked"] is True
|
|
|
|
|
|
def test_blocked_end_user_is_refused_without_the_callback(gateway: Gateway) -> None:
|
|
pytest.skip("BUG: blocked end-user request without callback returns 200 instead of 400")
|
|
with gateway.scenario() as scenario:
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
created: Final = gateway.request("POST", "/customer/new", {"user_id": end_user, "blocked": True})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(gateway.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
refused: Final = _chat(gateway, model, key, end_user)
|
|
_assert_blocked_refusal(refused, end_user)
|
|
|
|
|
|
def test_blocked_end_user_header_spelling_is_refused_with_the_callback(gateway: Gateway, tmp_path: Path) -> None:
|
|
pytest.skip("BUG: x-litellm-end-user-id request for a blocked customer returns 200 instead of 400")
|
|
end_user: Final = f"integration-end-user-{uuid.uuid4().hex}"
|
|
config: Final = _callback_config(tmp_path, blocked_user_list=[end_user])
|
|
with (
|
|
ExitStack() as stack,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
proxy: Final = stack.enter_context(owned_proxy(gateway, tmp_path / "listed", {}, config=config))
|
|
with proxy.scenario() as scenario:
|
|
created: Final = proxy.request("POST", "/customer/new", {"user_id": end_user, "max_budget": 100.0})
|
|
assert created.status_code == 200, created.text
|
|
scenario.cleanups.callback(proxy.request, "POST", "/customer/delete", {"user_ids": [end_user]})
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
refused: Final = proxy.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": f"blocked {uuid.uuid4().hex}"}]},
|
|
key=key,
|
|
headers={"x-litellm-end-user-id": end_user},
|
|
)
|
|
_assert_blocked_refusal(refused, end_user)
|
|
assert _observed(upstream) == []
|