scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

common_utils: add allowed_models to _upsert_budget_and_membership

Browse source
This commit is contained in:
Ishaan Jaffer 2026-04-01 15:35:07 -07:00
parent 5361579640
commit 0891c13429
1 changed files with 12 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
litellm/proxy/management_endpoints/common_utils.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING, Any, Dict, Optional, Union
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache
@ -354,6 +354,7 @@ async def _upsert_budget_and_membership(
user_api_key_dict: UserAPIKeyAuth,
tpm_limit: Optional[int] = None,
rpm_limit: Optional[int] = None,
allowed_models: Optional[List[str]] = None,
):
"""
Helper function to Create/Update or Delete the budget within the team membership
@ -366,11 +367,17 @@ async def _upsert_budget_and_membership(
user_api_key_dict: User API Key dictionary containing user information
tpm_limit: Tokens per minute limit for the team member
rpm_limit: Requests per minute limit for the team member
allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce.
If max_budget, tpm_limit, and rpm_limit are all None, the user's budget is removed from the team membership.
If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership.
If any of these values exist, a budget is updated or created and linked to the team membership.
"""
if max_budget is None and tpm_limit is None and rpm_limit is None:
if (
max_budget is None
and tpm_limit is None
and rpm_limit is None
and allowed_models is None
):
# disconnect the budget since all limits are None
await tx.litellm_teammembership.update(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
@ -389,6 +396,8 @@ async def _upsert_budget_and_membership(
create_data["tpm_limit"] = tpm_limit
if rpm_limit is not None:
create_data["rpm_limit"] = rpm_limit
if allowed_models is not None:
create_data["allowed_models"] = allowed_models
new_budget = await tx.litellm_budgettable.create(
data=create_data,