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

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

Browse source
This commit is contained in:
Krrish Dholakia 2024-02-24 22:06:00 -08:00
parent 1151bc268f
commit c33a472611
10 changed files with 347 additions and 93 deletions
Download patch file Download diff file Expand all files Collapse all files

33
litellm/proxy/_types.py
View file

@ -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]

149
litellm/proxy/proxy_server.py
View file

@ -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:
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:
@ -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)]
)

12
litellm/proxy/utils.py
View file

@ -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
@ -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

9
ui/litellm-dashboard/src/app/page.tsx
View file

@ -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}

8
ui/litellm-dashboard/src/components/create_key_button.tsx
View file

@ -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

28
ui/litellm-dashboard/src/components/dashboard_default_team.tsx
View file

@ -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
organization is used by default when creating new API Keys.
If you belong to multiple teams, this setting controls which team is
used by default when creating new API Keys.
</Paragraph>
{teams && teams.length > 0 ? (
<Select
id="distance"
value={value}
onValueChange={setValue}
className="mt-2"
<Select defaultValue="0">
{teams.map((team: any, index) => (
<SelectItem
value={String(index)}
onClick={() => setSelectedTeam(team)}
>
{teams.map((model) => (
<SelectItem value="model">{model}</SelectItem>
{team["team_alias"]}
</SelectItem>
))}
</Select>
) : (

2
ui/litellm-dashboard/src/components/leftnav.tsx
View file

@ -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>

21
ui/litellm-dashboard/src/components/networking.tsx
View file

@ -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
}),
});

140
ui/litellm-dashboard/src/components/teams.tsx
View file

@ -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 className="text-right">1</TableCell>
<TableCell>Uri, Schwyz, Unterwalden</TableCell>
<TableCell>National Hero</TableCell>
</TableRow>
<TableRow>
<TableCell>The Witcher</TableCell>
<TableCell className="text-right">129</TableCell>
<TableCell>Kaedwen</TableCell>
<TableCell>Legend</TableCell>
</TableRow>
<TableRow>
<TableCell>Mizutsune</TableCell>
<TableCell className="text-right">82</TableCell>
<TableCell>Japan</TableCell>
<TableCell>N/A</TableCell>
<TableCell>{team["team_alias"]}</TableCell>
<TableCell>{team["spend"]}</TableCell>
<TableCell>
{team["max_budget"] ? team["max_budget"] : "No limit"}
</TableCell>
<TableCell>
<Text>
TPM Limit:{" "}
{team.tpm_limit ? team.tpm_limit : "Unlimited"}{" "}
<br></br> RPM Limit:{" "}
{team.rpm_limit ? team.rpm_limit : "Unlimited"}
</Text>
</TableCell>
<TableCell>
<Icon icon={CogIcon} size="sm" />
</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
id="distance"
value={value}
onValueChange={setValue}
className="mt-2"
<Select defaultValue="0">
{teams.map((team: any, index) => (
<SelectItem
value={String(index)}
onClick={() => {
setSelectedTeam(team);
}}
>
{teams.map((model) => (
<SelectItem value="model">{model}</SelectItem>
{team["team_alias"]}
</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>Uri, Schwyz, Unterwalden</TableCell>
<TableCell>
<Icon icon={CogIcon} size="sm" />
</TableCell>
</TableRow>
<TableRow>
<TableCell>The Witcher</TableCell>
<TableCell>Kaedwen</TableCell>
<TableCell>
<Icon icon={CogIcon} size="sm" />
</TableCell>
</TableRow>
<TableRow>
<TableCell>Mizutsune</TableCell>
<TableCell>Japan</TableCell>
<TableCell>
{member["user_email"]
? member["user_email"]
: member["user_id"]
? member["user_id"]
: null}
</TableCell>
<TableCell>{member["role"]}</TableCell>
<TableCell>
<Icon icon={CogIcon} size="sm" />
</TableCell>
</TableRow>
))
: null}
</TableBody>
</Table>
</Card>

12
ui/litellm-dashboard/src/components/user_dashboard.tsx
View file

@ -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>