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

fix(proxy_server.py): allow user to set team tpm/rpm limits/budget/models

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-24 11:37:20 -08:00
parent 194c823783
commit 5b0eab2224
2 changed files with 51 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/_types.py
View file

@ -236,6 +236,10 @@ class NewTeamRequest(LiteLLMBase):
members: list = []
members_with_roles: List[Member] = []
metadata: Optional[dict] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
max_budget: Optional[float] = None
models: list = []
class UpdateTeamRequest(LiteLLMBase):
@ -252,25 +256,12 @@ class DeleteTeamRequest(LiteLLMBase):
class LiteLLM_TeamTable(NewTeamRequest):
max_budget: Optional[float] = None
spend: Optional[float] = None
models: list = []
max_parallel_requests: Optional[int] = None
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
budget_duration: Optional[str] = None
budget_reset_at: Optional[datetime] = None
class NewTeamResponse(LiteLLMBase):
team_id: str
admins: list
members: list
metadata: dict
created_at: datetime
updated_at: datetime
class TeamRequest(LiteLLMBase):
teams: List[str]

52
litellm/proxy/proxy_server.py
View file

@ -4314,7 +4314,7 @@ async def user_get_requests():
"/team/new",
tags=["team management"],
dependencies=[Depends(user_api_key_auth)],
response_model=NewTeamResponse,
response_model=LiteLLM_TeamTable,
)
async def new_team(
data: NewTeamRequest,
@ -4360,13 +4360,55 @@ async def new_team(
if data.team_id is None:
data.team_id = str(uuid.uuid4())
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}"
},
)
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}"
},
)
if (
data.max_budget is not None
and user_api_key_dict.max_budget is not None
and data.max_budget > user_api_key_dict.max_budget
):
raise HTTPException(
status_code=400,
detail={
"error": f"max budget higher than user max. User max budget={user_api_key_dict.max_budget}"
},
)
if data.models is not None:
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}"
},
)
complete_team_data = LiteLLM_TeamTable(
**data.json(),
max_budget=user_api_key_dict.max_budget,
models=user_api_key_dict.models,
max_parallel_requests=user_api_key_dict.max_parallel_requests,
tpm_limit=user_api_key_dict.tpm_limit,
rpm_limit=user_api_key_dict.rpm_limit,
budget_duration=user_api_key_dict.budget_duration,
budget_reset_at=user_api_key_dict.budget_reset_at,
)