From 0891c13429c57c8b8b90ade5325f811a86193998 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 1 Apr 2026 15:35:07 -0700 Subject: [PATCH] common_utils: add allowed_models to _upsert_budget_and_membership --- .../proxy/management_endpoints/common_utils.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index efc42d3355c..cc79eb30dbf 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -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,