diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bbe20d047..b2a8c8965 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f615232b7..44d4240b1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 5377fe90b..7e8e26fcd 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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 -} \ No newline at end of file + status String // approved, disapproved, pending +} diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 3cad1777c..cb29b0492 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -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={ diff --git a/schema.prisma b/schema.prisma index 8663db1b0..7e8e26fcd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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 -} \ No newline at end of file +}