scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(proxy): green up CI for team token budgets

Browse source 
- ruff format auth_checks / db_spend_update_writer / proxy_track_cost_callback
- document max_tokens_budget on /team/new and /team/update (api-docs gate)
- regenerate dashboard schema.d.ts from the updated OpenAPI spec
- reset-budget tests: team resets now also zero tokens_used, so assert the
  extra key for team writes only (key/user/org payloads are unchanged)
- memory test fixture: pass tokens_used=0 when hand-building the Prisma team
  row (the column is NOT NULL @default(0), exactly like spend)
This commit is contained in:
Srivatsa03 2026-10-11 20:50:48 -05:00
parent a4816fb9f0
commit cbf9c40b9b
7 changed files with 43 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/auth/auth_checks.py
View file

@ -6523,7 +6523,6 @@ async def team_max_budget_check(
)
async def team_max_tokens_check(
team_object: LiteLLM_TeamTable | None,
valid_token: UserAPIKeyAuth | None,

4
litellm/proxy/db/db_spend_update_writer.py
View file

@ -2168,9 +2168,7 @@ class DBSpendUpdateWriter:
### UPDATE TEAM TABLE ###
team_list_transactions: Final = db_spend_update_transactions["team_list_transactions"]
team_token_list_transactions: Final = (
db_spend_update_transactions.get("team_token_list_transactions") or {}
)
team_token_list_transactions: Final = db_spend_update_transactions.get("team_token_list_transactions") or {}
verbose_proxy_logger.debug("Team Spend transactions: %s", team_list_transactions)
if team_list_transactions is not None and len(team_list_transactions.keys()) > 0:
for i in range(n_retry_times + 1):

4
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -914,9 +914,7 @@ async def _update_database_and_spend_counters_in_batch(
increment=total_tokens,
)
except Exception:
verbose_proxy_logger.warning(
"Failed to increment team token-budget counter for team_id=%s", team_id
)
verbose_proxy_logger.warning("Failed to increment team token-budget counter for team_id=%s", team_id)
return True

2
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -1448,6 +1448,7 @@ async def new_team(
- rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement.
- tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement.
- max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
- max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped.
- soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
- budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
- models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
@ -2227,6 +2228,7 @@ async def update_team(
- rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit
- tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit
- max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
- max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped.
- soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
- budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
- models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.

18
tests/unit/proxy/common_utils/test_reset_budget_job.py
View file

@ -431,7 +431,9 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client):
assert write["where"] == {"team_id": "tid-1"}
assert write["data"]["spend"] == {"decrement": 500.0}
assert write["data"]["budget_reset_at"] > now
assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
# Team resets also hard-zero the token-budget accumulator for the new window.
assert write["data"]["tokens_used"] == 0
assert set(write["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"}
def test_reset_budget_for_enduser(reset_budget_job, mock_prisma_client):
@ -552,7 +554,11 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client):
assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}"
assert writes[0]["where"] == where
assert writes[0]["data"]["spend"] == {"decrement": decrement}
assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
# Teams carry the extra token-budget accumulator, zeroed for the new window.
expected_keys = {"spend", "budget_reset_at"}
if table_name == "team":
expected_keys = expected_keys | {"tokens_used"}
assert set(writes[0]["data"].keys()) == expected_keys
# The budget tier's cascade rides the same batch machinery.
assert _batch_writes(mock_prisma_client, "enduser") == [
@ -4108,9 +4114,13 @@ async def test_reset_budget_continues_other_categories_on_failure(frozen_reset_c
assert len(user_writes) == 1
assert user_writes[0]["where"] == {"user_id": "user2"}
assert len(team_writes) == 2
for c in key_writes + user_writes + team_writes:
for c in key_writes + user_writes:
assert set(c["data"].keys()) == {"spend", "budget_reset_at"}
assert c["data"]["spend"] == {"decrement": pre_reset_spend[next(iter(c["where"].values()))]}
# Team writes additionally zero the token-budget accumulator for the new window.
for c in team_writes:
assert set(c["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"}
assert c["data"]["spend"] == {"decrement": pre_reset_spend[next(iter(c["where"].values()))]}
@pytest.mark.asyncio
@ -4166,7 +4176,7 @@ async def test_reset_budget_teams_partial_failure():
team_writes = [c for c in batch_calls if c["table"] == "team"]
assert len(team_writes) == 1
assert team_writes[0]["where"] == {"team_id": "team2"}
assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at"}
assert set(team_writes[0]["data"].keys()) == {"spend", "budget_reset_at", "tokens_used"}
assert team_writes[0]["data"]["spend"] == {"decrement": pre_reset_spend}
failure_hook_calls = (

1
tests/unit/proxy/memory/test_memory_endpoints.py
View file

@ -197,6 +197,7 @@ def _make_team(team_id: str, *, admin_user_ids: List[str]) -> Any:
created_at=now,
updated_at=now,
spend=0.0,
tokens_used=0,
model_spend="{}",
model_max_budget="{}",
admins=[],

24
ui/litellm-dashboard/src/lib/http/schema.d.ts generated vendored
View file

@ -16912,6 +16912,7 @@ export interface paths {
* - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement.
* - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement.
* - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
* - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped.
* - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set, soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
* - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
* - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
@ -17143,6 +17144,7 @@ export interface paths {
* - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit
* - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit
* - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget
* - max_tokens_budget: Optional[int] - The maximum cumulative tokens the team may consume over budget_duration (mirrors max_budget). Null means token usage is tracked but not capped.
* - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set.
* - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets)
* - models: Optional[list] - A list of models associated with the team - all keys for this team_id will have at most, these models. If empty, assumes all models are allowed.
@ -33745,6 +33747,8 @@ export interface components {
max_budget?: number | null;
/** Max Parallel Requests */
max_parallel_requests?: number | null;
/** Max Tokens Budget */
max_tokens_budget?: number | null;
/**
* Members
* @default []
@ -33803,6 +33807,8 @@ export interface components {
team_id: string;
/** Team Member Permissions */
team_member_permissions?: string[] | null;
/** Tokens Used */
tokens_used?: number | null;
/** Tpd Limit */
tpd_limit?: number | null;
/** Tpm Limit */
@ -35276,6 +35282,8 @@ export interface components {
max_budget?: number | null;
/** Max Parallel Requests */
max_parallel_requests?: number | null;
/** Max Tokens Budget */
max_tokens_budget?: number | null;
/**
* Members
* @default []
@ -35334,6 +35342,8 @@ export interface components {
team_id: string;
/** Team Member Permissions */
team_member_permissions?: string[] | null;
/** Tokens Used */
tokens_used?: number | null;
/** Tpd Limit */
tpd_limit?: number | null;
/** Tpm Limit */
@ -38596,6 +38606,8 @@ export interface components {
guardrails?: string[] | null;
/** Max Budget */
max_budget?: number | null;
/** Max Tokens Budget */
max_tokens_budget?: number | null;
/** Mcp Rpm Limit */
mcp_rpm_limit?: {
[key: string]: number;
@ -39837,6 +39849,8 @@ export interface components {
guardrails?: string[] | null;
/** Max Budget */
max_budget?: number | null;
/** Max Tokens Budget */
max_tokens_budget?: number | null;
/** Mcp Rpm Limit */
mcp_rpm_limit?: {
[key: string]: number;
@ -45475,6 +45489,8 @@ export interface components {
max_budget?: number | null;
/** Max Parallel Requests */
max_parallel_requests?: number | null;
/** Max Tokens Budget */
max_tokens_budget?: number | null;
/**
* Members
* @default []
@ -45533,6 +45549,8 @@ export interface components {
team_id: string;
/** Team Member Permissions */
team_member_permissions?: string[] | null;
/** Tokens Used */
tokens_used?: number | null;
/** Tpd Limit */
tpd_limit?: number | null;
/** Tpm Limit */
@ -45612,6 +45630,8 @@ export interface components {
max_budget?: number | null;
/** Max Parallel Requests */
max_parallel_requests?: number | null;
/** Max Tokens Budget */
max_tokens_budget?: number | null;
/**
* Members
* @default []
@ -45675,6 +45695,8 @@ export interface components {
team_id: string;
/** Team Member Permissions */
team_member_permissions?: string[] | null;
/** Tokens Used */
tokens_used?: number | null;
/** Tpd Limit */
tpd_limit?: number | null;
/** Tpm Limit */
@ -47851,6 +47873,8 @@ export interface components {
guardrails?: string[] | null;
/** Max Budget */
max_budget?: number | null;
/** Max Tokens Budget */
max_tokens_budget?: number | null;
/** Mcp Rpm Limit */
mcp_rpm_limit?: {
[key: string]: number;