phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix: show all teams user is a part of in user_info

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-21 13:29:42 -08:00
parent 5a0f962beb
commit 846757e343
5 changed files with 85 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/_types.py
View file

@ -224,11 +224,17 @@ class UpdateUserRequest(GenerateRequestBase):
max_budget: Optional[float] = None
class Member(LiteLLMBase):
role: Literal["admin", "user"]
user_id: str
class NewTeamRequest(LiteLLMBase):
team_alias: Optional[str] = None
team_id: Optional[str] = None
admins: list = []
members: list = []
members_with_roles: List[Member] = []
metadata: Optional[dict] = None

62
litellm/proxy/proxy_server.py
View file

@ -4060,9 +4060,30 @@ async def user_info(
else:
user_info = None
## GET ALL TEAMS ##
teams = await prisma_client.get_data(
team_list = []
team_id_list = []
# _DEPRECATED_ check if user in 'member' field
teams_1 = await prisma_client.get_data(
user_id=user_id, table_name="team", query_type="find_all"
)
if teams_1 is not None and isinstance(teams_1, list):
team_list = teams_1
for team in teams_1:
team_id_list.append(team.team_id)
if user_info is not None:
# *NEW* get all teams in user 'teams' field
teams_2 = await prisma_client.get_data(
team_id_list=user_info.teams, table_name="team", query_type="find_all"
)
if teams_2 is not None and isinstance(teams_2, list):
for team in teams_2:
if team.team_id not in team_id_list:
team_list.append(team)
team_id_list.append(team.team_id)
## GET ALL KEYS ##
keys = await prisma_client.get_data(
user_id=user_id,
@ -4090,9 +4111,10 @@ async def user_info(
"user_id": user_id,
"user_info": user_info,
"keys": keys,
"teams": teams,
"teams": team_list,
}
except Exception as e:
traceback.print_exc()
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({str(e)})"),
@ -4274,12 +4296,31 @@ async def new_team(
Parameters:
- team_alias: Optional[str] - User defined team alias
- team_id: Optional[str] - The team id of the user. If none passed, we'll generate it.
- admins: list - A list of user IDs that will be owning the team
- members: list - A list of user IDs that will be members of the team
- members_with_roles: list - A list of dictionaries, mapping user_id to role in team (either 'admin' or 'user')
- metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" }
Returns:
- team_id: (str) Unique team id - used for tracking spend across multiple keys for same team id.
_deprecated_params:
- admins: list - A list of user_id's for the admin role
- users: list - A list of user_id's for the user role
Example Request:
```
curl --location 'http://0.0.0.0:8000/team/new' \
--header 'Authorization: Bearer sk-1234' \
--header 'Content-Type: application/json' \
--data '{
"team_alias": "my-new-team_2",
"members_with_roles": [{"role": "admin", "user_id": "user-1234"},
{"role": "user", "user_id": "user-2434"}]
}'
```
"""
global prisma_client
@ -4303,6 +4344,19 @@ async def new_team(
team_row = await prisma_client.insert_data(
data=complete_team_data.json(exclude_none=True), table_name="team"
)
## ADD TEAM ID TO USER TABLE ##
for user in complete_team_data.members_with_roles:
## add team id to user row ##
await prisma_client.update_data(
user_id=user.user_id,
data={"user_id": user.user_id, "teams": [team_row.team_id]},
update_key_values={
"teams": {
"push ": [team_row.team_id],
}
},
)
return team_row

6
litellm/proxy/schema.prisma
View file

@ -13,6 +13,7 @@ model LiteLLM_TeamTable {
team_alias String?
admins String[]
members String[]
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
spend Float @default(0.0)
@ -32,6 +33,7 @@ model LiteLLM_TeamTable {
model LiteLLM_UserTable {
user_id String @unique
team_id String?
teams String[] @default([])
user_role String?
max_budget Float?
spend Float @default(0.0)
@ -103,5 +105,5 @@ model LiteLLM_UserNotifications {
user_id String
models String[]
justification String
status String // approved, disapproved, pending
}
status String // approved, disapproved, pending
}

15
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"]
@ -697,7 +698,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 +776,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={

4
schema.prisma
View file

@ -13,6 +13,7 @@ model LiteLLM_TeamTable {
team_alias String?
admins String[]
members String[]
members_with_roles Json @default("{}")
metadata Json @default("{}")
max_budget Float?
spend Float @default(0.0)
@ -32,6 +33,7 @@ model LiteLLM_TeamTable {
model LiteLLM_UserTable {
user_id String @unique
team_id String?
teams String[] @default([])
user_role String?
max_budget Float?
spend Float @default(0.0)
@ -104,4 +106,4 @@ model LiteLLM_UserNotifications {
models String[]
justification String
status String // approved, disapproved, pending
}
}