mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
Merge 7b894385b3 into 4b1928387f
This commit is contained in:
commit
e01f7e686b
9 changed files with 36 additions and 1 deletions
|
|
@ -7060,6 +7060,7 @@ async def _check_team_member_budget(
|
|||
)
|
||||
|
||||
if team_member_budget is not None:
|
||||
valid_token.team_member_budget = team_member_budget # rebind-ok: carrier for router fallback checks
|
||||
team_member_spend = (loaded_membership.spend if loaded_membership is not None else 0.0) or 0.0
|
||||
|
||||
# Read from cross-pod counter (Redis-first) if available
|
||||
|
|
|
|||
|
|
@ -130,6 +130,18 @@ 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(
|
||||
|
|
|
|||
|
|
@ -73,6 +73,7 @@ class TeamGrants(TypedDict, total=False):
|
|||
team_member_spend: ReadOnly[float | None]
|
||||
team_member_tpm_limit: ReadOnly[int | None]
|
||||
team_member_rpm_limit: ReadOnly[int | None]
|
||||
team_member_budget: ReadOnly[float | None]
|
||||
|
||||
|
||||
_TEAM_GRANT_AUTH_FIELDS: Final[frozenset[str]] = frozenset(
|
||||
|
|
@ -210,4 +211,9 @@ def team_grants(
|
|||
team_member_rpm_limit=(
|
||||
team_membership.safe_get_team_member_rpm_limit() if team_membership is not None else None
|
||||
),
|
||||
team_member_budget=(
|
||||
team_membership.litellm_budget_table.max_budget
|
||||
if team_membership is not None and team_membership.litellm_budget_table is not None
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -2976,6 +2976,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
|
|||
team_member_budget: Final = team_member_info.litellm_budget_table.effective_max_budget(
|
||||
now=datetime.now(timezone.utc),
|
||||
)
|
||||
valid_token.team_member_budget = team_member_budget # rebind-ok: carrier for router fallback checks
|
||||
if team_member_budget is not None and team_member_budget > 0:
|
||||
# Read from cross-pod counter (Redis-first) if available
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
|
|
|||
|
|
@ -5771,6 +5771,7 @@ class PrismaClient:
|
|||
tm.spend AS team_member_spend,
|
||||
b_tm.tpm_limit AS team_member_tpm_limit,
|
||||
b_tm.rpm_limit AS team_member_rpm_limit,
|
||||
b_tm.max_budget AS team_member_budget,
|
||||
m.aliases AS team_model_aliases,
|
||||
-- Added comma to separate b.* columns
|
||||
b.max_budget AS litellm_budget_table_max_budget,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -50,7 +50,7 @@ def _membership() -> LiteLLM_TeamMembership:
|
|||
user_id=USER_ID,
|
||||
team_id=TEAM_ID,
|
||||
spend=3.25,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=500, rpm_limit=5),
|
||||
litellm_budget_table=LiteLLM_BudgetTable(tpm_limit=500, rpm_limit=5, max_budget=100.0),
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -85,6 +85,7 @@ def test_team_grants_cover_every_team_field_the_key_path_gets():
|
|||
assert token.team_member_spend == 3.25
|
||||
assert token.team_member_tpm_limit == 500
|
||||
assert token.team_member_rpm_limit == 5
|
||||
assert token.team_member_budget == 100.0
|
||||
|
||||
|
||||
def test_team_grants_without_team_leave_token_defaults():
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -49273,6 +49273,8 @@ export interface components {
|
|||
/** Team Max Budget */
|
||||
team_max_budget?: number | null;
|
||||
team_member?: components["schemas"]["Member"] | null;
|
||||
/** Team Member Budget */
|
||||
team_member_budget?: number | null;
|
||||
/** Team Member Rpm Limit */
|
||||
team_member_rpm_limit?: number | null;
|
||||
/** Team Member Spend */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue