forked from phoenix/litellm-mirror
feat(proxy_server.py): add new team_member delete endpoint
This commit is contained in:
Krrish Dholakia
parent
2fac06de08
commit
6b8b0f40cf
5 changed files with 216 additions and 2 deletions
|
|
@ -2070,6 +2070,36 @@
|
|||
"output_cost_per_token": 0.00000028,
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "chat"
|
||||
},
|
||||
"perplexity/sonar-small-chat": {
|
||||
"max_tokens": 16384,
|
||||
"input_cost_per_token": 0.00000007,
|
||||
"output_cost_per_token": 0.00000028,
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "chat"
|
||||
},
|
||||
"perplexity/sonar-small-online": {
|
||||
"max_tokens": 12000,
|
||||
"input_cost_per_token": 0,
|
||||
"output_cost_per_token": 0.00000028,
|
||||
"input_cost_per_request": 0.005,
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "chat"
|
||||
},
|
||||
"perplexity/sonar-medium-chat": {
|
||||
"max_tokens": 16384,
|
||||
"input_cost_per_token": 0.0000006,
|
||||
"output_cost_per_token": 0.0000018,
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "chat"
|
||||
},
|
||||
"perplexity/sonar-medium-online": {
|
||||
"max_tokens": 12000,
|
||||
"input_cost_per_token": 0,
|
||||
"output_cost_per_token": 0.0000018,
|
||||
"input_cost_per_request": 0.005,
|
||||
"litellm_provider": "perplexity",
|
||||
"mode": "chat"
|
||||
},
|
||||
"anyscale/mistralai/Mistral-7B-Instruct-v0.1": {
|
||||
"max_tokens": 16384,
|
||||
|
|
|
|||
|
|
@ -262,7 +262,19 @@ class NewTeamRequest(LiteLLMBase):
|
|||
|
||||
class TeamMemberAddRequest(LiteLLMBase):
|
||||
team_id: str
|
||||
member: Optional[Member] = None
|
||||
member: Member
|
||||
|
||||
|
||||
class TeamMemberDeleteRequest(LiteLLMBase):
|
||||
team_id: str
|
||||
user_id: Optional[str] = None
|
||||
user_email: Optional[str] = None
|
||||
|
||||
@root_validator(pre=True)
|
||||
def check_user_info(cls, values):
|
||||
if values.get("user_id") is None and values.get("user_email") is None:
|
||||
raise ValueError("Either user id or user email must be provided")
|
||||
return values
|
||||
|
||||
|
||||
class UpdateTeamRequest(LiteLLMBase):
|
||||
|
|
|
|||
|
|
@ -5143,6 +5143,117 @@ async def team_member_add(
|
|||
return team_row
|
||||
|
||||
|
||||
@router.post(
|
||||
"/team/member_delete",
|
||||
tags=["team management"],
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
)
|
||||
async def team_member_delete(
|
||||
data: TeamMemberDeleteRequest,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
):
|
||||
"""
|
||||
[BETA]
|
||||
|
||||
delete members (either via user_email or user_id) from a team
|
||||
|
||||
If user doesn't exist, an exception will be raised
|
||||
```
|
||||
curl -X POST 'http://0.0.0.0:8000/team/update' \
|
||||
|
||||
-H 'Authorization: Bearer sk-1234' \
|
||||
|
||||
-H 'Content-Type: application/json' \
|
||||
|
||||
-D '{
|
||||
"team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849",
|
||||
"member": {"role": "user", "user_id": "krrish247652@berri.ai"}
|
||||
}'
|
||||
```
|
||||
"""
|
||||
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.user_id is None and data.user_email is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Either user_id or user_email needs to be passed in"},
|
||||
)
|
||||
|
||||
existing_team_row = await prisma_client.get_data( # type: ignore
|
||||
team_id=data.team_id, table_name="team", query_type="find_unique"
|
||||
)
|
||||
|
||||
## DELETE MEMBER FROM TEAM
|
||||
new_team_members = []
|
||||
for m in existing_team_row.members_with_roles:
|
||||
if (
|
||||
data.user_id is not None
|
||||
and m["user_id"] is not None
|
||||
and data.user_id == m["user_id"]
|
||||
):
|
||||
continue
|
||||
elif (
|
||||
data.user_email is not None
|
||||
and m["user_email"] is not None
|
||||
and data.user_email == m["user_email"]
|
||||
):
|
||||
continue
|
||||
new_team_members.append(m)
|
||||
existing_team_row.members_with_roles = new_team_members
|
||||
complete_team_data = LiteLLM_TeamTable(
|
||||
**existing_team_row.model_dump(),
|
||||
)
|
||||
|
||||
team_row = await prisma_client.update_data(
|
||||
update_key_values=complete_team_data.json(exclude_none=True),
|
||||
data=complete_team_data.json(exclude_none=True),
|
||||
table_name="team",
|
||||
team_id=data.team_id,
|
||||
)
|
||||
|
||||
## DELETE TEAM ID from USER ROW, IF EXISTS ##
|
||||
# get user row
|
||||
key_val = {}
|
||||
if data.user_id is not None:
|
||||
key_val["user_id"] = data.user_id
|
||||
elif data.user_email is not None:
|
||||
key_val["user_email"] = data.user_email
|
||||
existing_user_rows = await prisma_client.get_data(
|
||||
key_val=key_val,
|
||||
table_name="user",
|
||||
query_type="find_all",
|
||||
)
|
||||
user_data = { # type: ignore
|
||||
"teams": [],
|
||||
"models": team_row["data"].models,
|
||||
}
|
||||
if existing_user_rows is not None and (
|
||||
isinstance(existing_user_rows, list) and len(existing_user_rows) > 0
|
||||
):
|
||||
for existing_user in existing_user_rows:
|
||||
team_list = []
|
||||
if hasattr(existing_user, "teams"):
|
||||
team_list = existing_user.teams
|
||||
team_list.remove(data.team_id)
|
||||
user_data["user_id"] = existing_user.user_id
|
||||
await prisma_client.update_data(
|
||||
user_id=existing_user.user_id,
|
||||
data=user_data,
|
||||
update_key_values_custom_query={
|
||||
"teams": {
|
||||
"set": [team_row["team_id"]],
|
||||
}
|
||||
},
|
||||
table_name="user",
|
||||
)
|
||||
|
||||
return team_row["data"]
|
||||
|
||||
|
||||
@router.post(
|
||||
"/team/delete", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -4881,7 +4881,7 @@ def get_optional_params(
|
|||
extra_body # openai client supports `extra_body` param
|
||||
)
|
||||
else: # assume passing in params for openai/azure openai
|
||||
print_verbose(f"UNMAPPED PROVIDER, ASSUMING IT'S OPENAI/AZUREs")
|
||||
print_verbose(f"UNMAPPED PROVIDER, ASSUMING IT'S OPENAI/AZURE")
|
||||
supported_params = [
|
||||
"functions",
|
||||
"function_call",
|
||||
|
|
|
|||
|
|
@ -35,6 +35,25 @@ async def new_user(session, i, user_id=None, budget=None, budget_duration=None):
|
|||
return await response.json()
|
||||
|
||||
|
||||
async def delete_member(session, i, team_id, user_id):
|
||||
url = "http://0.0.0.0:4000/team/member_delete"
|
||||
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
||||
data = {"team_id": team_id, "user_id": user_id}
|
||||
|
||||
async with session.post(url, headers=headers, json=data) as response:
|
||||
status = response.status
|
||||
response_text = await response.text()
|
||||
|
||||
print(f"Response {i} (Status code: {status}):")
|
||||
print(response_text)
|
||||
print()
|
||||
|
||||
if status != 200:
|
||||
raise Exception(f"Request {i} did not return a 200 status code: {status}")
|
||||
|
||||
return await response.json()
|
||||
|
||||
|
||||
async def generate_key(
|
||||
session,
|
||||
i,
|
||||
|
|
@ -290,3 +309,45 @@ async def test_team_delete():
|
|||
response = await chat_completion(session=session, key=key)
|
||||
## Delete team
|
||||
await delete_team(session=session, i=0, team_id=team_data["team_id"])
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_member_delete():
|
||||
"""
|
||||
- Create team
|
||||
- Add member
|
||||
- Get team info (check if member in team)
|
||||
- Delete member
|
||||
- Get team info (check if member in team)
|
||||
"""
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Create Team
|
||||
## Create admin
|
||||
admin_user = f"{uuid.uuid4()}"
|
||||
await new_user(session=session, i=0, user_id=admin_user)
|
||||
## Create normal user
|
||||
normal_user = f"{uuid.uuid4()}"
|
||||
print(f"normal_user: {normal_user}")
|
||||
await new_user(session=session, i=0, user_id=normal_user)
|
||||
## Create team with 1 admin and 1 user
|
||||
member_list = [
|
||||
{"role": "admin", "user_id": admin_user},
|
||||
{"role": "user", "user_id": normal_user},
|
||||
]
|
||||
team_data = await new_team(session=session, i=0, member_list=member_list)
|
||||
print(f"team_data: {team_data}")
|
||||
member_id_list = []
|
||||
for member in team_data["members_with_roles"]:
|
||||
member_id_list.append(member["user_id"])
|
||||
|
||||
assert normal_user in member_id_list
|
||||
# Delete member
|
||||
updated_team_data = await delete_member(
|
||||
session=session, i=0, team_id=team_data["team_id"], user_id=normal_user
|
||||
)
|
||||
print(f"updated_team_data: {updated_team_data}")
|
||||
member_id_list = []
|
||||
for member in updated_team_data["members_with_roles"]:
|
||||
member_id_list.append(member["user_id"])
|
||||
|
||||
assert normal_user not in member_id_list
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue