mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
44d99475a5
* test(integration): stop GET /__observations from draining the upstream log GET /__observations used to empty the queue, so a second read in a test saw nothing and a "never reached upstream" check could pass whether or not the request was forwarded. GET now returns every request since the last reset and DELETE /__observations is the reset. An autouse fixture resets the worker upstream before each test, which also covers tests that read without the gateway fixture. Discarded reads used as resets become DELETEs, helpers that accumulated drained reads now read one snapshot, and tests with several phases reset between them. The three vacuous negative checks (rag allowlist, guardrail opt-out chaos, xAI batch output) now assert the absence on the same read that shows a request which did reach the upstream. * style(integration): space the observations unit test like its neighbours
1026 lines
49 KiB
Python
1026 lines
49 KiB
Python
import os
|
|
import uuid
|
|
from collections.abc import Callable, Iterator
|
|
from datetime import datetime, timedelta, timezone
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
|
from integration._support.database import read_rows
|
|
from pydantic import BaseModel, JsonValue
|
|
from redis import Redis
|
|
|
|
|
|
def _python_isoformat(value: datetime) -> str:
|
|
return value.isoformat()
|
|
|
|
|
|
def _javascript_iso_string(value: datetime) -> str:
|
|
return value.strftime("%Y-%m-%dT%H:%M:%S.000Z")
|
|
|
|
|
|
class _ObservedRequest(BaseModel):
|
|
body: dict[str, JsonValue]
|
|
|
|
|
|
class _Observations(BaseModel):
|
|
requests: list[_ObservedRequest]
|
|
|
|
|
|
class _ProxyError(BaseModel):
|
|
type: str
|
|
|
|
|
|
class _ErrorResponse(BaseModel):
|
|
error: _ProxyError
|
|
|
|
|
|
class _BudgetRecord(BaseModel):
|
|
budget_id: str
|
|
max_budget: float | None = None
|
|
tpm_limit: int | None = None
|
|
rpm_limit: int | None = None
|
|
budget_duration: str | None = None
|
|
allowed_models: list[str] | None = None
|
|
temp_budget_increase: float | None = None
|
|
temp_budget_expiry: datetime | None = None
|
|
|
|
|
|
class _MemberBudgetState(BaseModel):
|
|
user_id: str
|
|
team_id: str
|
|
budget_id: str | None = None
|
|
spend: float | None = None
|
|
max_budget: float | None = None
|
|
tpm_limit: int | None = None
|
|
rpm_limit: int | None = None
|
|
budget_duration: str | None = None
|
|
allowed_models: list[str] | None = None
|
|
temp_budget_increase: float | None = None
|
|
temp_budget_expiry: datetime | None = None
|
|
|
|
|
|
class _TeamRosterMember(BaseModel):
|
|
role: str
|
|
user_id: str | None = None
|
|
user_email: str | None = None
|
|
|
|
|
|
class _TeamState(BaseModel):
|
|
team_id: str
|
|
models: list[str] | None = None
|
|
members_with_roles: tuple[_TeamRosterMember, ...] = ()
|
|
metadata: dict[str, JsonValue] | None = None
|
|
|
|
|
|
class _TeamMembershipInfo(BaseModel):
|
|
user_id: str
|
|
team_id: str
|
|
budget_id: str | None = None
|
|
spend: float | None = None
|
|
budget_source: str | None = None
|
|
litellm_budget_table: _BudgetRecord | None = None
|
|
|
|
|
|
class _TeamInfo(BaseModel):
|
|
team_info: _TeamState
|
|
team_memberships: list[_TeamMembershipInfo]
|
|
|
|
|
|
class _MemberMe(BaseModel):
|
|
user_id: str
|
|
team_id: str
|
|
budget_id: str | None = None
|
|
role: str | None = None
|
|
litellm_budget_table: _BudgetRecord | None = None
|
|
|
|
|
|
class _MemberUpdateResponse(BaseModel):
|
|
team_id: str
|
|
user_id: str
|
|
user_email: str | None = None
|
|
max_budget_in_team: float | None = None
|
|
tpm_limit: int | None = None
|
|
rpm_limit: int | None = None
|
|
budget_duration: str | None = None
|
|
allowed_models: list[str] | None = None
|
|
temp_budget_increase: float | None = None
|
|
temp_budget_expiry: datetime | None = None
|
|
|
|
|
|
class _BulkBudgetResult(BaseModel):
|
|
user_id: str | None = None
|
|
user_email: str | None = None
|
|
success: bool
|
|
error: str | None = None
|
|
budget_id: str | None = None
|
|
max_budget: float | None = None
|
|
max_budget_source: str | None = None
|
|
tpm_limit: int | None = None
|
|
rpm_limit: int | None = None
|
|
budget_duration: str | None = None
|
|
allowed_models: tuple[str, ...] | None = None
|
|
|
|
|
|
class _BulkBudgetUpdateResponse(BaseModel):
|
|
data: tuple[_BulkBudgetResult, ...]
|
|
|
|
|
|
@pytest.fixture
|
|
def upstream(gateway: Gateway) -> Iterator[httpx.Client]:
|
|
with httpx.Client(base_url=gateway.upstream_url, timeout=15, trust_env=False) as client:
|
|
yield client
|
|
|
|
|
|
def _priced_model(scenario: Scenario) -> tuple[str, str]:
|
|
provider_model: Final = f"team-member-spend-{uuid.uuid4().hex}"
|
|
model: Final = scenario.model(
|
|
model=f"openai/{provider_model}",
|
|
input_cost_per_token=0.001,
|
|
output_cost_per_token=0.002,
|
|
)
|
|
return model, provider_model
|
|
|
|
|
|
def _free_model(scenario: Scenario) -> tuple[str, str]:
|
|
provider_model: Final = f"team-member-free-{uuid.uuid4().hex}"
|
|
model: Final = scenario.model(
|
|
model=f"openai/{provider_model}",
|
|
input_cost_per_token=0,
|
|
output_cost_per_token=0,
|
|
)
|
|
return model, provider_model
|
|
|
|
|
|
def _observed_models(upstream: httpx.Client) -> tuple[str, ...]:
|
|
response: Final = upstream.get("/__observations")
|
|
assert response.status_code == 200, response.text
|
|
observations: Final = _Observations.model_validate_json(response.text)
|
|
return tuple(string_value(request.body["model"]) for request in observations.requests)
|
|
|
|
|
|
def _chat(gateway: Gateway, model: str, key: str, text: str) -> httpx.Response:
|
|
return gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": text}]},
|
|
key=key,
|
|
)
|
|
|
|
|
|
def _team_info(gateway: Gateway, team_id: str) -> _TeamInfo:
|
|
response: Final = gateway.request("GET", "/team/info", params={"team_id": team_id})
|
|
assert response.status_code == 200, response.text
|
|
return _TeamInfo.model_validate_json(response.text)
|
|
|
|
|
|
def _member_me(gateway: Gateway, team_id: str, key: str) -> _MemberMe:
|
|
response: Final = gateway.request("GET", f"/team/{team_id}/members/me", key=key)
|
|
assert response.status_code == 200, response.text
|
|
return _MemberMe.model_validate_json(response.text)
|
|
|
|
|
|
def _budget_state(team_id: str, user_id: str) -> tuple[_MemberBudgetState, ...]:
|
|
return tuple(
|
|
_MemberBudgetState.model_validate(row)
|
|
for row in read_rows(
|
|
"SELECT tm.user_id, tm.team_id, tm.budget_id, tm.spend, b.max_budget, b.tpm_limit, b.rpm_limit, "
|
|
"b.budget_duration, b.allowed_models, b.temp_budget_increase, "
|
|
"b.temp_budget_expiry::text AS temp_budget_expiry "
|
|
'FROM "LiteLLM_TeamMembership" tm LEFT JOIN "LiteLLM_BudgetTable" b ON b.budget_id = tm.budget_id '
|
|
"WHERE tm.team_id = %s AND tm.user_id = %s",
|
|
(team_id, user_id),
|
|
)
|
|
)
|
|
|
|
|
|
def _budget_record(budget_id: str) -> tuple[_BudgetRecord, ...]:
|
|
return tuple(
|
|
_BudgetRecord.model_validate(row)
|
|
for row in read_rows(
|
|
"SELECT budget_id, max_budget, tpm_limit, rpm_limit, budget_duration, allowed_models, "
|
|
'temp_budget_increase, temp_budget_expiry FROM "LiteLLM_BudgetTable" WHERE budget_id = %s',
|
|
(budget_id,),
|
|
)
|
|
)
|
|
|
|
|
|
def _team_default_budget_id(team_id: str) -> str:
|
|
rows: Final = read_rows('SELECT metadata FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
|
|
assert len(rows) == 1, rows
|
|
metadata: Final = object_value(rows[0]["metadata"])
|
|
return string_value(metadata["team_member_budget_id"])
|
|
|
|
|
|
def _member_update(gateway: Gateway, body: dict[str, JsonValue]) -> _MemberUpdateResponse:
|
|
response: Final = gateway.request("POST", "/team/member_update", body)
|
|
assert response.status_code == 200, response.text
|
|
return _MemberUpdateResponse.model_validate_json(response.text)
|
|
|
|
|
|
def _budget_refused(response: httpx.Response) -> None:
|
|
assert response.status_code == 422, response.text
|
|
assert _ErrorResponse.model_validate_json(response.text).error.type == "budget_exceeded", response.text
|
|
|
|
|
|
@pytest.mark.covers("spend.team_member.member_without_budget_gets_membership_row_and_spend")
|
|
def test_member_added_without_any_budget_is_charged_on_its_membership_row(gateway: Gateway) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
team: Final = scenario.team(models=[model])
|
|
user: Final = scenario.user()
|
|
added: Final = gateway.request(
|
|
"POST", "/team/member_add", {"team_id": team, "member": {"user_id": user, "role": "user"}}
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
memberships: Final = added.json()["updated_team_memberships"]
|
|
assert [
|
|
{"user_id": row["user_id"], "team_id": row["team_id"], "budget_id": row["budget_id"], "spend": row["spend"]}
|
|
for row in memberships
|
|
] == [{"user_id": user, "team_id": team, "budget_id": None, "spend": 0}], added.text
|
|
assert read_rows(
|
|
'SELECT budget_id, spend, total_spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s',
|
|
(team, user),
|
|
) == [{"budget_id": None, "spend": 0.0, "total_spend": 0.0}]
|
|
key: Final = scenario.key(team_id=team, user_id=user, models=[model])
|
|
assert gateway.chat(model, key=key, text=f"member spend {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
|
|
charged: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT spend, total_spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s',
|
|
(team, user),
|
|
),
|
|
lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06,
|
|
seconds=70,
|
|
)
|
|
assert float(charged[0]["spend"]) == pytest.approx(0.06)
|
|
assert float(charged[0]["total_spend"]) == pytest.approx(0.06)
|
|
team_rows: Final = eventually(
|
|
lambda: read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id=%s', (team,)),
|
|
lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06,
|
|
seconds=70,
|
|
)
|
|
assert float(team_rows[0]["spend"]) == pytest.approx(0.06)
|
|
info: Final = gateway.get("/team/info", {"team_id": team})
|
|
listed: Final = info["team_memberships"]
|
|
assert isinstance(listed, list)
|
|
exposed: Final = [
|
|
(object_value(row)["user_id"], object_value(row)["spend"])
|
|
for row in listed
|
|
if object_value(row)["user_id"] == user
|
|
]
|
|
assert len(exposed) == 1 and exposed[0][1] == pytest.approx(0.06), info
|
|
|
|
|
|
@pytest.mark.covers("spend.team_member.stale_low_redis_counter_still_blocks_member_over_budget")
|
|
def test_member_over_budget_is_blocked_when_redis_counter_reads_stale_low(gateway: Gateway) -> None:
|
|
with (
|
|
gateway.scenario() as scenario,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"])) as cache,
|
|
):
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
team: Final = scenario.team(models=[model])
|
|
user: Final = scenario.user()
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{"team_id": team, "member": {"user_id": user, "role": "user"}, "max_budget_in_team": 0.05},
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
key: Final = scenario.key(team_id=team, user_id=user, models=[model])
|
|
assert gateway.chat(model, key=key, text=f"member budget {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
|
|
charged: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT spend FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s', (team, user)
|
|
),
|
|
lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06,
|
|
seconds=70,
|
|
)
|
|
assert float(charged[0]["spend"]) == pytest.approx(0.06)
|
|
counter_key: Final = f"spend:team_member:{user}:{team}"
|
|
counted: Final = eventually(lambda: cache.get(counter_key), lambda value: value is not None, seconds=10)
|
|
assert float(counted) == pytest.approx(0.06), counted
|
|
cache.set(counter_key, "0.01")
|
|
upstream.delete("/__observations").raise_for_status()
|
|
denied: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": f"stale counter {uuid.uuid4().hex}"}]},
|
|
key=key,
|
|
)
|
|
assert denied.status_code == 422 and denied.json()["error"]["type"] == "budget_exceeded", denied.text
|
|
assert upstream.get("/__observations").json()["requests"] == []
|
|
assert float(cache.get(counter_key)) == pytest.approx(0.06), denied.text
|
|
|
|
|
|
def test_member_update_lowering_the_budget_below_spend_blocks_the_next_call(
|
|
gateway: Gateway, peer: Gateway, upstream: httpx.Client
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model, provider_model = _priced_model(scenario)
|
|
team: Final = scenario.team(models=[model])
|
|
user: Final = scenario.user()
|
|
added: Final = gateway.request(
|
|
"POST", "/team/member_add", {"team_id": team, "member": {"user_id": user, "role": "user"}}
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
initial_budget: Final = _member_update(
|
|
gateway,
|
|
{"team_id": team, "user_id": user, "max_budget_in_team": 1},
|
|
)
|
|
assert initial_budget == _MemberUpdateResponse(team_id=team, user_id=user, max_budget_in_team=1), initial_budget
|
|
key: Final = scenario.key(team_id=team, user_id=user)
|
|
personal_key: Final = scenario.key(user_id=user)
|
|
|
|
first: Final = _chat(gateway, model, key, f"accrued member spend {uuid.uuid4().hex}")
|
|
assert first.status_code == 200, first.text
|
|
eventually(
|
|
lambda: _budget_state(team, user),
|
|
lambda values: len(values) == 1 and values[0].spend is not None and values[0].spend >= 0.06,
|
|
seconds=70,
|
|
)
|
|
peer_first: Final = _chat(peer, model, key, f"accrued member spend on peer {uuid.uuid4().hex}")
|
|
assert peer_first.status_code == 200, peer_first.text
|
|
observed_before_update: Final = _observed_models(upstream)
|
|
assert observed_before_update == (provider_model, provider_model), repr(observed_before_update)
|
|
upstream.delete("/__observations").raise_for_status()
|
|
accrued: Final = eventually(
|
|
lambda: _budget_state(team, user),
|
|
lambda values: len(values) == 1 and values[0].spend is not None and values[0].spend >= 0.12,
|
|
seconds=70,
|
|
)
|
|
assert accrued[0].spend == pytest.approx(0.12), accrued
|
|
cached_membership: Final = _member_me(gateway, team, key)
|
|
assert cached_membership.litellm_budget_table is not None, repr(cached_membership)
|
|
assert cached_membership.litellm_budget_table.max_budget == 1, repr(cached_membership)
|
|
|
|
updated: Final = _member_update(
|
|
gateway,
|
|
{
|
|
"team_id": team,
|
|
"user_id": user,
|
|
"max_budget_in_team": 0.0001,
|
|
"budget_duration": "30d",
|
|
"tpm_limit": 1000,
|
|
"rpm_limit": 100,
|
|
},
|
|
)
|
|
assert updated == _MemberUpdateResponse(
|
|
team_id=team,
|
|
user_id=user,
|
|
max_budget_in_team=0.0001,
|
|
budget_duration="30d",
|
|
tpm_limit=1000,
|
|
rpm_limit=100,
|
|
), repr(updated)
|
|
denied: Final = _chat(gateway, model, key, f"lowered member budget {uuid.uuid4().hex}")
|
|
_budget_refused(denied)
|
|
peer_denied: Final = _chat(peer, model, key, f"lowered member budget on peer {uuid.uuid4().hex}")
|
|
_budget_refused(peer_denied)
|
|
observed_after_denial: Final = _observed_models(upstream)
|
|
assert observed_after_denial == (), repr(observed_after_denial)
|
|
assert observed_before_update + observed_after_denial == (provider_model, provider_model), observed_after_denial
|
|
|
|
state: Final = _budget_state(team, user)
|
|
assert len(state) == 1, state
|
|
assert state[0].user_id == user, state
|
|
assert state[0].team_id == team, state
|
|
assert state[0].spend == pytest.approx(0.12), state
|
|
assert state[0].max_budget == 0.0001, state
|
|
assert state[0].tpm_limit == 1000, state
|
|
assert state[0].rpm_limit == 100, state
|
|
assert state[0].budget_duration == "30d", state
|
|
team_info: Final = _team_info(gateway, team)
|
|
listed: Final = next(member for member in team_info.team_memberships if member.user_id == user)
|
|
assert listed.budget_id == state[0].budget_id, repr(listed)
|
|
assert listed.litellm_budget_table is not None, repr(listed)
|
|
assert listed.litellm_budget_table.max_budget == 0.0001, repr(listed)
|
|
assert listed.litellm_budget_table.tpm_limit == 1000, repr(listed)
|
|
assert listed.litellm_budget_table.rpm_limit == 100, repr(listed)
|
|
assert listed.litellm_budget_table.budget_duration == "30d", repr(listed)
|
|
member_info: Final = _member_me(gateway, team, personal_key)
|
|
assert member_info.budget_id == state[0].budget_id, repr(member_info)
|
|
assert member_info.litellm_budget_table is not None, repr(member_info)
|
|
assert member_info.litellm_budget_table.max_budget == 0.0001, repr(member_info)
|
|
assert member_info.litellm_budget_table.tpm_limit == 1000, repr(member_info)
|
|
assert member_info.litellm_budget_table.rpm_limit == 100, repr(member_info)
|
|
assert member_info.litellm_budget_table.budget_duration == "30d", repr(member_info)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"spell_expiry",
|
|
[
|
|
pytest.param(_python_isoformat, id="python-isoformat"),
|
|
pytest.param(_javascript_iso_string, id="javascript-toisostring"),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"budget_source",
|
|
[
|
|
pytest.param("own-budget", id="own-budget"),
|
|
pytest.param("team-default", id="team-default"),
|
|
],
|
|
)
|
|
def test_member_update_temp_budget_increase_readmits_a_blocked_member(
|
|
gateway: Gateway,
|
|
upstream: httpx.Client,
|
|
spell_expiry: Callable[[datetime], str],
|
|
budget_source: str,
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model, provider_model = _priced_model(scenario)
|
|
team: Final = (
|
|
scenario.team(models=[model], team_member_budget=0.05)
|
|
if budget_source == "team-default"
|
|
else scenario.team(models=[model])
|
|
)
|
|
user: Final = scenario.user()
|
|
member_add_body: Final = {"team_id": team, "member": {"user_id": user, "role": "user"}}
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
member_add_body | ({"max_budget_in_team": 0.05} if budget_source == "own-budget" else {}),
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
default_budget_id: Final = _team_default_budget_id(team) if budget_source == "team-default" else None
|
|
default_budget_before: Final = (
|
|
_budget_record(default_budget_id) if default_budget_id is not None else ()
|
|
)
|
|
if default_budget_id is not None:
|
|
assert len(default_budget_before) == 1, repr(default_budget_before)
|
|
assert default_budget_before[0].max_budget == 0.05, default_budget_before
|
|
assert default_budget_before[0].temp_budget_increase is None, default_budget_before
|
|
assert default_budget_before[0].temp_budget_expiry is None, default_budget_before
|
|
initial_state: Final = _budget_state(team, user)
|
|
assert len(initial_state) == 1, initial_state
|
|
if budget_source == "own-budget":
|
|
assert initial_state[0].max_budget == 0.05, initial_state
|
|
else:
|
|
assert default_budget_id is not None
|
|
assert initial_state[0].budget_id == default_budget_id, initial_state
|
|
key: Final = scenario.key(team_id=team, user_id=user)
|
|
|
|
first: Final = _chat(gateway, model, key, f"initial member budget {uuid.uuid4().hex}")
|
|
assert first.status_code == 200, first.text
|
|
observed_before_block: Final = _observed_models(upstream)
|
|
assert observed_before_block == (provider_model,), repr(observed_before_block)
|
|
upstream.delete("/__observations").raise_for_status()
|
|
accrued: Final = eventually(
|
|
lambda: _budget_state(team, user),
|
|
lambda values: len(values) == 1 and values[0].spend is not None and values[0].spend >= 0.06,
|
|
seconds=70,
|
|
)
|
|
assert accrued[0].spend == pytest.approx(0.06), accrued
|
|
blocked: Final = _chat(gateway, model, key, f"exceeded member budget {uuid.uuid4().hex}")
|
|
_budget_refused(blocked)
|
|
observed_after_block: Final = _observed_models(upstream)
|
|
assert observed_after_block == (), repr(observed_after_block)
|
|
|
|
expiry: Final = (datetime.now(timezone.utc) + timedelta(days=1)).replace(microsecond=0)
|
|
updated: Final = _member_update(
|
|
gateway,
|
|
{
|
|
"team_id": team,
|
|
"user_id": user,
|
|
"temp_budget_increase": 1.0,
|
|
"temp_budget_expiry": spell_expiry(expiry),
|
|
},
|
|
)
|
|
assert updated.temp_budget_expiry is not None, repr(updated)
|
|
assert updated.model_copy(
|
|
update={"temp_budget_expiry": updated.temp_budget_expiry.replace(tzinfo=timezone.utc)}
|
|
) == _MemberUpdateResponse(
|
|
team_id=team,
|
|
user_id=user,
|
|
temp_budget_increase=1.0,
|
|
temp_budget_expiry=expiry,
|
|
), repr(updated)
|
|
state: Final = _budget_state(team, user)
|
|
assert len(state) == 1, state
|
|
if budget_source == "own-budget":
|
|
assert state[0].max_budget == 0.05, state
|
|
else:
|
|
assert default_budget_id is not None
|
|
assert state[0].budget_id != default_budget_id, state
|
|
assert state[0].max_budget is None, state
|
|
default_budget_after: Final = _budget_record(default_budget_id)
|
|
assert default_budget_after == default_budget_before, repr(default_budget_after)
|
|
assert default_budget_after[0].max_budget == 0.05, default_budget_after
|
|
assert default_budget_after[0].temp_budget_increase is None, default_budget_after
|
|
assert default_budget_after[0].temp_budget_expiry is None, default_budget_after
|
|
assert state[0].temp_budget_increase == 1.0, state
|
|
assert state[0].temp_budget_expiry is not None, state
|
|
assert state[0].temp_budget_expiry.replace(tzinfo=timezone.utc) == expiry, state
|
|
team_info: Final = _team_info(gateway, team)
|
|
listed: Final = next(member for member in team_info.team_memberships if member.user_id == user)
|
|
assert listed.budget_id == state[0].budget_id, repr(listed)
|
|
assert listed.litellm_budget_table is not None, repr(listed)
|
|
assert listed.litellm_budget_table.budget_id == state[0].budget_id, repr(listed)
|
|
assert listed.litellm_budget_table.temp_budget_increase == 1.0, repr(listed)
|
|
assert listed.litellm_budget_table.temp_budget_expiry is not None, repr(listed)
|
|
assert listed.litellm_budget_table.temp_budget_expiry.replace(tzinfo=timezone.utc) == expiry, repr(listed)
|
|
member_info: Final = _member_me(gateway, team, key)
|
|
assert member_info.budget_id == state[0].budget_id, repr(member_info)
|
|
assert member_info.litellm_budget_table is not None, repr(member_info)
|
|
assert member_info.litellm_budget_table.budget_id == state[0].budget_id, repr(member_info)
|
|
assert member_info.litellm_budget_table.temp_budget_increase == 1.0, repr(member_info)
|
|
assert member_info.litellm_budget_table.temp_budget_expiry is not None, repr(member_info)
|
|
assert member_info.litellm_budget_table.temp_budget_expiry.replace(tzinfo=timezone.utc) == expiry, repr(
|
|
member_info
|
|
)
|
|
|
|
readmitted: Final = _chat(gateway, model, key, f"temporary member budget {uuid.uuid4().hex}")
|
|
assert readmitted.status_code == 200, readmitted.text
|
|
observed_after_readmission: Final = _observed_models(upstream)
|
|
assert observed_after_readmission == (provider_model,), repr(observed_after_readmission)
|
|
assert observed_before_block + observed_after_readmission == (provider_model, provider_model), repr(
|
|
(observed_before_block, observed_after_readmission)
|
|
)
|
|
|
|
|
|
def test_member_rpm_limit_update_changes_runtime_admission(gateway: Gateway, upstream: httpx.Client) -> None:
|
|
pytest.skip("BUG: member_update leaves a warmed key using the old RPM limit")
|
|
with gateway.scenario() as scenario:
|
|
model, provider_model = _free_model(scenario)
|
|
team: Final = scenario.team(models=[model])
|
|
user: Final = scenario.user()
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{"team_id": team, "member": {"user_id": user, "role": "user"}},
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
key: Final = scenario.key(team_id=team, user_id=user)
|
|
|
|
warm: Final = _chat(gateway, model, key, f"warm member key before rpm limit {uuid.uuid4().hex}")
|
|
assert warm.status_code == 200, warm.text
|
|
updated: Final = _member_update(gateway, {"team_id": team, "user_id": user, "rpm_limit": 1})
|
|
assert updated.rpm_limit == 1, repr(updated)
|
|
|
|
admitted: Final = _chat(gateway, model, key, f"member rpm admitted request {uuid.uuid4().hex}")
|
|
assert admitted.status_code == 200, admitted.text
|
|
limited: Final = _chat(gateway, model, key, f"member rpm limited request {uuid.uuid4().hex}")
|
|
assert limited.status_code == 429, limited.text
|
|
assert _ErrorResponse.model_validate_json(limited.text).error.type == "throttling_error", limited.text
|
|
observed: Final = _observed_models(upstream)
|
|
assert observed == (provider_model, provider_model), repr(observed)
|
|
|
|
|
|
def test_member_rpm_limit_update_changes_runtime_admission_for_cold_key(
|
|
gateway: Gateway, upstream: httpx.Client
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model, provider_model = _free_model(scenario)
|
|
team: Final = scenario.team(models=[model])
|
|
user: Final = scenario.user()
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{"team_id": team, "member": {"user_id": user, "role": "user"}},
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
key: Final = scenario.key(team_id=team, user_id=user)
|
|
|
|
updated: Final = _member_update(gateway, {"team_id": team, "user_id": user, "rpm_limit": 1})
|
|
assert updated.rpm_limit == 1, repr(updated)
|
|
|
|
admitted: Final = _chat(gateway, model, key, f"cold member rpm admitted request {uuid.uuid4().hex}")
|
|
assert admitted.status_code == 200, admitted.text
|
|
limited: Final = _chat(gateway, model, key, f"cold member rpm limited request {uuid.uuid4().hex}")
|
|
assert limited.status_code == 429, limited.text
|
|
assert _ErrorResponse.model_validate_json(limited.text).error.type == "throttling_error", limited.text
|
|
observed: Final = _observed_models(upstream)
|
|
assert observed == (provider_model,), repr(observed)
|
|
|
|
|
|
def test_member_tpm_limit_update_changes_runtime_admission(gateway: Gateway, upstream: httpx.Client) -> None:
|
|
pytest.skip("BUG: member_update leaves a warmed key using the old TPM limit")
|
|
with gateway.scenario() as scenario:
|
|
model, provider_model = _free_model(scenario)
|
|
team: Final = scenario.team(models=[model])
|
|
user: Final = scenario.user()
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{"team_id": team, "member": {"user_id": user, "role": "user"}},
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
key: Final = scenario.key(team_id=team, user_id=user)
|
|
|
|
warm: Final = _chat(gateway, model, key, f"warm member key before tpm limit {uuid.uuid4().hex}")
|
|
assert warm.status_code == 200, warm.text
|
|
updated: Final = _member_update(gateway, {"team_id": team, "user_id": user, "tpm_limit": 30})
|
|
assert updated.tpm_limit == 30, repr(updated)
|
|
|
|
admitted: Final = _chat(gateway, model, key, f"first member tpm-limited request {uuid.uuid4().hex}")
|
|
assert admitted.status_code == 200, admitted.text
|
|
limited: Final = _chat(gateway, model, key, f"second member tpm-limited request {uuid.uuid4().hex}")
|
|
assert limited.status_code == 429, limited.text
|
|
assert _ErrorResponse.model_validate_json(limited.text).error.type == "throttling_error", limited.text
|
|
observed: Final = _observed_models(upstream)
|
|
assert observed == (provider_model, provider_model), repr(observed)
|
|
|
|
|
|
def test_member_tpm_limit_update_changes_runtime_admission_for_cold_key(
|
|
gateway: Gateway, upstream: httpx.Client
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model, provider_model = _free_model(scenario)
|
|
team: Final = scenario.team(models=[model])
|
|
user: Final = scenario.user()
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{"team_id": team, "member": {"user_id": user, "role": "user"}},
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
key: Final = scenario.key(team_id=team, user_id=user)
|
|
|
|
updated: Final = _member_update(gateway, {"team_id": team, "user_id": user, "tpm_limit": 30})
|
|
assert updated.tpm_limit == 30, repr(updated)
|
|
|
|
admitted: Final = _chat(gateway, model, key, f"cold member tpm admitted request {uuid.uuid4().hex}")
|
|
assert admitted.status_code == 200, admitted.text
|
|
limited: Final = _chat(gateway, model, key, f"cold member tpm limited request {uuid.uuid4().hex}")
|
|
assert limited.status_code == 429, limited.text
|
|
assert _ErrorResponse.model_validate_json(limited.text).error.type == "throttling_error", limited.text
|
|
observed: Final = _observed_models(upstream)
|
|
assert observed == (provider_model,), repr(observed)
|
|
|
|
|
|
def test_member_update_by_email_sets_models_role_and_forks_the_team_default_budget(
|
|
gateway: Gateway, upstream: httpx.Client
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model_one, provider_one = _free_model(scenario)
|
|
model_two, provider_two = _free_model(scenario)
|
|
team: Final = scenario.team(models=[model_one, model_two], team_member_budget=25)
|
|
email_one: Final = f"member-{uuid.uuid4().hex}@example.com"
|
|
user_one: Final = scenario.user(user_email=email_one, user_role="internal_user")
|
|
user_two: Final = scenario.user(user_role="internal_user")
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{
|
|
"team_id": team,
|
|
"member": [
|
|
{"user_email": email_one, "role": "user"},
|
|
{"user_id": user_two, "role": "user"},
|
|
],
|
|
},
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
default_budget_id: Final = _team_default_budget_id(team)
|
|
default_budget_before: Final = _budget_record(default_budget_id)
|
|
assert len(default_budget_before) == 1, repr(default_budget_before)
|
|
initial_one: Final = _budget_state(team, user_one)
|
|
initial_two: Final = _budget_state(team, user_two)
|
|
assert len(initial_one) == 1 and initial_one[0].budget_id == default_budget_id, initial_one
|
|
assert len(initial_two) == 1 and initial_two[0].budget_id == default_budget_id, initial_two
|
|
|
|
user_one_key: Final = scenario.key(team_id=team, user_id=user_one)
|
|
model_scope: Final = _member_update(
|
|
gateway,
|
|
{"team_id": team, "user_email": email_one, "allowed_models": [model_one]},
|
|
)
|
|
assert model_scope.user_id == user_one, repr(model_scope)
|
|
assert model_scope.allowed_models == [model_one], repr(model_scope)
|
|
scoped_state: Final = _budget_state(team, user_one)
|
|
assert len(scoped_state) == 1, scoped_state
|
|
assert scoped_state[0].allowed_models == [model_one], scoped_state
|
|
assert scoped_state[0].budget_id != default_budget_id, scoped_state
|
|
other_member_state: Final = _budget_state(team, user_two)
|
|
assert other_member_state[0].budget_id == default_budget_id, other_member_state
|
|
assert _budget_record(default_budget_id) == default_budget_before, repr(default_budget_before)
|
|
|
|
allowed: Final = _chat(gateway, model_one, user_one_key, f"email allowed model {uuid.uuid4().hex}")
|
|
assert allowed.status_code == 200, allowed.text
|
|
denied: Final = _chat(gateway, model_two, user_one_key, f"email restricted model {uuid.uuid4().hex}")
|
|
assert denied.status_code == 403, denied.text
|
|
assert _ErrorResponse.model_validate_json(denied.text).error.type == "team_model_access_denied", denied.text
|
|
observed_before_clear: Final = _observed_models(upstream)
|
|
assert observed_before_clear == (provider_one,), repr(observed_before_clear)
|
|
upstream.delete("/__observations").raise_for_status()
|
|
|
|
cleared_models: Final = _member_update(
|
|
gateway,
|
|
{"team_id": team, "user_email": email_one, "allowed_models": []},
|
|
)
|
|
assert cleared_models.user_id == user_one, repr(cleared_models)
|
|
assert cleared_models.allowed_models == [], repr(cleared_models)
|
|
cleared_state: Final = _budget_state(team, user_one)
|
|
assert len(cleared_state) == 1, cleared_state
|
|
assert cleared_state[0].allowed_models == [], cleared_state
|
|
unrestricted: Final = _chat(gateway, model_two, user_one_key, f"email cleared model scope {uuid.uuid4().hex}")
|
|
assert unrestricted.status_code == 200, unrestricted.text
|
|
observed_after_clear: Final = _observed_models(upstream)
|
|
assert observed_after_clear == (provider_two,), repr(observed_after_clear)
|
|
|
|
budget_update: Final = _member_update(
|
|
gateway,
|
|
{"team_id": team, "user_email": email_one, "max_budget_in_team": 12},
|
|
)
|
|
assert budget_update.user_id == user_one, repr(budget_update)
|
|
assert budget_update.max_budget_in_team == 12, repr(budget_update)
|
|
user_one_state: Final = _budget_state(team, user_one)
|
|
user_two_state: Final = _budget_state(team, user_two)
|
|
assert len(user_one_state) == 1 and user_one_state[0].budget_id != default_budget_id, user_one_state
|
|
assert user_one_state[0].max_budget == 12, user_one_state
|
|
assert user_two_state == initial_two, user_two_state
|
|
assert _budget_record(default_budget_id) == default_budget_before, repr(default_budget_before)
|
|
team_info: Final = _team_info(gateway, team)
|
|
listed_one: Final = next(member for member in team_info.team_memberships if member.user_id == user_one)
|
|
assert listed_one.litellm_budget_table is not None, repr(listed_one)
|
|
assert listed_one.litellm_budget_table.max_budget == 12, repr(listed_one)
|
|
assert listed_one.litellm_budget_table.allowed_models == [], repr(listed_one)
|
|
member_info: Final = _member_me(gateway, team, user_one_key)
|
|
assert member_info.litellm_budget_table is not None, repr(member_info)
|
|
assert member_info.litellm_budget_table.max_budget == 12, repr(member_info)
|
|
assert member_info.litellm_budget_table.allowed_models == [], repr(member_info)
|
|
roster_member: Final = next(
|
|
member for member in team_info.team_info.members_with_roles if member.user_id == user_one
|
|
)
|
|
assert roster_member.role == "user", repr(roster_member)
|
|
|
|
new_member: Final = scenario.user()
|
|
user_key: Final = scenario.key(user_id=user_one)
|
|
denied_admin_action: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{"team_id": team, "member": {"user_id": new_member, "role": "user"}},
|
|
key=user_key,
|
|
)
|
|
assert denied_admin_action.status_code == 403, denied_admin_action.text
|
|
role_update: Final = _member_update(
|
|
gateway,
|
|
{"team_id": team, "user_email": email_one, "role": "admin"},
|
|
)
|
|
assert role_update.user_id == user_one, repr(role_update)
|
|
role_info: Final = _team_info(gateway, team)
|
|
admin_roster: Final = next(
|
|
member for member in role_info.team_info.members_with_roles if member.user_id == user_one
|
|
)
|
|
assert admin_roster.role == "admin", repr(admin_roster)
|
|
roster_rows: Final = read_rows(
|
|
'SELECT members_with_roles FROM "LiteLLM_TeamTable" WHERE team_id = %s',
|
|
(team,),
|
|
)
|
|
assert len(roster_rows) == 1, roster_rows
|
|
database_roster: Final = _TeamState.model_validate(
|
|
{"team_id": team, "members_with_roles": roster_rows[0]["members_with_roles"]}
|
|
).members_with_roles
|
|
assert next(member for member in database_roster if member.user_id == user_one).role == "admin", repr(
|
|
database_roster
|
|
)
|
|
allowed_admin_action: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{"team_id": team, "member": {"user_id": new_member, "role": "user"}},
|
|
key=user_key,
|
|
)
|
|
assert allowed_admin_action.status_code == 200, allowed_admin_action.text
|
|
assert observed_before_clear + observed_after_clear == (provider_one, provider_two), repr(
|
|
(observed_before_clear, observed_after_clear)
|
|
)
|
|
|
|
|
|
def test_bulk_update_is_a_merge_patch_per_member(gateway: Gateway, upstream: httpx.Client) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model_one, provider_one = _free_model(scenario)
|
|
model_two, _ = _free_model(scenario)
|
|
spend_model, provider_spend = _priced_model(scenario)
|
|
team: Final = scenario.team(models=[model_one, model_two, spend_model], team_member_budget=25)
|
|
user_one: Final = scenario.user()
|
|
email_two: Final = f"bulk-member-{uuid.uuid4().hex}@example.com"
|
|
user_two: Final = scenario.user(user_email=email_two, user_role="internal_user")
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{
|
|
"team_id": team,
|
|
"member": [
|
|
{"user_id": user_one, "role": "user"},
|
|
{"user_email": email_two, "role": "user"},
|
|
],
|
|
},
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
default_budget_id: Final = _team_default_budget_id(team)
|
|
default_budget_before: Final = _budget_record(default_budget_id)
|
|
key_one: Final = scenario.key(team_id=team, user_id=user_one)
|
|
key_two: Final = scenario.key(team_id=team, user_id=user_two)
|
|
personal_key_one: Final = scenario.key(user_id=user_one)
|
|
|
|
first_chat: Final = _chat(gateway, spend_model, key_one, f"bulk member spend {uuid.uuid4().hex}")
|
|
assert first_chat.status_code == 200, first_chat.text
|
|
observed_before_update: Final = _observed_models(upstream)
|
|
assert observed_before_update == (provider_spend,), repr(observed_before_update)
|
|
upstream.delete("/__observations").raise_for_status()
|
|
accrued: Final = eventually(
|
|
lambda: _budget_state(team, user_one),
|
|
lambda values: len(values) == 1 and values[0].spend is not None and values[0].spend >= 0.06,
|
|
seconds=70,
|
|
)
|
|
assert accrued[0].spend == pytest.approx(0.06), accrued
|
|
|
|
update: Final = gateway.request(
|
|
"POST",
|
|
f"/management/v1/teams/{team}/members/bulk_update",
|
|
{
|
|
"members": [
|
|
{"user_id": user_one, "max_budget_in_team": 0.0001, "rpm_limit": 7},
|
|
{"user_email": email_two, "rpm_limit": 5, "allowed_models": [model_one]},
|
|
]
|
|
},
|
|
)
|
|
assert update.status_code == 200, update.text
|
|
update_response: Final = _BulkBudgetUpdateResponse.model_validate_json(update.text)
|
|
state_one: Final = _budget_state(team, user_one)
|
|
state_two: Final = _budget_state(team, user_two)
|
|
assert len(state_one) == 1 and len(state_two) == 1, (state_one, state_two)
|
|
assert state_one[0].budget_id is not None, state_one
|
|
assert state_two[0].budget_id is not None, state_two
|
|
assert state_one[0].budget_id != default_budget_id, state_one
|
|
assert state_two[0].budget_id != default_budget_id, state_two
|
|
assert update_response.data == (
|
|
_BulkBudgetResult(
|
|
user_id=user_one,
|
|
success=True,
|
|
budget_id=state_one[0].budget_id,
|
|
max_budget=0.0001,
|
|
max_budget_source="member",
|
|
rpm_limit=7,
|
|
allowed_models=tuple(state_one[0].allowed_models or ()),
|
|
),
|
|
_BulkBudgetResult(
|
|
user_id=user_two,
|
|
user_email=email_two,
|
|
success=True,
|
|
budget_id=state_two[0].budget_id,
|
|
max_budget=25,
|
|
max_budget_source="member",
|
|
rpm_limit=5,
|
|
allowed_models=(model_one,),
|
|
),
|
|
), update.text
|
|
assert state_one[0].max_budget == 0.0001 and state_one[0].rpm_limit == 7, state_one
|
|
assert state_two[0].max_budget == 25 and state_two[0].rpm_limit == 5, state_two
|
|
assert state_two[0].allowed_models == [model_one], state_two
|
|
assert _budget_record(default_budget_id) == default_budget_before, repr(default_budget_before)
|
|
info: Final = _team_info(gateway, team)
|
|
info_one: Final = next(member for member in info.team_memberships if member.user_id == user_one)
|
|
assert info_one.litellm_budget_table is not None, repr(info_one)
|
|
assert info_one.litellm_budget_table.max_budget == 0.0001, repr(info_one)
|
|
assert info_one.litellm_budget_table.rpm_limit == 7, repr(info_one)
|
|
me_one: Final = _member_me(gateway, team, personal_key_one)
|
|
assert me_one.litellm_budget_table is not None, repr(me_one)
|
|
assert me_one.litellm_budget_table.max_budget == 0.0001, repr(me_one)
|
|
assert me_one.litellm_budget_table.rpm_limit == 7, repr(me_one)
|
|
|
|
over_budget: Final = _chat(gateway, spend_model, key_one, f"bulk member blocked {uuid.uuid4().hex}")
|
|
_budget_refused(over_budget)
|
|
observed_after_over_budget: Final = _observed_models(upstream)
|
|
assert observed_after_over_budget == (), repr(observed_after_over_budget)
|
|
allowed_two: Final = _chat(gateway, model_one, key_two, f"bulk member scoped allowed {uuid.uuid4().hex}")
|
|
assert allowed_two.status_code == 200, allowed_two.text
|
|
denied_two: Final = _chat(gateway, model_two, key_two, f"bulk member scoped denied {uuid.uuid4().hex}")
|
|
assert denied_two.status_code == 403, denied_two.text
|
|
assert _ErrorResponse.model_validate_json(denied_two.text).error.type == "team_model_access_denied", (
|
|
denied_two.text
|
|
)
|
|
observed_after_first_patch: Final = _observed_models(upstream)
|
|
assert observed_after_first_patch == (provider_one,), repr(observed_after_first_patch)
|
|
upstream.delete("/__observations").raise_for_status()
|
|
|
|
cleared_max: Final = gateway.request(
|
|
"POST",
|
|
f"/management/v1/teams/{team}/members/bulk_update",
|
|
{"members": [{"user_id": user_one, "max_budget_in_team": None}]},
|
|
)
|
|
assert cleared_max.status_code == 200, cleared_max.text
|
|
cleared_max_response: Final = _BulkBudgetUpdateResponse.model_validate_json(cleared_max.text)
|
|
cleared_max_state: Final = _budget_state(team, user_one)
|
|
assert len(cleared_max_state) == 1, cleared_max_state
|
|
assert cleared_max_response.data == (
|
|
_BulkBudgetResult(
|
|
user_id=user_one,
|
|
success=True,
|
|
budget_id=cleared_max_state[0].budget_id,
|
|
max_budget=25,
|
|
max_budget_source="team_default",
|
|
rpm_limit=7,
|
|
allowed_models=tuple(cleared_max_state[0].allowed_models or ()),
|
|
),
|
|
), cleared_max.text
|
|
assert cleared_max_state[0].budget_id == state_one[0].budget_id, cleared_max_state
|
|
assert cleared_max_state[0].max_budget is None, cleared_max_state
|
|
assert cleared_max_state[0].rpm_limit == 7, cleared_max_state
|
|
assert _budget_record(default_budget_id) == default_budget_before, repr(default_budget_before)
|
|
me_after_max_clear: Final = _member_me(gateway, team, personal_key_one)
|
|
assert me_after_max_clear.litellm_budget_table is not None, repr(me_after_max_clear)
|
|
assert me_after_max_clear.litellm_budget_table.max_budget is None, repr(me_after_max_clear)
|
|
assert me_after_max_clear.litellm_budget_table.rpm_limit == 7, repr(me_after_max_clear)
|
|
admitted_again: Final = _chat(gateway, spend_model, key_one, f"bulk member readmitted {uuid.uuid4().hex}")
|
|
assert admitted_again.status_code == 200, admitted_again.text
|
|
observed_after_max_clear: Final = _observed_models(upstream)
|
|
assert observed_after_max_clear == (provider_spend,), repr(observed_after_max_clear)
|
|
|
|
assert observed_before_update + observed_after_first_patch + observed_after_max_clear == (
|
|
provider_spend,
|
|
provider_one,
|
|
provider_spend,
|
|
), repr((observed_before_update, observed_after_first_patch, observed_after_max_clear))
|
|
|
|
|
|
def test_bulk_update_restores_team_default_when_last_member_limit_is_cleared(
|
|
gateway: Gateway, upstream: httpx.Client
|
|
) -> None:
|
|
with gateway.scenario() as scenario:
|
|
model, provider_model = _priced_model(scenario)
|
|
team: Final = scenario.team(models=[model], team_member_budget=0.0001)
|
|
user: Final = scenario.user()
|
|
added: Final = gateway.request(
|
|
"POST",
|
|
"/team/member_add",
|
|
{"team_id": team, "member": {"user_id": user, "role": "user"}},
|
|
)
|
|
assert added.status_code == 200, added.text
|
|
default_budget_id: Final = _team_default_budget_id(team)
|
|
default_budget_before: Final = _budget_record(default_budget_id)
|
|
initial_state: Final = _budget_state(team, user)
|
|
assert len(initial_state) == 1 and initial_state[0].budget_id == default_budget_id, initial_state
|
|
key: Final = scenario.key(team_id=team, user_id=user, models=[model])
|
|
increased: Final = gateway.request(
|
|
"POST",
|
|
f"/management/v1/teams/{team}/members/bulk_update",
|
|
{"members": [{"user_id": user, "max_budget_in_team": 5}]},
|
|
)
|
|
assert increased.status_code == 200, increased.text
|
|
increased_response: Final = _BulkBudgetUpdateResponse.model_validate_json(increased.text)
|
|
increased_state: Final = _budget_state(team, user)
|
|
assert len(increased_state) == 1, increased_state
|
|
assert increased_state[0].budget_id is not None and increased_state[0].budget_id != default_budget_id, (
|
|
increased_state
|
|
)
|
|
assert increased_response.data == (
|
|
_BulkBudgetResult(
|
|
user_id=user,
|
|
success=True,
|
|
budget_id=increased_state[0].budget_id,
|
|
max_budget=5,
|
|
max_budget_source="member",
|
|
allowed_models=tuple(increased_state[0].allowed_models or ()),
|
|
),
|
|
), increased.text
|
|
first_chat: Final = _chat(gateway, model, key, f"member raised limit {uuid.uuid4().hex}")
|
|
assert first_chat.status_code == 200, first_chat.text
|
|
first_observations: Final = _observed_models(upstream)
|
|
assert first_observations == (provider_model,), repr(first_observations)
|
|
upstream.delete("/__observations").raise_for_status()
|
|
accrued: Final = eventually(
|
|
lambda: _budget_state(team, user),
|
|
lambda values: len(values) == 1 and values[0].spend is not None and values[0].spend >= 0.06,
|
|
seconds=70,
|
|
)
|
|
assert accrued[0].spend == pytest.approx(0.06), accrued
|
|
second_chat: Final = _chat(gateway, model, key, f"member still under limit {uuid.uuid4().hex}")
|
|
assert second_chat.status_code == 200, second_chat.text
|
|
second_observations: Final = _observed_models(upstream)
|
|
assert second_observations == (provider_model,), repr(second_observations)
|
|
upstream.delete("/__observations").raise_for_status()
|
|
cleared: Final = gateway.request(
|
|
"POST",
|
|
f"/management/v1/teams/{team}/members/bulk_update",
|
|
{"members": [{"user_id": user, "max_budget_in_team": None}]},
|
|
)
|
|
assert cleared.status_code == 200, cleared.text
|
|
cleared_response: Final = _BulkBudgetUpdateResponse.model_validate_json(cleared.text)
|
|
final_state: Final = _budget_state(team, user)
|
|
assert len(final_state) == 1, final_state
|
|
assert final_state[0].budget_id is None, final_state
|
|
assert cleared_response.data == (
|
|
_BulkBudgetResult(
|
|
user_id=user,
|
|
success=True,
|
|
budget_id=None,
|
|
max_budget=0.0001,
|
|
max_budget_source="team_default",
|
|
tpm_limit=None,
|
|
rpm_limit=None,
|
|
budget_duration=None,
|
|
allowed_models=None,
|
|
),
|
|
), cleared.text
|
|
assert read_rows(
|
|
'SELECT budget_id FROM "LiteLLM_TeamMembership" WHERE team_id=%s AND user_id=%s',
|
|
(team, user),
|
|
) == [{"budget_id": None}], cleared.text
|
|
assert _budget_record(default_budget_id) == default_budget_before, repr(default_budget_before)
|
|
refused: Final = _chat(gateway, model, key, f"member default cap reached {uuid.uuid4().hex}")
|
|
_budget_refused(refused)
|
|
observed_after_refusal: Final = _observed_models(upstream)
|
|
assert observed_after_refusal == (), repr(observed_after_refusal)
|
|
assert first_observations + second_observations == (provider_model, provider_model), repr(
|
|
(first_observations, second_observations)
|
|
)
|