scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101

feat: team scoped model overrides

Browse source
This commit is contained in:
Harshit28j 2026-02-28 01:05:52 +05:30
parent 6600c86dbd
commit a1b79edfeb
8 changed files with 162 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/proxy/_types.py
View file

@ -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):

32
litellm/proxy/auth/auth_checks.py
View file

@ -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",

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

@ -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"

70
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

@ -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,

4
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -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,
)

18
litellm/proxy/management_helpers/utils.py
View file

@ -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},
)
)

2
litellm/proxy/schema.prisma
View file

@ -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])
}

2
litellm/proxy/utils.py
View file

@ -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,