diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index 159b311a9..9c8927caf 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -119,6 +119,51 @@ curl --location 'http://0.0.0.0:8000/key/generate' \ --data '{"models": ["azure-models"], "user_id": "krrish3@berri.ai"}' ``` + + +You can: +- Add budgets to Teams + + +#### **Add budgets to users** +```shell +curl --location 'http://localhost:8000/team/new' \ +--header 'Authorization: Bearer ' \ +--header 'Content-Type: application/json' \ +--data-raw '{ + "team_alias": "my-new-team_4", + "members_with_roles": [{"role": "admin", "user_id": "5c4a0aa3-a1e1-43dc-bd87-3c2da8382a3a"}], + "rpm_limit": 99 +}' +``` + +[**See Swagger**](https://litellm-api.up.railway.app/#/team%20management/new_team_team_new_post) + +**Sample Response** + +```shell +{ + "team_alias": "my-new-team_4", + "team_id": "13e83b19-f851-43fe-8e93-f96e21033100", + "admins": [], + "members": [], + "members_with_roles": [ + { + "role": "admin", + "user_id": "5c4a0aa3-a1e1-43dc-bd87-3c2da8382a3a" + } + ], + "metadata": {}, + "tpm_limit": null, + "rpm_limit": 99, + "max_budget": null, + "models": [], + "spend": 0.0, + "max_parallel_requests": null, + "budget_duration": null, + "budget_reset_at": null +} +``` diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 029f35867..778a012b6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -230,7 +230,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): @@ -240,6 +247,15 @@ class NewTeamRequest(LiteLLMBase): members: list = [] members_with_roles: List[Member] = [] metadata: Optional[dict] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + max_budget: Optional[float] = None + models: list = [] + + +class TeamMemberAddRequest(LiteLLMBase): + team_id: str + member: Optional[Member] = None class UpdateTeamRequest(LiteLLMBase): @@ -256,23 +272,29 @@ class DeleteTeamRequest(LiteLLMBase): class LiteLLM_TeamTable(NewTeamRequest): - max_budget: Optional[float] = None spend: Optional[float] = None - models: list = [] max_parallel_requests: Optional[int] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None 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") -class NewTeamResponse(LiteLLMBase): - team_id: str - admins: list - members: list - metadata: dict - created_at: datetime - updated_at: datetime + return values class TeamRequest(LiteLLMBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 939dd244a..56c49ae32 100644 --- a/litellm/proxy/proxy_server.py +++ b/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): @@ -4401,7 +4420,7 @@ async def unblock_user(data: BlockUsers): "/team/new", tags=["team management"], dependencies=[Depends(user_api_key_auth)], - response_model=NewTeamResponse, + response_model=LiteLLM_TeamTable, ) async def new_team( data: NewTeamRequest, @@ -4447,13 +4466,66 @@ async def new_team( if data.team_id is None: data.team_id = str(uuid.uuid4()) + if ( + data.tpm_limit is not None + and user_api_key_dict.tpm_limit is not None + and data.tpm_limit > user_api_key_dict.tpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"tpm limit higher than user max. User tpm limit={user_api_key_dict.tpm_limit}" + }, + ) + + if ( + data.rpm_limit is not None + and user_api_key_dict.rpm_limit is not None + and data.rpm_limit > user_api_key_dict.rpm_limit + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"rpm limit higher than user max. User rpm limit={user_api_key_dict.rpm_limit}" + }, + ) + + if ( + data.max_budget is not None + and user_api_key_dict.max_budget is not None + and data.max_budget > user_api_key_dict.max_budget + ): + raise HTTPException( + status_code=400, + detail={ + "error": f"max budget higher than user max. User max budget={user_api_key_dict.max_budget}" + }, + ) + + if data.models is not None: + for m in data.models: + if m not in user_api_key_dict.models: + raise HTTPException( + status_code=400, + detail={ + "error": f"Model not in allowed user models. User allowed models={user_api_key_dict.models}" + }, + ) + + 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_budget=user_api_key_dict.max_budget, - models=user_api_key_dict.models, max_parallel_requests=user_api_key_dict.max_parallel_requests, - tpm_limit=user_api_key_dict.tpm_limit, - rpm_limit=user_api_key_dict.rpm_limit, budget_duration=user_api_key_dict.budget_duration, budget_reset_at=user_api_key_dict.budget_reset_at, ) @@ -4468,13 +4540,13 @@ async def new_team( await prisma_client.update_data( user_id=user.user_id, data={"user_id": user.user_id, "teams": [team_row.team_id]}, - update_key_values={ + update_key_values_custom_query={ "teams": { "push ": [team_row.team_id], } }, ) - return team_row + return team_row.model_dump() @router.post( @@ -4485,6 +4557,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 ``` @@ -4524,7 +4599,8 @@ async def update_team( existing_user_id_list = [] ## Get new users for user in existing_team_row.members_with_roles: - existing_user_id_list.append(user["user_id"]) + 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) if data.members_with_roles is not None: @@ -4532,26 +4608,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"]]}, - update_key_values={ + data={ + "user_id": user.user_id, + "teams": [team_row["team_id"]], + "models": team_row["data"].models, + }, + update_key_values_custom_query={ "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: - deleted_user_id_list.append(user.user_id) + new_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: @@ -4570,6 +4652,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_custom_query={ + "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)] ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9f9774412..868977885 100644 --- a/litellm/proxy/utils.py +++ b/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={ - "user_id": user_id, # type: ignore - } + where=key_val # 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 @@ -869,12 +873,15 @@ class PrismaClient: query_type: Literal["update", "update_many"] = "update", table_name: Optional[Literal["user", "key", "config", "spend", "team"]] = None, update_key_values: Optional[dict] = None, + update_key_values_custom_query: Optional[dict] = None, ): """ Update existing data """ 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 @@ -902,7 +909,10 @@ class PrismaClient: if user_id is None: user_id = db_data["user_id"] if update_key_values is None: - update_key_values = db_data + if update_key_values_custom_query is not None: + update_key_values = update_key_values_custom_query + else: + update_key_values = db_data update_user_row = await self.db.litellm_usertable.upsert( where={"user_id": user_id}, # type: ignore data={ @@ -944,7 +954,6 @@ class PrismaClient: 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={ diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 424ddfcec..416022ba3 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -5,14 +5,16 @@ import Navbar from "../components/navbar"; import UserDashboard from "../components/user_dashboard"; import ModelDashboard from "@/components/model_dashboard"; import ViewUserDashboard from "@/components/view_users"; +import Teams from "@/components/teams"; import ChatUI from "@/components/chat_ui"; import Sidebar from "../components/leftnav"; import Usage from "../components/usage"; import { jwtDecode } from "jwt-decode"; const CreateKeyPage = () => { - const [userRole, setUserRole] = useState(''); + const [userRole, setUserRole] = useState(""); const [userEmail, setUserEmail] = useState(null); + const [teams, setTeams] = useState(null); const searchParams = useSearchParams(); const userID = searchParams.get("userID"); @@ -74,14 +76,20 @@ const CreateKeyPage = () => {
- + {page == "api-keys" ? ( ) : page == "models" ? ( { token={token} accessToken={accessToken} /> - ) - : page == "users" ? ( + ) : page == "users" ? ( - ) - : ( + ) : page == "teams" ? ( + + ) : ( { + const [userRole, setUserRole] = useState(""); + const [userEmail, setUserEmail] = useState(null); + const [teams, setTeams] = useState(null); + const searchParams = useSearchParams(); + + const userID = searchParams.get("userID"); + const token = searchParams.get("token"); + + const [page, setPage] = useState("team"); + const [accessToken, setAccessToken] = useState(null); + + useEffect(() => { + if (token) { + const decoded = jwtDecode(token) as { [key: string]: any }; + if (decoded) { + // cast decoded to dictionary + console.log("Decoded token:", decoded); + + console.log("Decoded key:", decoded.key); + // set accessToken + setAccessToken(decoded.key); + + // check if userRole is defined + if (decoded.user_role) { + const formattedUserRole = formatUserRole(decoded.user_role); + console.log("Decoded user_role:", formattedUserRole); + setUserRole(formattedUserRole); + } else { + console.log("User role not defined"); + } + + if (decoded.user_email) { + setUserEmail(decoded.user_email); + } else { + console.log(`User Email is not set ${decoded}`); + } + } + } + }, [token]); + + function formatUserRole(userRole: string) { + if (!userRole) { + return "Undefined Role"; + } + console.log(`Received user role: ${userRole}`); + switch (userRole.toLowerCase()) { + case "app_owner": + return "App Owner"; + case "demo_app_owner": + return "App Owner"; + case "app_admin": + return "Admin"; + case "app_user": + return "App User"; + default: + return "Unknown Role"; + } + } + + return ( + Loading...
}> +
+ +
+ + + + + + + Team Name + Spend (USD) + Budget (USD) + TPM / RPM Limits + Settings + + + + + + Wilhelm Tell + 1 + Uri, Schwyz, Unterwalden + National Hero + + + + + + The Witcher + 129 + Kaedwen + Legend + + + + + + Mizutsune + 82 + Japan + N/A + + + + + +
+
+ + + + + + +
+
+
+ + ); +}; +export default TeamSettingsPage; diff --git a/ui/litellm-dashboard/src/components/create_key_button.tsx b/ui/litellm-dashboard/src/components/create_key_button.tsx index 03ef9a89e..ededfec98 100644 --- a/ui/litellm-dashboard/src/components/create_key_button.tsx +++ b/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 = ({ userID, + teamID, userRole, accessToken, data, @@ -36,7 +38,6 @@ const CreateKey: React.FC = ({ 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 = ({ - + + {teams.map((team: any, index) => ( + setSelectedTeam(team)} + > + {team["team_alias"]} + + ))} + + ) : ( + + No team created. Defaulting to personal account. + + )} +
+ ); +}; + +export default DashboardTeam; diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 8a13d2213..a37658f6e 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -1,5 +1,6 @@ import { Layout, Menu } from "antd"; import Link from "next/link"; +import { List } from "postcss/lib/list"; const { Sider } = Layout; @@ -7,15 +8,20 @@ const { Sider } = Layout; interface SidebarProps { setPage: React.Dispatch>; userRole: string; + defaultSelectedKey: string[] | null; } -const Sidebar: React.FC = ({ setPage, userRole }) => { +const Sidebar: React.FC = ({ + setPage, + userRole, + defaultSelectedKey, +}) => { return ( setPage("api-keys")}> @@ -30,13 +36,16 @@ const Sidebar: React.FC = ({ setPage, userRole }) => { setPage("usage")}> Usage - { - userRole == "Admin" ? - setPage("users")}> - Users - - : null - } + {userRole == "Admin" ? ( + setPage("users")}> + Users + + ) : null} + {userRole == "Admin" ? ( + setPage("teams")}> + Teams + + ) : null} diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index 3546bcc5f..bb13d100f 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -29,14 +29,19 @@ const Navbar: React.FC = ({ userID, userRole, userEmail }) => { const isLocal = process.env.NODE_ENV === "development"; const imageUrl = isLocal ? "http://localhost:4000/get_image" : "/get_image"; - return (