phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Merge pull request #2119 from BerriAI/litellm_updated_team_endpoints

Browse source 
Enable `/team/update`, `/team/delete` endpoints + create teams with user defined roles
This commit is contained in:
Krish Dholakia 2024-02-21 17:24:58 -08:00 • committed by GitHub
commit 0733bf1e7a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 464 additions and 35 deletions
Download patch file Download diff file Expand all files Collapse all files

82
litellm/proxy/utils.py
View file

@ -532,6 +532,7 @@ class PrismaClient:
user_id: Optional[str] = None,
user_id_list: Optional[list] = None,
team_id: Optional[str] = None,
team_id_list: Optional[list] = None,
key_val: Optional[dict] = None,
table_name: Optional[
Literal["user", "key", "config", "spend", "team", "user_notification"]
@ -641,13 +642,12 @@ class PrismaClient:
}
)
elif query_type == "find_all" and user_id_list is not None:
user_id_values = str(tuple(user_id_list))
user_id_values = ", ".join(f"'{item}'" for item in user_id_list)
sql_query = f"""
SELECT *
FROM "LiteLLM_UserTable"
WHERE "user_id" IN {user_id_values}
WHERE "user_id" IN ({user_id_values})
"""
# Execute the raw query
# The asterisk before `user_id_list` unpacks the list into separate arguments
response = await self.db.query_raw(sql_query)
@ -697,7 +697,13 @@ class PrismaClient:
)
elif query_type == "find_all" and user_id is not None:
response = await self.db.litellm_teamtable.find_many(
where={"members": {"has": user_id}}
where={
"members": {"has": user_id},
},
)
elif query_type == "find_all" and team_id_list is not None:
response = await self.db.litellm_teamtable.find_many(
where={"team_id": {"in": team_id_list}}
)
return response
elif table_name == "user_notification":
@ -769,6 +775,12 @@ class PrismaClient:
return new_user_row
elif table_name == "team":
db_data = self.jsonify_object(data=data)
if db_data.get("members_with_roles", None) is not None and isinstance(
db_data["members_with_roles"], list
):
db_data["members_with_roles"] = json.dumps(
db_data["members_with_roles"]
)
new_team_row = await self.db.litellm_teamtable.upsert(
where={"team_id": data["team_id"]},
data={
@ -915,6 +927,19 @@ class PrismaClient:
update_key_values = db_data
if "team_id" not in db_data and team_id is not None:
db_data["team_id"] = team_id
if "members_with_roles" in db_data and isinstance(
db_data["members_with_roles"], list
):
db_data["members_with_roles"] = json.dumps(
db_data["members_with_roles"]
)
if "members_with_roles" in update_key_values and isinstance(
update_key_values["members_with_roles"], list
):
update_key_values["members_with_roles"] = json.dumps(
update_key_values["members_with_roles"]
)
update_team_row = await self.db.litellm_teamtable.upsert(
where={"team_id": team_id}, # type: ignore
data={
@ -929,7 +954,7 @@ class PrismaClient:
+ f"DB Team Table - update succeeded {update_team_row}"
+ "\033[0m"
)
return {"team_id": team_id, "data": db_data}
return {"team_id": team_id, "data": update_team_row}
elif (
table_name is not None
and table_name == "key"
@ -1001,22 +1026,45 @@ class PrismaClient:
max_time=10, # maximum total time to retry for
on_backoff=on_backoff, # specifying the function to call on backoff
)
async def delete_data(self, tokens: List):
async def delete_data(
self,
tokens: Optional[List] = None,
team_id_list: Optional[List] = None,
table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None,
):
"""
Allow user to delete a key(s)
"""
try:
hashed_tokens = []
for token in tokens:
if isinstance(token, str) and token.startswith("sk-"):
hashed_token = self.hash_token(token=token)
else:
hashed_token = token
hashed_tokens.append(hashed_token)
await self.db.litellm_verificationtoken.delete_many(
where={"token": {"in": hashed_tokens}}
)
return {"deleted_keys": tokens}
if tokens is not None and isinstance(tokens, List):
hashed_tokens = []
for token in tokens:
if isinstance(token, str) and token.startswith("sk-"):
hashed_token = self.hash_token(token=token)
else:
hashed_token = token
hashed_tokens.append(hashed_token)
await self.db.litellm_verificationtoken.delete_many(
where={"token": {"in": hashed_tokens}}
)
return {"deleted_keys": tokens}
elif (
table_name == "team"
and team_id_list is not None
and isinstance(team_id_list, List)
):
await self.db.litellm_teamtable.delete_many(
where={"team_id": {"in": team_id_list}}
)
return {"deleted_teams": team_id_list}
elif (
table_name == "key"
and team_id_list is not None
and isinstance(team_id_list, List)
):
await self.db.litellm_verificationtoken.delete_many(
where={"team_id": {"in": team_id_list}}
)
except Exception as e:
asyncio.create_task(
self.proxy_logging_obj.failure_handler(original_exception=e)