mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(fallback_budget): enforce team member budget on router fallback targets
This commit is contained in:
parent
f81e657903
commit
b59ba74733
3 changed files with 27 additions and 0 deletions
|
|
@ -130,6 +130,22 @@ async def is_token_within_budget_for_model(
|
|||
if is_model_cost_zero(model=model, llm_router=llm_router):
|
||||
return True
|
||||
|
||||
# Gating check: team member overall budget within their team.
|
||||
# Prevents zero-cost models with paid fallbacks from exceeding the member's team budget cap.
|
||||
team_member_budget: Final = valid_token.team_member_budget
|
||||
if (
|
||||
team_member_budget is not None
|
||||
and valid_token.user_id is not None
|
||||
and valid_token.team_id is not None
|
||||
):
|
||||
team_member_spend: Final = await _counter_spend(
|
||||
counter_key=f"spend:team_member:{valid_token.user_id}:{valid_token.team_id}",
|
||||
fallback_spend=valid_token.team_member_spend or 0.0,
|
||||
max_budget=team_member_budget,
|
||||
)
|
||||
if team_member_spend >= team_member_budget:
|
||||
return False
|
||||
|
||||
key_budget: Final = valid_token.max_budget
|
||||
if key_budget is not None and valid_token.token is not None:
|
||||
key_spend: Final = await _counter_spend(
|
||||
|
|
|
|||
|
|
@ -78,6 +78,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
|
|||
team_member_spend: float | None = None
|
||||
team_member_tpm_limit: int | None = None
|
||||
team_member_rpm_limit: int | None = None
|
||||
team_member_budget: float | None = None
|
||||
|
||||
# End User Params
|
||||
end_user_id: str | None = None
|
||||
|
|
|
|||
|
|
@ -115,6 +115,16 @@ async def test_paid_target_refused_when_over_user_budget():
|
|||
token = _token(user_spend=1900.0, user_max_budget=50.0)
|
||||
assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_paid_target_refused_when_over_team_member_budget():
|
||||
token = _token(user_id="u1", team_id="t1", team_member_budget=50.0, team_member_spend=100.0)
|
||||
assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is False
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_paid_target_allowed_when_under_team_member_budget():
|
||||
token = _token(user_id="u1", team_id="t1", team_member_budget=50.0, team_member_spend=10.0)
|
||||
assert await is_token_within_budget_for_model(model="paid-model", valid_token=token, llm_router=_router()) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_refuses_an_exhausted_paid_target() -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue