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

Fix/new org team validate against org (#17333)

Browse source 
* fix: skip user budget/model validation for org-scoped teams

When creating a team with organization_id, budget and model constraints
should be validated against the organization's limits, not the user's
personal limits. This allows org admins with restrictive personal
budgets to create teams within their organization's more generous limits.

Adds 4 unit tests to verify:
- Org-scoped teams bypass user budget validation
- Org-scoped teams bypass user model validation
- Standalone teams still validate against user limits

* fix: enforce user budget/model limits for standalone teams in update_team

- Add user-level budget and model validation to update_team endpoint for standalone teams,
  matching the existing pattern in new_team
- Org-scoped teams correctly bypass user validation and use organization limits instead
- Add 5 new comprehensive tests covering standalone/org team budget/model validation

* fix: Add direct TPM/RPM org limit validation and consolidate user team limit checks

- Add direct TPM/RPM comparison against org limits in _check_org_team_limits()
- Consolidate budget/models/TPM/RPM user validation into _check_user_team_limits() helper
- Ensure user limits only apply to standalone teams (organization_id=None)
- Org-scoped teams now validate TPM/RPM against org limits (not user limits)
- Add 8 tests for TPM/RPM validation scenarios (org and user limits)
- Reduce code duplication between new_team() and update_team()
This commit is contained in:
rioiart 2025-12-02 04:51:42 +01:00 • committed by GitHub
parent 1cdfb3da8f
commit 70126d9130
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 1901 additions and 197 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -461,11 +461,65 @@ async def _check_org_team_limits(
prisma_client: PrismaClient,
) -> None:
"""
Check if the organization team is allocating guaranteed throughput limits. If so, raise an error if we're overallocating.
Only runs check if tpm_limit_type or rpm_limit_type is "guaranteed_throughput"
Check organization team limits including:
- Team budget vs organization's max_budget
- Team models vs organization's allowed models
- Guaranteed throughput limits (tpm/rpm) if applicable
"""
# Validate team budget against organization's max_budget
if (
data.max_budget is not None
and org_table.litellm_budget_table is not None
and org_table.litellm_budget_table.max_budget is not None
and data.max_budget > org_table.litellm_budget_table.max_budget
):
raise HTTPException(
status_code=400,
detail={
"error": f"Team max_budget ({data.max_budget}) exceeds organization's max_budget ({org_table.litellm_budget_table.max_budget}). Organization: {org_table.organization_id}"
},
)
# Validate team models against organization's allowed models
if data.models is not None and len(org_table.models) > 0:
for m in data.models:
if m not in org_table.models:
raise HTTPException(
status_code=400,
detail={
"error": f"Model '{m}' not in organization's allowed models. Organization allowed models={org_table.models}. Organization: {org_table.organization_id}"
},
)
# Validate team TPM/RPM against organization's TPM/RPM limits (direct comparison)
if (
data.tpm_limit is not None
and org_table.litellm_budget_table is not None
and org_table.litellm_budget_table.tpm_limit is not None
and data.tpm_limit > org_table.litellm_budget_table.tpm_limit
):
raise HTTPException(
status_code=400,
detail={
"error": f"Team tpm_limit ({data.tpm_limit}) exceeds organization's tpm_limit ({org_table.litellm_budget_table.tpm_limit}). Organization: {org_table.organization_id}"
},
)
if (
data.rpm_limit is not None
and org_table.litellm_budget_table is not None
and org_table.litellm_budget_table.rpm_limit is not None
and data.rpm_limit > org_table.litellm_budget_table.rpm_limit
):
raise HTTPException(
status_code=400,
detail={
"error": f"Team rpm_limit ({data.rpm_limit}) exceeds organization's rpm_limit ({org_table.litellm_budget_table.rpm_limit}). Organization: {org_table.organization_id}"
},
)
# Check guaranteed throughput limits (only if applicable)
rpm_limit_type = getattr(data, "rpm_limit_type", None) or (
data.metadata.get("rpm_limit_type", None) if data.metadata else None
)
@ -503,6 +557,80 @@ async def _check_org_team_limits(
)
async def _check_user_team_limits(
data: Union[NewTeamRequest, UpdateTeamRequest],
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
user_api_key_cache: Any,
) -> None:
"""
Check user team limits for standalone teams (not org-scoped).
This validates:
- Team budget vs user's max_budget
- Team models vs user's allowed models
Should only be called for standalone teams (when organization_id is None).
For org-scoped teams, use _check_org_team_limits() instead.
"""
# Validate team budget against user's max_budget
if data.max_budget is not None and user_api_key_dict.user_id is not None:
user_obj = await get_user_object(
user_id=user_api_key_dict.user_id,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
if (
user_obj is not None
and user_obj.max_budget is not None
and data.max_budget > user_obj.max_budget
):
raise HTTPException(
status_code=400,
detail={
"error": f"max budget higher than user max. User max budget={user_obj.max_budget}. User role={user_api_key_dict.user_role}"
},
)
# Validate team models against user's allowed models
if data.models is not None and len(user_api_key_dict.models) > 0:
for m in data.models:
if m not in user_api_key_dict.models:
raise HTTPException(
status_code=400,
detail={
"error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}. User id={user_api_key_dict.user_id}"
},
)
# Validate team TPM/RPM against user's TPM/RPM limits
if (
data.tpm_limit is not None
and user_api_key_dict.tpm_limit is not None
and data.tpm_limit > user_api_key_dict.tpm_limit
):
raise HTTPException(
status_code=400,
detail={
"error": f"tpm limit higher than user max. User tpm limit={user_api_key_dict.tpm_limit}. User role={user_api_key_dict.user_role}"
},
)
if (
data.rpm_limit is not None
and user_api_key_dict.rpm_limit is not None
and data.rpm_limit > user_api_key_dict.rpm_limit
):
raise HTTPException(
status_code=400,
detail={
"error": f"rpm limit higher than user max. User rpm limit={user_api_key_dict.rpm_limit}. User role={user_api_key_dict.user_role}"
},
)
#### TEAM MANAGEMENT ####
@router.post(
"/team/new",
@ -665,61 +793,16 @@ async def new_team( # noqa: PLR0915
user_api_key_dict.user_role is None
or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
): # don't restrict proxy admin
if (
data.tpm_limit is not None
and user_api_key_dict.tpm_limit is not None
and data.tpm_limit > user_api_key_dict.tpm_limit
):
raise HTTPException(
status_code=400,
detail={
"error": f"tpm limit higher than user max. User tpm limit={user_api_key_dict.tpm_limit}. User role={user_api_key_dict.user_role}"
},
)
if (
data.rpm_limit is not None
and user_api_key_dict.rpm_limit is not None
and data.rpm_limit > user_api_key_dict.rpm_limit
):
raise HTTPException(
status_code=400,
detail={
"error": f"rpm limit higher than user max. User rpm limit={user_api_key_dict.rpm_limit}. User role={user_api_key_dict.user_role}"
},
)
if data.max_budget is not None and user_api_key_dict.user_id is not None:
# Fetch user object to get max_budget
user_obj = await get_user_object(
user_id=user_api_key_dict.user_id,
# Only validate user budget/models/tpm/rpm for standalone teams (not org-scoped)
# For org-scoped teams, validation is done by _check_org_team_limits()
if data.organization_id is None:
await _check_user_team_limits(
data=data,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
user_id_upsert=False,
)
if (
user_obj is not None
and user_obj.max_budget is not None
and data.max_budget > user_obj.max_budget
):
raise HTTPException(
status_code=400,
detail={
"error": f"max budget higher than user max. User max budget={user_obj.max_budget}. User role={user_api_key_dict.user_role}"
},
)
if data.models is not None and len(user_api_key_dict.models) > 0:
for m in data.models:
if m not in user_api_key_dict.models:
raise HTTPException(
status_code=400,
detail={
"error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}. User id={user_api_key_dict.user_id}"
},
)
if user_api_key_dict.user_id is not None:
creating_user_in_list = False
for member in data.members_with_roles:
@ -1151,168 +1234,187 @@ async def update_team(
}'
```
"""
from litellm.proxy.auth.auth_checks import _cache_team_object
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
try:
from litellm.proxy.auth.auth_checks import _cache_team_object
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
llm_router,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)
if data.team_id is None:
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
verbose_proxy_logger.debug("/team/update - %s", data)
existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": data.team_id}
)
if existing_team_row is None:
raise HTTPException(
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
if (
data.organization_id is not None and len(data.organization_id) > 0
): # allow unsetting the organization_id
await fetch_and_validate_organization(
organization_id=data.organization_id,
existing_team_row=existing_team_row,
llm_router=llm_router,
prisma_client=prisma_client,
)
elif data.organization_id is not None and len(data.organization_id) == 0:
# unsetting the organization_id
data.organization_id = None
# check org team limits - if updating team that belongs to an org
org_id_to_check = (
data.organization_id
if data.organization_id is not None
else existing_team_row.organization_id
)
if (
org_id_to_check is not None
and isinstance(org_id_to_check, str)
and prisma_client is not None
):
org_table = await get_org_object(
org_id=org_id_to_check,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
if org_table is not None:
await _check_org_team_limits(
org_table=org_table,
data=data,
prisma_client=prisma_client,
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
updated_kv = data.json(exclude_unset=True)
if data.team_id is None:
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})
verbose_proxy_logger.debug("/team/update - %s", data)
# Check budget_duration and budget_reset_at
if data.budget_duration is not None:
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": data.team_id}
)
reset_at = get_budget_reset_time(budget_duration=data.budget_duration)
if existing_team_row is None:
raise HTTPException(
status_code=404,
detail={"error": f"Team not found, passed team_id={data.team_id}"},
)
# set the budget_reset_at in DB
updated_kv["budget_reset_at"] = reset_at
if (
data.organization_id is not None and len(data.organization_id) > 0
): # allow unsetting the organization_id
await fetch_and_validate_organization(
organization_id=data.organization_id,
existing_team_row=existing_team_row,
llm_router=llm_router,
prisma_client=prisma_client,
)
elif data.organization_id is not None and len(data.organization_id) == 0:
# unsetting the organization_id
data.organization_id = None
if TeamMemberBudgetHandler.should_create_budget(
team_member_budget=data.team_member_budget,
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
):
updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table(
team_table=existing_team_row,
user_api_key_dict=user_api_key_dict,
updated_kv=updated_kv,
# check org team limits - if updating team that belongs to an org
org_id_to_check = (
data.organization_id
if data.organization_id is not None
else existing_team_row.organization_id
)
if (
org_id_to_check is not None
and isinstance(org_id_to_check, str)
and prisma_client is not None
):
org_table = await get_org_object(
org_id=org_id_to_check,
user_api_key_cache=user_api_key_cache,
prisma_client=prisma_client,
)
if org_table is not None:
await _check_org_team_limits(
org_table=org_table,
data=data,
prisma_client=prisma_client,
)
# Check user limits for standalone teams (not org-scoped)
# Skip for PROXY_ADMIN users
if (
user_api_key_dict.user_role is None
or user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN
):
# Only validate user budget/models for standalone teams
# For org-scoped teams, validation is done by _check_org_team_limits() above
if org_id_to_check is None:
await _check_user_team_limits(
data=data,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
user_api_key_cache=user_api_key_cache,
)
updated_kv = data.json(exclude_unset=True)
# Check budget_duration and budget_reset_at
if data.budget_duration is not None:
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
reset_at = get_budget_reset_time(budget_duration=data.budget_duration)
# set the budget_reset_at in DB
updated_kv["budget_reset_at"] = reset_at
if TeamMemberBudgetHandler.should_create_budget(
team_member_budget=data.team_member_budget,
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
)
else:
TeamMemberBudgetHandler._clean_team_member_fields(updated_kv)
# Check object permission
if data.object_permission is not None:
updated_kv = await handle_update_object_permission(
data_json=updated_kv,
existing_team_row=existing_team_row,
)
# update team metadata fields
_team_metadata_fields = LiteLLM_ManagementEndpoint_MetadataFields_Premium
for field in _team_metadata_fields:
if field in updated_kv and updated_kv[field] is not None:
_update_metadata_field(
):
updated_kv = await TeamMemberBudgetHandler.upsert_team_member_budget_table(
team_table=existing_team_row,
user_api_key_dict=user_api_key_dict,
updated_kv=updated_kv,
field_name=field,
team_member_budget=data.team_member_budget,
team_member_rpm_limit=data.team_member_rpm_limit,
team_member_tpm_limit=data.team_member_tpm_limit,
)
else:
TeamMemberBudgetHandler._clean_team_member_fields(updated_kv)
# Check object permission
if data.object_permission is not None:
updated_kv = await handle_update_object_permission(
data_json=updated_kv,
existing_team_row=existing_team_row,
)
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,
# update team metadata fields
_team_metadata_fields = LiteLLM_ManagementEndpoint_MetadataFields_Premium
for field in _team_metadata_fields:
if field in updated_kv and updated_kv[field] is not None:
_update_metadata_field(
updated_kv=updated_kv,
field_name=field,
)
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,
)
if "model_aliases" in updated_kv:
updated_kv.pop("model_aliases")
_model_id = await _update_model_table(
data=data,
model_id=existing_team_row.model_id,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
if _model_id is not None:
updated_kv["model_id"] = _model_id
updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
team_row: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.update(
where={"team_id": data.team_id},
data=updated_kv,
include={"litellm_model_table": True}, # type: ignore
)
)
if team_row is None or team_row.team_id is None:
raise HTTPException(
status_code=400,
detail={"error": "Team doesn't exist. Got={}".format(team_row)},
)
if "model_aliases" in updated_kv:
updated_kv.pop("model_aliases")
_model_id = await _update_model_table(
data=data,
model_id=existing_team_row.model_id,
prisma_client=prisma_client,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
if _model_id is not None:
updated_kv["model_id"] = _model_id
updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
team_row: Optional[LiteLLM_TeamTable] = (
await prisma_client.db.litellm_teamtable.update(
where={"team_id": data.team_id},
data=updated_kv,
include={"litellm_model_table": True}, # type: ignore
)
)
if team_row is None or team_row.team_id is None:
raise HTTPException(
status_code=400,
detail={"error": "Team doesn't exist. Got={}".format(team_row)},
verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id)
await _cache_team_object(
team_id=team_row.team_id,
team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id)
await _cache_team_object(
team_id=team_row.team_id,
team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()),
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=proxy_logging_obj,
)
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
if litellm.store_audit_logs is True:
await _create_team_update_audit_log(
existing_team_row=existing_team_row,
updated_kv=updated_kv,
team_id=data.team_id,
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
if litellm.store_audit_logs is True:
await _create_team_update_audit_log(
existing_team_row=existing_team_row,
updated_kv=updated_kv,
team_id=data.team_id,
litellm_changed_by=litellm_changed_by,
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
)
return {"team_id": team_row.team_id, "data": team_row}
return {"team_id": team_row.team_id, "data": team_row}
except Exception as e:
raise handle_exception_on_proxy(e)
async def handle_update_object_permission(

1602
tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
View file

File diff suppressed because it is too large Load diff