build(ui): enable admin to create teams, add members, create keys for teams

litellm/proxy/_types.py

@@ -226,7 +226,14 @@ class UpdateUserRequest(GenerateRequestBase):
 
 class Member(LiteLLMBase):
     role: Literal["admin", "user"]
-    user_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 NewTeamRequest(LiteLLMBase):
@@ -242,6 +249,11 @@ class NewTeamRequest(LiteLLMBase):
     models: list = []
 
 
+class TeamMemberAddRequest(LiteLLMBase):
+    team_id: str
+    member: Optional[Member] = None
+
+
 class UpdateTeamRequest(LiteLLMBase):
     team_id: str  # required
     team_alias: Optional[str] = None
@@ -261,6 +273,25 @@ class LiteLLM_TeamTable(NewTeamRequest):
     budget_duration: Optional[str] = None
     budget_reset_at: Optional[datetime] = None
 
+    @root_validator(pre=True)
+    def set_model_info(cls, values):
+        dict_fields = [
+            "metadata",
+            "aliases",
+            "config",
+            "permissions",
+            "model_max_budget",
+        ]
+        for field in dict_fields:
+            value = values.get(field)
+            if value is not None and isinstance(value, str):
+                try:
+                    values[field] = json.loads(value)
+                except json.JSONDecodeError:
+                    raise ValueError(f"Field {field} should be a valid dictionary")
+
+        return values
+
 
 class TeamRequest(LiteLLMBase):
     teams: List[str]

litellm/proxy/proxy_server.py

@@ -4063,6 +4063,7 @@ async def user_info(
         default=False,
         description="set to true to View all users. When using view_all, don't pass user_id",
     ),
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
     Use this to get user information. (user row + all user key info)
@@ -4108,6 +4109,22 @@ async def user_info(
                 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)
