Skip to content

Commit

Permalink
Document team admins + Enforce assigning team admins as an enterprise…
Browse files Browse the repository at this point in the history
… feature (#7359)

* fix(team_endpoints.py): enforce assigning team admins as an enterprise feature

* fix(proxy/_types.py): fix common proxy error to link to trial key

* fix: fix linting errors
  • Loading branch information
krrishdholakia authored Dec 22, 2024
1 parent a8ae2f5 commit 20feacc
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 11 deletions.
2 changes: 1 addition & 1 deletion litellm/proxy/_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -1892,7 +1892,7 @@ class CommonProxyErrors(str, enum.Enum):
db_not_connected_error = "DB not connected"
no_llm_router = "No models configured on proxy"
not_allowed_access = "Admin-only endpoint. Not allowed to access this."
not_premium_user = "You must be a LiteLLM Enterprise user to use this feature. If you have a license please set `LITELLM_LICENSE` in your env. If you want to obtain a license meet with us here: https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat. \nPricing: https://www.litellm.ai/#pricing"
not_premium_user = "You must be a LiteLLM Enterprise user to use this feature. If you have a license please set `LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/#trial. \nPricing: https://www.litellm.ai/#pricing"


class SpendCalculateRequest(LiteLLMPydanticObjectBase):
Expand Down
64 changes: 54 additions & 10 deletions litellm/proxy/management_endpoints/team_endpoints.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
import traceback
import uuid
from datetime import datetime, timedelta, timezone
from typing import List, Optional, Union
from typing import List, Optional, Union, cast

import fastapi
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
Expand Down Expand Up @@ -547,6 +547,49 @@ async def update_team(
return {"team_id": team_row.team_id, "data": team_row}


def _check_team_member_admin_add(
member: Union[Member, List[Member]],
premium_user: bool,
):
if isinstance(member, Member) and member.role == "admin":
if premium_user is not True:
raise ValueError(
f"Assigning team admins is a premium feature. {CommonProxyErrors.not_premium_user.value}"
)
elif isinstance(member, List):
for m in member:
if m.role == "admin":
if premium_user is not True:
raise ValueError(
f"Assigning team admins is a premium feature. Got={m}. {CommonProxyErrors.not_premium_user.value}. "
)


def team_call_validation_checks(
prisma_client: Optional[PrismaClient],
data: TeamMemberAddRequest,
premium_user: bool,
):
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})

if data.team_id is None:
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})

if data.member is None:
raise HTTPException(
status_code=400, detail={"error": "No member/members passed in"}
)

try:
_check_team_member_admin_add(
member=data.member,
premium_user=premium_user,
)
except Exception as e:
raise HTTPException(status_code=400, detail={"error": str(e)})


@router.post(
"/team/member_add",
tags=["team management"],
Expand Down Expand Up @@ -578,21 +621,22 @@ async def team_member_add(
"""
from litellm.proxy.proxy_server import (
litellm_proxy_admin_name,
premium_user,
prisma_client,
proxy_logging_obj,
user_api_key_cache,
)

if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No db connected"})

if data.team_id is None:
raise HTTPException(status_code=400, detail={"error": "No team id passed in"})

if data.member is None:
raise HTTPException(
status_code=400, detail={"error": "No member/members passed in"}
try:
team_call_validation_checks(
prisma_client=prisma_client,
data=data,
premium_user=premium_user,
)
except HTTPException as e:
raise e

prisma_client = cast(PrismaClient, prisma_client)

existing_team_row = await get_team_object(
team_id=data.team_id,
Expand Down

0 comments on commit 20feacc

Please sign in to comment.