From a1b79edfeb23a2df99c6f7cfca0cac3ba1eaee88 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Sat, 28 Feb 2026 01:05:52 +0530 Subject: [PATCH] feat: team scoped model overrides --- litellm/proxy/_types.py | 8 +++ litellm/proxy/auth/auth_checks.py | 32 ++++++++- .../management_endpoints/common_utils.py | 46 ++++++++---- .../key_management_endpoints.py | 70 +++++++++++++++++++ .../management_endpoints/team_endpoints.py | 4 ++ litellm/proxy/management_helpers/utils.py | 18 +++-- litellm/proxy/schema.prisma | 2 + litellm/proxy/utils.py | 2 + 8 files changed, 162 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index de1609baf62..c39d9b5e359 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1509,6 +1509,7 @@ class TeamBase(LiteLLMPydanticObjectBase): budget_duration: Optional[str] = None models: list = [] + default_models: list = [] blocked: bool = False router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None @@ -1580,6 +1581,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): max_budget: Optional[float] = None soft_budget: Optional[float] = None models: Optional[list] = None + default_models: Optional[list] = None blocked: Optional[bool] = None budget_duration: Optional[str] = None tags: Optional[list] = None @@ -2251,6 +2253,8 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_max_budget: Optional[float] = None team_soft_budget: Optional[float] = None team_models: List = [] + team_default_models: List[str] = [] + team_member_models: List[str] = [] team_blocked: bool = False soft_budget: Optional[float] = None team_model_aliases: Optional[Dict] = None @@ -3407,6 +3411,7 @@ class LiteLLM_TeamMembership(LiteLLMPydanticObjectBase): team_id: str budget_id: Optional[str] = None spend: Optional[float] = 0.0 + models: List[str] = [] litellm_budget_table: Optional[LiteLLM_BudgetTable] def safe_get_team_member_rpm_limit(self) -> Optional[int]: @@ -3520,6 +3525,7 @@ class TeamMemberAddRequest(MemberAddRequest): default=None, description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", ) + models: Optional[List[str]] = None class TeamMemberDeleteRequest(MemberDeleteRequest): @@ -3535,6 +3541,7 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): rpm_limit: Optional[int] = Field( default=None, description="Requests per minute limit for this team member" ) + models: Optional[List[str]] = None class TeamMemberUpdateResponse(MemberUpdateResponse): @@ -3542,6 +3549,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse): max_budget_in_team: Optional[float] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + models: Optional[List[str]] = None class TeamModelAddRequest(BaseModel): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 500a39d9455..a82a1849194 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -65,6 +65,7 @@ from litellm.utils import get_utc_datetime from .auth_checks_organization import organization_role_based_access_check from .auth_utils import get_model_from_request +from litellm.proxy.management_endpoints.common_utils import _is_team_model_overrides_enabled if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -273,6 +274,7 @@ async def common_checks( team_object=team_object, llm_router=llm_router, team_model_aliases=valid_token.team_model_aliases if valid_token else None, + valid_token=valid_token, ): raise ProxyException( message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", @@ -2557,11 +2559,20 @@ def can_org_access_model( ) +def compute_effective_team_models( + team_default_models: List[str], + team_member_models: List[str], +) -> List[str]: + """Union of team defaults and per-user overrides, deduplicated.""" + return list(set(team_default_models) | set(team_member_models)) + + async def can_team_access_model( model: Union[str, List[str]], team_object: Optional[LiteLLM_TeamTable], llm_router: Optional[Router], team_model_aliases: Optional[Dict[str, str]] = None, + valid_token: Optional[UserAPIKeyAuth] = None, ) -> Literal[True]: """ Returns True if the team can access a specific model. @@ -2569,11 +2580,30 @@ async def can_team_access_model( 1. First checks native team-level model permissions (current implementation) 2. If not allowed natively, falls back to access_group_ids on the team """ + models_to_check: List[str] = team_object.models if team_object else [] + if _is_team_model_overrides_enabled() and valid_token: + # Compute effective models: team defaults + per-user overrides + effective_models = compute_effective_team_models( + team_default_models=valid_token.team_default_models, + team_member_models=valid_token.team_member_models, + ) + + # If effective_models is empty, and feature is enabled, deny access + if len(effective_models) == 0: + raise ProxyException( + message=f"Team not allowed to access model. No models available for user in this team. Model={model}.", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + + models_to_check = effective_models + try: return _can_object_call_model( model=model, llm_router=llm_router, - models=team_object.models if team_object else [], + models=models_to_check, team_model_aliases=team_model_aliases, team_id=team_object.team_id if team_object else None, object_type="team", diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index 999c53e6b82..6b3aa6202a0 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,3 +1,4 @@ +import os from typing import TYPE_CHECKING, Any, Dict, Optional, Union from litellm._logging import verbose_proxy_logger @@ -312,6 +313,7 @@ async def _upsert_budget_and_membership( user_api_key_dict: UserAPIKeyAuth, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, + models: Optional[List[str]] = None, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -324,15 +326,20 @@ 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 + models: Model names the team member is allowed to call (per-user overrides) If max_budget, tpm_limit, and rpm_limit 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: # disconnect the budget since all limits are None + update_data: Dict[str, Any] = {"litellm_budget_table": {"disconnect": True}} + if models is not None: + update_data["models"] = models + await tx.litellm_teammembership.update( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, - data={"litellm_budget_table": {"disconnect": True}}, + data=update_data, ) return @@ -352,6 +359,25 @@ async def _upsert_budget_and_membership( data=create_data, include={"team_membership": True}, ) + + update_payload: Dict[str, Any] = { + "litellm_budget_table": { + "connect": {"budget_id": new_budget.budget_id}, + }, + } + if models is not None: + update_payload["models"] = models + + create_payload: Dict[str, Any] = { + "user_id": user_id, + "team_id": team_id, + "litellm_budget_table": { + "connect": {"budget_id": new_budget.budget_id}, + }, + } + if models is not None: + create_payload["models"] = models + # upsert the team membership with the new/updated budget await tx.litellm_teammembership.upsert( where={ @@ -361,18 +387,8 @@ async def _upsert_budget_and_membership( } }, data={ - "create": { - "user_id": user_id, - "team_id": team_id, - "litellm_budget_table": { - "connect": {"budget_id": new_budget.budget_id}, - }, - }, - "update": { - "litellm_budget_table": { - "connect": {"budget_id": new_budget.budget_id}, - }, - }, + "create": create_payload, + "update": update_payload, }, ) @@ -429,3 +445,7 @@ def _update_metadata_fields(updated_kv: dict) -> None: for field in LiteLLM_ManagementEndpoint_MetadataFields: if field in updated_kv and updated_kv[field] is not None: _update_metadata_field(updated_kv=updated_kv, field_name=field) + + +def _is_team_model_overrides_enabled() -> bool: + return os.getenv("LITELLM_TEAM_MODEL_OVERRIDES", "false").lower() == "true" diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b489369071f..5b58ffde3db 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -890,6 +890,62 @@ async def _check_team_key_limits( ) +async def _validate_key_models_against_effective_team_models( + team_id: str, + user_id: Optional[str], + requested_models: List[str], + team_table: LiteLLM_TeamTableCachedObj, + prisma_client: PrismaClient, +) -> None: + """ + Validate that the requested models for a key are a subset of the effective team models. + Effective models = team.default_models ∪ membership.models + """ + from litellm.proxy.management_endpoints.common_utils import ( + _is_team_model_overrides_enabled, + ) + + if not _is_team_model_overrides_enabled(): + return + + if not requested_models: + return + + # 1. Fetch team membership if user_id is provided + member_models: List[str] = [] + if user_id: + membership = await prisma_client.db.litellm_teammembership.find_unique( + where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}} + ) + if membership: + member_models = membership.models or [] + + # 2. Compute effective models + from litellm.proxy.auth.auth_checks import compute_effective_team_models + + effective_models = compute_effective_team_models( + team_default_models=team_table.default_models or [], + team_member_models=member_models, + ) + + # 3. If effective models are defined, validate requested models are a subset + if effective_models: + for m in requested_models: + if m not in effective_models: + raise HTTPException( + status_code=400, + detail={ + "error": f"Model '{m}' is not available for this user in Team={team_id}. Available models = {effective_models}" + }, + ) + else: + # If no effective models (empty defaults AND empty member overrides) + # and feature is enabled, we only allow models that are in the global team models. + # But according to plan: "If Effective models is empty, follow EXISTING behavior (team.models)." + # So we don't raise error here. + pass + + async def _check_project_key_limits( project_id: str, data: Union[GenerateKeyRequest, UpdateKeyRequest], @@ -1201,6 +1257,13 @@ async def generate_key_fn( data=data, prisma_client=prisma_client, ) + await _validate_key_models_against_effective_team_models( + team_id=data.team_id, + user_id=data.user_id, + requested_models=data.models or [], + team_table=team_table, + prisma_client=prisma_client, + ) # Validate key against project limits if project_id is set if data.project_id is not None: @@ -1354,6 +1417,13 @@ async def generate_service_account_key_fn( data=data, prisma_client=prisma_client, ) + await _validate_key_models_against_effective_team_models( + team_id=data.team_id, + user_id=data.user_id, + requested_models=data.models or [], + team_table=team_table, + prisma_client=prisma_client, + ) key_generation_check( team_table=team_table, diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 80d50f31a17..498d41afa25 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1688,6 +1688,7 @@ async def _process_team_members( litellm_proxy_admin_name=litellm_proxy_admin_name, team_id=data.team_id, default_team_budget_id=default_team_budget_id, + models=data.models, ) except Exception as e: raise HTTPException( @@ -1712,6 +1713,7 @@ async def _process_team_members( litellm_proxy_admin_name=litellm_proxy_admin_name, team_id=data.team_id, default_team_budget_id=default_team_budget_id, + models=data.models, ) except Exception as e: raise HTTPException( @@ -2336,6 +2338,7 @@ async def team_member_update( user_api_key_dict=user_api_key_dict, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + models=data.models, ) ### update team member role @@ -2368,6 +2371,7 @@ async def team_member_update( max_budget_in_team=data.max_budget_in_team, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + models=data.models, ) diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 67a1ea659f9..c10074268ed 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -146,6 +146,7 @@ async def add_new_member( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, default_team_budget_id: Optional[str] = None, + models: Optional[List[str]] = None, ) -> Tuple[LiteLLM_UserTable, Optional[LiteLLM_TeamMembership]]: """ Add a new member to a team @@ -220,14 +221,19 @@ async def add_new_member( else: _budget_id = default_team_budget_id - if _budget_id and returned_user is not None and returned_user.user_id is not None: + if returned_user is not None and returned_user.user_id is not None: + create_data: Dict[str, Any] = { + "team_id": team_id, + "user_id": returned_user.user_id, + } + if _budget_id: + create_data["budget_id"] = _budget_id + if models is not None: + create_data["models"] = models + _returned_team_membership = ( await prisma_client.db.litellm_teammembership.create( - data={ - "team_id": team_id, - "user_id": returned_user.user_id, - "budget_id": _budget_id, - }, + data=create_data, # type: ignore include={"litellm_budget_table": True}, ) ) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 440c9c1d829..a71127fcb88 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -119,6 +119,7 @@ model LiteLLM_TeamTable { soft_budget Float? spend Float @default(0.0) models String[] + default_models String[] @default([]) max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? @@ -538,6 +539,7 @@ model LiteLLM_TeamMembership { team_id String spend Float @default(0.0) budget_id String? + models String[] @default([]) litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) @@id([user_id, team_id]) } diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f6613b5548f..11a5d6e39e2 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -2833,6 +2833,7 @@ class PrismaClient: t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, t.models AS team_models, + t.default_models AS team_default_models, t.metadata AS team_metadata, t.blocked AS team_blocked, t.team_alias AS team_alias, @@ -2841,6 +2842,7 @@ class PrismaClient: t.object_permission_id AS team_object_permission_id, t.organization_id as org_id, tm.spend AS team_member_spend, + tm.models AS team_member_models, m.aliases AS team_model_aliases, -- Added comma to separate b.* columns b.max_budget AS litellm_budget_table_max_budget,