+        elif user_api_key_dict.user_id is not None:
+            caller_user_info = await prisma_client.get_data(
+                user_id=user_api_key_dict.user_id
+            )
+            # *NEW* get all teams in user 'teams' field
+            teams_2 = await prisma_client.get_data(
+                team_id_list=caller_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:
@@ -4137,12 +4154,14 @@ async def user_info(
                 # if using pydantic v1
                 key = key.dict()
             key.pop("token", None)
-        return {
+
+        response_data = {
             "user_id": user_id,
             "user_info": user_info,
             "keys": keys,
             "teams": team_list,
         }
+        return response_data
     except Exception as e:
         traceback.print_exc()
         if isinstance(e, HTTPException):
@@ -4406,6 +4425,17 @@ async def new_team(
                     },
                 )
 
+    if user_api_key_dict.user_id is not None:
+        creating_user_in_list = False
+        for member in data.members_with_roles:
+            if member.user_id == user_api_key_dict.user_id:
+                creating_user_in_list = True
+
+        if creating_user_in_list == False:
+            data.members_with_roles.append(
+                Member(role="admin", user_id=user_api_key_dict.user_id)
+            )
+
     complete_team_data = LiteLLM_TeamTable(
         **data.json(),
         max_parallel_requests=user_api_key_dict.max_parallel_requests,
@@ -4440,6 +4470,9 @@ async def update_team(
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
+    [BETA]
+    [DEPRECATED] - use the `/team/member_add` and `/team/member_remove` endpoints instead 
+
     You can now add / delete users from a team via /team/update
 
     ```
@@ -4479,6 +4512,7 @@ async def update_team(
     existing_user_id_list = []
     ## Get new users
     for user in existing_team_row.members_with_roles:
+        if user["user_id"] is not None:
             existing_user_id_list.append(user["user_id"])
 
     ## Update new user rows with team id (info used by /user/info to show all teams, user is a part of)
@@ -4487,26 +4521,32 @@ async def update_team(
             if user.user_id not in existing_user_id_list:
                 await prisma_client.update_data(
                     user_id=user.user_id,
-                    data={"user_id": user.user_id, "teams": [team_row["team_id"]]},
+                    data={
+                        "user_id": user.user_id,
+                        "teams": [team_row["team_id"]],
+                        "models": team_row["data"].models,
+                    },
                     update_key_values={
                         "teams": {
                             "push": [team_row["team_id"]],
                         }
                     },
+                    table_name="user",
                 )
 
     ## REMOVE DELETED USERS ##
     ### Get list of deleted users (old list - new list)
     deleted_user_id_list = []
-    existing_user_id_list = []
+    new_user_id_list = []
     ## Get old user list
-    for user in existing_team_row.members_with_roles:
-        existing_user_id_list.append(user["user_id"])
-    ## Get diff
     if data.members_with_roles is not None:
         for user in data.members_with_roles:
-            if user.user_id not in existing_user_id_list:
+            new_user_id_list.append(user.user_id)
-                deleted_user_id_list.append(user.user_id)
+    ## Get diff
+    if existing_team_row.members_with_roles is not None:
+        for user in existing_team_row.members_with_roles:
+            if user["user_id"] not in new_user_id_list:
+                deleted_user_id_list.append(user["user_id"])
 
     ## SET UPDATED LIST
     if len(deleted_user_id_list) > 0:
@@ -4525,6 +4565,99 @@ async def update_team(
     return team_row
 
 
+@router.post(
+    "/team/member_add",
+    tags=["team management"],
+    dependencies=[Depends(user_api_key_auth)],
+)
+async def team_member_add(
+    data: TeamMemberAddRequest,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """ 
+    [BETA]
+
+    Add new members (either via user_email or user_id) to a team
+
+    If user doesn't exist, new user row will also be added to User Table
+
+    ```
+    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.member is None:
+        raise HTTPException(status_code=400, detail={"error": "No member passed in"})
+
+    existing_team_row = await prisma_client.get_data(  # type: ignore
+        team_id=data.team_id, table_name="team", query_type="find_unique"
+    )
+
+    new_member = data.member
+
+    existing_team_row.members_with_roles.append(new_member)
+
+    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,
+    )
+
+    ## ADD USER, IF NEW ##
+    user_data = {  # type: ignore
+        "teams": [team_row["team_id"]],
+        "models": team_row["data"].models,
+    }
+    if new_member.user_id is not None:
+        user_data["user_id"] = new_member.user_id  # type: ignore
+        await prisma_client.update_data(
+            user_id=new_member.user_id,
+            data=user_data,
+            update_key_values={
+                "teams": {
+                    "push": [team_row["team_id"]],
+                }
+            },
+            table_name="user",
+        )
+    elif new_member.user_email is not None:
+        user_data["user_id"] = str(uuid.uuid4())
+        user_data["user_email"] = new_member.user_email
+        ## user email is not unique acc. to prisma schema -> future improvement
+        ### for now: check if it exists in db, if not - insert it
+        existing_user_row = await prisma_client.get_data(
+            key_val={"user_email": new_member.user_email},
+            table_name="user",
+            query_type="find_all",
+        )
+        if existing_user_row is None or (
+            isinstance(existing_user_row, list) and len(existing_user_row) == 0
+        ):
+
+            await prisma_client.insert_data(data=user_data, table_name="user")
+
+    return team_row
+
+
 @router.post(
     "/team/delete", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
 )

litellm/proxy/utils.py

@@ -635,11 +635,15 @@ class PrismaClient:
                 table_name is not None and table_name == "user"
             ):
                 if query_type == "find_unique":
+                    if key_val is None:
+                        key_val = {"user_id": user_id}
                     response = await self.db.litellm_usertable.find_unique(  # type: ignore
-                        where={
+                        where=key_val  # type: ignore
-                            "user_id": user_id,  # type: ignore
-                        }
                     )
+                elif query_type == "find_all" and key_val is not None:
+                    response = await self.db.litellm_usertable.find_many(
+                        where=key_val  # type: ignore
+                    )  # type: ignore
                 elif query_type == "find_all" and reset_at is not None:
                     response = await self.db.litellm_usertable.find_many(
                         where={  # type:ignore
@@ -875,6 +879,8 @@ class PrismaClient:
         """
         try:
             db_data = self.jsonify_object(data=data)
+            if update_key_values is not None:
+                update_key_values = self.jsonify_object(data=update_key_values)
             if token is not None:
                 print_verbose(f"token: {token}")
                 # check if plain text or hash

ui/litellm-dashboard/src/app/page.tsx

@@ -14,7 +14,7 @@ import { jwtDecode } from "jwt-decode";
 const CreateKeyPage = () => {
   const [userRole, setUserRole] = useState("");
   const [userEmail, setUserEmail] = useState<null | string>(null);
-  const [teams, setTeams] = useState<null | string[]>(null);
+  const [teams, setTeams] = useState<null | any[]>(null);
   const searchParams = useSearchParams();
 
   const userID = searchParams.get("userID");
@@ -113,7 +113,12 @@ const CreateKeyPage = () => {
               accessToken={accessToken}
             />
           ) : page == "teams" ? (
-            <Teams teams={teams} searchParams={searchParams} />
+            <Teams
+              teams={teams}
+              setTeams={setTeams}
+              searchParams={searchParams}
+              accessToken={accessToken}
+            />
           ) : (
             <Usage
               userID={userID}

ui/litellm-dashboard/src/components/create_key_button.tsx

@@ -18,6 +18,7 @@ const { Option } = Select;
 
 interface CreateKeyProps {
   userID: string;
+  teamID: string | null;
   userRole: string | null;
   accessToken: string;
   data: any[] | null;
@@ -27,6 +28,7 @@ interface CreateKeyProps {
 
 const CreateKey: React.FC<CreateKeyProps> = ({
   userID,
+  teamID,
   userRole,
   accessToken,
   data,
@@ -36,7 +38,6 @@ const CreateKey: React.FC<CreateKeyProps> = ({
   const [form] = Form.useForm();
   const [isModalVisible, setIsModalVisible] = useState(false);
   const [apiKey, setApiKey] = useState(null);
-
   const handleOk = () => {
     setIsModalVisible(false);
     form.resetFields();
@@ -89,7 +90,10 @@ const CreateKey: React.FC<CreateKeyProps> = ({
                 <Input />
               </Form.Item>
               <Form.Item label="Team ID" name="team_id">
-                <Input placeholder="ai_team" />
+                <Input
+                  placeholder="ai_team"
+                  defaultValue={teamID ? teamID : ""}
+                />
               </Form.Item>
               <Form.Item label="Models" name="models">
                 <Select

ui/litellm-dashboard/src/components/dashboard_default_team.tsx

@@ -3,29 +3,33 @@ import { Typography } from "antd";
 import { Select, SelectItem } from "@tremor/react";
 
 interface DashboardTeamProps {
-  teams: string[] | null;
+  teams: Object[] | null;
+  setSelectedTeam: React.Dispatch<React.SetStateAction<any | null>>;
 }
 
-const DashboardTeam: React.FC<DashboardTeamProps> = ({ teams }) => {
+const DashboardTeam: React.FC<DashboardTeamProps> = ({
+  teams,
+  setSelectedTeam,
+}) => {
   const { Title, Paragraph } = Typography;
   const [value, setValue] = useState("");
-  console.log(`received teams ${teams}`);
+
   return (
     <div className="mt-10">
       <Title level={4}>Default Team</Title>
       <Paragraph>
-        If you belong to multiple teams, this setting controls which
+        If you belong to multiple teams, this setting controls which team is
-        organization is used by default when creating new API Keys.
+        used by default when creating new API Keys.
       </Paragraph>
       {teams && teams.length > 0 ? (
-        <Select
+        <Select defaultValue="0">
-          id="distance"
+          {teams.map((team: any, index) => (
-          value={value}
+            <SelectItem
-          onValueChange={setValue}
+              value={String(index)}
-          className="mt-2"
+              onClick={() => setSelectedTeam(team)}
             >
-          {teams.map((model) => (
+              {team["team_alias"]}
-            <SelectItem value="model">{model}</SelectItem>
+            </SelectItem>
           ))}
         </Select>
       ) : (

ui/litellm-dashboard/src/components/leftnav.tsx

@@ -41,9 +41,11 @@ const Sidebar: React.FC<SidebarProps> = ({
               Users
             </Menu.Item>
           ) : null}
+          {userRole == "Admin" ? (
             <Menu.Item key="6" onClick={() => setPage("teams")}>
               Teams
             </Menu.Item>
+          ) : null}
         </Menu>
       </Sider>
     </Layout>

ui/litellm-dashboard/src/components/networking.tsx

@@ -196,6 +196,7 @@ export const userInfoCall = async (
     }
 
     const data = await response.json();
+    console.log("API Response:", data);
     message.info("Received user data");
     return data;
     // Handle success - you might want to update some state or UI based on the created key
@@ -503,14 +504,23 @@ export const teamCreateCall = async (
   }
 };
 
-export const teamUpdateCall = async (
+export interface Member {
+  role: string;
+  user_id: string | null;
+  user_email: string | null;
+}
+
+export const teamMemberAddCall = async (
   accessToken: string,
-  formValues: Record<string, any> // Assuming formValues is an object
+  teamId: string,
+  formValues: Member // Assuming formValues is an object
 ) => {
   try {
-    console.log("Form Values in teamCreateCall:", formValues); // Log the form values before making the API call
+    console.log("Form Values in teamMemberAddCall:", formValues); // Log the form values before making the API call
 
-    const url = proxyBaseUrl ? `${proxyBaseUrl}/team/update` : `/team/update`;
+    const url = proxyBaseUrl
+      ? `${proxyBaseUrl}/team/member_add`
+      : `/team/member_add`;
     const response = await fetch(url, {
       method: "POST",
       headers: {
@@ -518,7 +528,8 @@ export const teamUpdateCall = async (
         "Content-Type": "application/json",
       },
       body: JSON.stringify({
-        ...formValues, // Include formValues in the request body
+        team_id: teamId,
+        member: formValues, // Include formValues in the request body
       }),
     });
 

ui/litellm-dashboard/src/components/teams.tsx

@@ -22,19 +22,32 @@ import {
   Icon,
   Button,
   Col,
+  Text,
   Grid,
 } from "@tremor/react";
 import { CogIcon } from "@heroicons/react/outline";
 interface TeamProps {
-  teams: string[] | null;
+  teams: any[] | null;
   searchParams: any;
+  accessToken: string | null;
+  setTeams: React.Dispatch<React.SetStateAction<Object[] | null>>;
 }
+import { teamCreateCall, teamMemberAddCall, Member } from "./networking";
 
-const Team: React.FC<TeamProps> = ({ teams, searchParams }) => {
+const Team: React.FC<TeamProps> = ({
+  teams,
+  searchParams,
+  accessToken,
+  setTeams,
+}) => {
   const [form] = Form.useForm();
   const [memberForm] = Form.useForm();
   const { Title, Paragraph } = Typography;
   const [value, setValue] = useState("");
+
+  const [selectedTeam, setSelectedTeam] = useState<null | any>(
+    teams ? teams[0] : null
+  );
   const [isTeamModalVisible, setIsTeamModalVisible] = useState(false);
   const [isAddMemberModalVisible, setIsAddMemberModalVisible] = useState(false);
   const handleOk = () => {
@@ -59,7 +72,17 @@ const Team: React.FC<TeamProps> = ({ teams, searchParams }) => {
 
   const handleCreate = async (formValues: Record<string, any>) => {
     try {
-      console.log("reaches here");
+      if (accessToken != null) {
+        message.info("Making API Call");
+        const response: any = await teamCreateCall(accessToken, formValues);
+        if (teams !== null) {
+          setTeams([...teams, response]);
+        } else {
+          setTeams([response]);
+        }
+        console.log(`response for team create call: ${response}`);
+        setIsTeamModalVisible(false);
+      }
     } catch (error) {
       console.error("Error creating the key:", error);
     }
@@ -67,7 +90,36 @@ const Team: React.FC<TeamProps> = ({ teams, searchParams }) => {
 
   const handleMemberCreate = async (formValues: Record<string, any>) => {
     try {
-      console.log("reaches here");
+      if (accessToken != null && teams != null) {
+        message.info("Making API Call");
+        const user_role: Member = {
+          role: "user",
+          user_email: formValues.user_email,
+          user_id: null,
+        };
+        const response: any = await teamMemberAddCall(
+          accessToken,
+          selectedTeam["team_id"],
+          user_role
+        );
+        console.log(`response for team create call: ${response["data"]}`);
+        // Checking if the team exists in the list and updating or adding accordingly
+        const foundIndex = teams.findIndex((team) => {
+          console.log(
+            `team.team_id=${team.team_id}; response.data.team_id=${response.data.team_id}`
+          );
+          return team.team_id === response.data.team_id;
+        });
+        console.log(`foundIndex: ${foundIndex}`);
+        if (foundIndex !== -1) {
+          // If the team is found, update it
+          const updatedTeams = [...teams]; // Copy the current state
+          updatedTeams[foundIndex] = response.data; // Update the specific team
+          setTeams(updatedTeams); // Set the new state
+          setSelectedTeam(response.data);
+        }
+        setIsAddMemberModalVisible(false);
+      }
     } catch (error) {
       console.error("Error creating the key:", error);
     }
@@ -90,24 +142,28 @@ const Team: React.FC<TeamProps> = ({ teams, searchParams }) => {
               </TableHead>
 
               <TableBody>
+                {teams && teams.length > 0
+                  ? teams.map((team: any) => (
                       <TableRow>
-                  <TableCell>Wilhelm Tell</TableCell>
+                        <TableCell>{team["team_alias"]}</TableCell>
-                  <TableCell className="text-right">1</TableCell>
+                        <TableCell>{team["spend"]}</TableCell>
-                  <TableCell>Uri, Schwyz, Unterwalden</TableCell>
+                        <TableCell>
-                  <TableCell>National Hero</TableCell>
+                          {team["max_budget"] ? team["max_budget"] : "No limit"}
-                </TableRow>
+                        </TableCell>
-                <TableRow>
+                        <TableCell>
-                  <TableCell>The Witcher</TableCell>
+                          <Text>
-                  <TableCell className="text-right">129</TableCell>
+                            TPM Limit:{" "}
-                  <TableCell>Kaedwen</TableCell>
+                            {team.tpm_limit ? team.tpm_limit : "Unlimited"}{" "}
-                  <TableCell>Legend</TableCell>
+                            <br></br> RPM Limit:{" "}
-                </TableRow>
+                            {team.rpm_limit ? team.rpm_limit : "Unlimited"}
-                <TableRow>
+                          </Text>
-                  <TableCell>Mizutsune</TableCell>
+                        </TableCell>
-                  <TableCell className="text-right">82</TableCell>
+                        <TableCell>
-                  <TableCell>Japan</TableCell>
+                          <Icon icon={CogIcon} size="sm" />
-                  <TableCell>N/A</TableCell>
+                        </TableCell>
                       </TableRow>
+                    ))
+                  : null}
               </TableBody>
             </Table>
           </Card>
@@ -180,14 +236,16 @@ const Team: React.FC<TeamProps> = ({ teams, searchParams }) => {
             members you see.
           </Paragraph>
           {teams && teams.length > 0 ? (
-            <Select
+            <Select defaultValue="0">
-              id="distance"
+              {teams.map((team: any, index) => (
-              value={value}
+                <SelectItem
-              onValueChange={setValue}
+                  value={String(index)}
-              className="mt-2"
+                  onClick={() => {
+                    setSelectedTeam(team);
+                  }}
                 >
-              {teams.map((model) => (
+                  {team["team_alias"]}
-                <SelectItem value="model">{model}</SelectItem>
+                </SelectItem>
               ))}
             </Select>
           ) : (
@@ -208,27 +266,23 @@ const Team: React.FC<TeamProps> = ({ teams, searchParams }) => {
               </TableHead>
 
               <TableBody>
+                {selectedTeam
+                  ? selectedTeam["members_with_roles"].map((member: any) => (
                       <TableRow>
-                  <TableCell>Wilhelm Tell</TableCell>
+                        <TableCell>
-                  <TableCell>Uri, Schwyz, Unterwalden</TableCell>
+                          {member["user_email"]
-                  <TableCell>
+                            ? member["user_email"]
-                    <Icon icon={CogIcon} size="sm" />
+                            : member["user_id"]
-                  </TableCell>
+                            ? member["user_id"]
-                </TableRow>
+                            : null}
-                <TableRow>
+                        </TableCell>
-                  <TableCell>The Witcher</TableCell>
+                        <TableCell>{member["role"]}</TableCell>
-                  <TableCell>Kaedwen</TableCell>
-                  <TableCell>
-                    <Icon icon={CogIcon} size="sm" />
-                  </TableCell>
-                </TableRow>
-                <TableRow>
-                  <TableCell>Mizutsune</TableCell>
-                  <TableCell>Japan</TableCell>
                         <TableCell>
                           <Icon icon={CogIcon} size="sm" />
                         </TableCell>
                       </TableRow>
+                    ))
+                  : null}
               </TableBody>
             </Table>
           </Card>

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -25,10 +25,10 @@ interface UserDashboardProps {
   userID: string | null;
   userRole: string | null;
   userEmail: string | null;
-  teams: string[] | null;
+  teams: any[] | null;
   setUserRole: React.Dispatch<React.SetStateAction<string>>;
   setUserEmail: React.Dispatch<React.SetStateAction<string | null>>;
-  setTeams: React.Dispatch<React.SetStateAction<string[] | null>>;
+  setTeams: React.Dispatch<React.SetStateAction<Object[] | null>>;
 }
 
 const UserDashboard: React.FC<UserDashboardProps> = ({
@@ -53,6 +53,9 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
   const token = searchParams.get("token");
   const [accessToken, setAccessToken] = useState<string | null>(null);
   const [userModels, setUserModels] = useState<string[]>([]);
+  const [selectedTeam, setSelectedTeam] = useState<any | null>(
+    teams ? teams[0] : null
+  );
   // check if window is not undefined
   if (typeof window !== "undefined") {
     window.addEventListener("beforeunload", function () {
@@ -119,7 +122,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
             console.log(
               `received teams in user dashboard: ${Object.keys(
                 response
-              )}; team type: ${Array.isArray(response.teams)}`
+              )}; team values: ${Object.entries(response.teams)}`
             );
             setUserSpendData(response["user_info"]);
             setData(response["keys"]); // Assuming this is the correct path to your data
@@ -196,13 +199,14 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
           />
           <CreateKey
             userID={userID}
+            teamID={selectedTeam ? selectedTeam["team_id"] : null}
             userRole={userRole}
             userModels={userModels}
             accessToken={accessToken}
             data={data}
             setData={setData}
           />
-          <DashboardTeam teams={teams} />
+          <DashboardTeam teams={teams} setSelectedTeam={setSelectedTeam} />
         </Col>
       </Grid>
     </div>