feat: e2e flow complete - admin can invite new users to proxy via invite links

Completes https://github.com/BerriAI/litellm/issues/3863
This commit is contained in:
Krrish Dholakia 2024-05-27 23:19:30 -07:00
parent 073bca78d4
commit 293d5cf1f2
6 changed files with 614 additions and 242 deletions

View file

@ -7737,12 +7737,12 @@ async def block_user(data: BlockUsers):
"/end_user/unblock",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
include_in_schema=False,
)
@router.post(
"/customer/unblock",
tags=["Customer Management"],
dependencies=[Depends(user_api_key_auth)],
include_in_schema=False,
)
async def unblock_user(data: BlockUsers):
"""
@ -10718,6 +10718,90 @@ async def login(request: Request):
)
@app.get("/onboarding/{invite_link}", include_in_schema=False)
async def onboarding(invite_link: str):
"""
- Get the invite link
- Validate it's still 'valid'
- Invalidate the link (prevents abuse)
- Get user from db
- Pass in user_email if set
"""
global prisma_client
### VALIDATE INVITE LINK ###
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
invite_obj = await prisma_client.db.litellm_invitationlink.find_unique(
where={"id": invite_link}
)
if invite_obj is None:
raise HTTPException(
status_code=401, detail={"error": "Invitation link does not exist in db."}
)
#### CHECK IF EXPIRED
if invite_obj.expires_at < litellm.utils.get_utc_datetime():
raise HTTPException(
status_code=401, detail={"error": "Invitation link has expired."}
)
#### INVALIDATE LINK
current_time = litellm.utils.get_utc_datetime()
_ = await prisma_client.db.litellm_invitationlink.update(
where={"id": invite_link},
data={
"accepted_at": current_time,
"updated_at": current_time,
"is_accepted": True,
"updated_by": invite_obj.user_id, # type: ignore
},
)
### GET USER OBJECT ###
user_obj = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": invite_obj.user_id}
)
if user_obj is None:
raise HTTPException(
status_code=401, detail={"error": "User does not exist in db."}
)
user_email = user_obj.user_email
response = await generate_key_helper_fn(
**{"user_role": user_obj.user_role or "app_owner", "duration": "2hr", "key_max_budget": 5, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_obj.user_id, "team_id": "litellm-dashboard"} # type: ignore
)
key = response["token"] # type: ignore
litellm_dashboard_ui = os.getenv("PROXY_BASE_URL", "")
if litellm_dashboard_ui.endswith("/"):
litellm_dashboard_ui += "ui/onboarding"
else:
litellm_dashboard_ui += "/ui/onboarding"
import jwt
jwt_token = jwt.encode(
{
"user_id": user_obj.user_id,
"key": key,
"user_email": user_obj.user_email,
"user_role": "app_owner",
"login_method": "username_password",
"premium_user": premium_user,
},
"secret",
algorithm="HS256",
)
litellm_dashboard_ui += "?token={}&user_email={}".format(jwt_token, user_email)
return RedirectResponse(url=litellm_dashboard_ui, status_code=303)
@app.get("/get_image", include_in_schema=False)
def get_image():
"""Get logo to show on admin UI"""

View file

@ -0,0 +1,92 @@
"use client";
import React, { Suspense, useEffect, useState } from "react";
import { useSearchParams, useRouter } from "next/navigation";
import { Card, Title, Text, TextInput } from "@tremor/react";
import {
invitationClaimCall,
userUpdateUserCall,
} from "@/components/networking";
import { jwtDecode } from "jwt-decode";
import { Form, Button as Button2, message } from "antd";
export default function Onboarding() {
const [form] = Form.useForm();
const searchParams = useSearchParams();
const token = searchParams.get("token");
const user_email = searchParams.get("user_email");
const [accessToken, setAccessToken] = useState<string | null>(null);
const [defaultUserEmail, setDefaultUserEmail] = useState<string>("");
const router = useRouter();
useEffect(() => {
if (!token) {
return;
}
const decoded = jwtDecode(token) as { [key: string]: any };
setAccessToken(decoded.key);
}, [token]);
useEffect(() => {
if (!user_email) {
return;
}
setDefaultUserEmail(user_email);
}, [user_email]);
const handleSubmit = (formValues: Record<string, any>) => {
if (!accessToken || !token) {
return;
}
userUpdateUserCall(accessToken, formValues, null).then((data) => {
let litellm_dashboard_ui = "";
const user_id = data.data?.user_id || data.user_id;
litellm_dashboard_ui += "?userID=" + user_id + "&token=" + token;
router.replace(litellm_dashboard_ui);
});
};
return (
<div className="mx-auto max-w-md mt-10">
<Card>
<Title className="text-sm mb-5 text-center">🚅 LiteLLM</Title>
<Title className="text-xl">Sign up</Title>
<Text>Claim your user account to login to Admin UI.</Text>
<Form
className="mt-10 mb-5 mx-auto"
layout="vertical"
onFinish={handleSubmit}
>
<>
<Form.Item
label="Email Address"
name="user_email"
rules={[{ required: true, message: "Set the user email" }]}
help="required"
>
<TextInput
placeholder={defaultUserEmail}
type="email"
className="max-w-md"
/>
</Form.Item>
<Form.Item
label="Password"
name="password"
rules={[{ required: true, message: "Set the user password" }]}
help="required"
>
<TextInput placeholder="" type="password" className="max-w-md" />
</Form.Item>
</>
<div className="mt-10">
<Button2 htmlType="submit">Sign Up</Button2>
</div>
</Form>
</Card>
</div>
);
}

View file

@ -4,6 +4,7 @@
*/
import React, { useState, useEffect } from "react";
import { Typography } from "antd";
import { useRouter } from "next/navigation";
import {
Button as Button2,
Modal,
@ -13,6 +14,7 @@ import {
InputNumber,
message,
} from "antd";
import { CopyToClipboard } from "react-copy-to-clipboard";
import { Select, SelectItem } from "@tremor/react";
import {
Table,
@ -28,6 +30,7 @@ import {
Text,
Grid,
Callout,
Divider,
} from "@tremor/react";
import { PencilAltIcon } from "@heroicons/react/outline";
interface AdminPanelProps {
@ -36,41 +39,60 @@ interface AdminPanelProps {
setTeams: React.Dispatch<React.SetStateAction<Object[] | null>>;
showSSOBanner: boolean;
}
interface InvitationLink {
id: string;
user_id: string;
is_accepted: boolean;
accepted_at: Date | null;
expires_at: Date;
created_at: Date;
created_by: string;
updated_at: Date;
updated_by: string;
}
import {
userUpdateUserCall,
Member,
userGetAllUsersCall,
User,
setCallbacksCall,
invitationCreateCall,
} from "./networking";
const AdminPanel: React.FC<AdminPanelProps> = ({
searchParams,
accessToken,
showSSOBanner
showSSOBanner,
}) => {
const [form] = Form.useForm();
const [memberForm] = Form.useForm();
const { Title, Paragraph } = Typography;
const [value, setValue] = useState("");
const [admins, setAdmins] = useState<null | any[]>(null);
const [invitationLinkData, setInvitationLinkData] =
useState<InvitationLink | null>(null);
const [isInvitationLinkModalVisible, setInvitationLinkModalVisible] =
useState(false);
const [isAddMemberModalVisible, setIsAddMemberModalVisible] = useState(false);
const [isAddAdminModalVisible, setIsAddAdminModalVisible] = useState(false);
const [isUpdateMemberModalVisible, setIsUpdateModalModalVisible] = useState(false);
const [isUpdateMemberModalVisible, setIsUpdateModalModalVisible] =
useState(false);
const [isAddSSOModalVisible, setIsAddSSOModalVisible] = useState(false);
const [isInstructionsModalVisible, setIsInstructionsModalVisible] = useState(false);
const [isInstructionsModalVisible, setIsInstructionsModalVisible] =
useState(false);
const router = useRouter();
const [baseUrl, setBaseUrl] = useState("");
let nonSssoUrl;
try {
nonSssoUrl = window.location.origin;
} catch (error) {
nonSssoUrl = '<your-proxy-url>';
nonSssoUrl = "<your-proxy-url>";
}
nonSssoUrl += '/fallback/login';
nonSssoUrl += "/fallback/login";
const handleAddSSOOk = () => {
setIsAddSSOModalVisible(false);
form.resetFields();
};
@ -96,7 +118,15 @@ const handleInstructionsCancel = () => {
setIsInstructionsModalVisible(false);
};
const roles = ["proxy_admin", "proxy_admin_viewer"]
const roles = ["proxy_admin", "proxy_admin_viewer"];
useEffect(() => {
if (router) {
const { protocol, host } = window.location;
const baseUrl = `${protocol}//${host}`;
setBaseUrl(baseUrl);
}
}, [router]);
useEffect(() => {
// Fetch model info and set the default selected model
@ -167,12 +197,13 @@ const handleInstructionsCancel = () => {
const handleMemberUpdateCancel = () => {
setIsUpdateModalModalVisible(false);
memberForm.resetFields();
}
};
// Define the type for the handleMemberCreate function
type HandleMemberCreate = (formValues: Record<string, any>) => Promise<void>;
const addMemberForm = (handleMemberCreate: HandleMemberCreate,) => {
return <Form
const addMemberForm = (handleMemberCreate: HandleMemberCreate) => {
return (
<Form
form={form}
onFinish={handleMemberCreate}
labelCol={{ span: 8 }}
@ -198,10 +229,16 @@ const handleInstructionsCancel = () => {
<Button2 htmlType="submit">Add member</Button2>
</div>
</Form>
}
);
};
const modifyMemberForm = (handleMemberUpdate: HandleMemberCreate, currentRole: string, userID: string) => {
return <Form
const modifyMemberForm = (
handleMemberUpdate: HandleMemberCreate,
currentRole: string,
userID: string
) => {
return (
<Form
form={form}
onFinish={handleMemberUpdate}
labelCol={{ span: 8 }}
@ -209,13 +246,16 @@ const handleInstructionsCancel = () => {
labelAlign="left"
>
<>
<Form.Item rules={[{ required: true, message: 'Required' }]} label="User Role" name="user_role" labelCol={{ span: 10 }} labelAlign="left">
<Form.Item
rules={[{ required: true, message: "Required" }]}
label="User Role"
name="user_role"
labelCol={{ span: 10 }}
labelAlign="left"
>
<Select value={currentRole}>
{roles.map((role, index) => (
<SelectItem
key={index}
value={role}
>
<SelectItem key={index} value={role}>
{role}
</SelectItem>
))}
@ -236,13 +276,18 @@ const handleInstructionsCancel = () => {
<Button2 htmlType="submit">Update role</Button2>
</div>
</Form>
}
);
};
const handleMemberUpdate = async (formValues: Record<string, any>) => {
try {
if (accessToken != null && admins != null) {
message.info("Making API Call");
const response: any = await userUpdateUserCall(accessToken, formValues, null);
const response: any = await userUpdateUserCall(
accessToken,
formValues,
null
);
console.log(`response for team create call: ${response}`);
// Checking if the team exists in the list and updating or adding accordingly
const foundIndex = admins.findIndex((user) => {
@ -258,19 +303,23 @@ const handleInstructionsCancel = () => {
// If new user is found, update it
setAdmins(admins); // Set the new state
}
message.success("Refresh tab to see updated user role")
message.success("Refresh tab to see updated user role");
setIsUpdateModalModalVisible(false);
}
} catch (error) {
console.error("Error creating the key:", error);
}
}
};
const handleMemberCreate = async (formValues: Record<string, any>) => {
try {
if (accessToken != null && admins != null) {
message.info("Making API Call");
const response: any = await userUpdateUserCall(accessToken, formValues, "proxy_admin_viewer");
const response: any = await userUpdateUserCall(
accessToken,
formValues,
"proxy_admin_viewer"
);
console.log(`response for team create call: ${response}`);
// Checking if the team exists in the list and updating or adding accordingly
const foundIndex = admins.findIndex((user) => {
@ -301,12 +350,23 @@ const handleInstructionsCancel = () => {
user_email: formValues.user_email,
user_id: formValues.user_id,
};
const response: any = await userUpdateUserCall(accessToken, formValues, "proxy_admin");
const response: any = await userUpdateUserCall(
accessToken,
formValues,
"proxy_admin"
);
// Give admin an invite link for inviting user to proxy
const user_id = response.data?.user_id || response.user_id;
invitationCreateCall(accessToken, user_id).then((data) => {
setInvitationLinkData(data);
setInvitationLinkModalVisible(true);
});
console.log(`response for team create call: ${response}`);
// Checking if the team exists in the list and updating or adding accordingly
const foundIndex = admins.findIndex((user) => {
console.log(
`user.user_id=${user.user_id}; response.user_id=${response.user_id}`
`user.user_id=${user.user_id}; response.user_id=${user_id}`
);
return user.user_id === response.user_id;
});
@ -336,22 +396,24 @@ const handleInstructionsCancel = () => {
},
};
setCallbacksCall(accessToken, payload);
}
};
console.log(`admins: ${admins?.length}`);
return (
<div className="w-full m-2 mt-2 p-8">
<Title level={4}>Admin Access </Title>
<Paragraph>
{
showSSOBanner && <a href="https://docs.litellm.ai/docs/proxy/ui#restrict-ui-access">Requires SSO Setup</a>
}
{showSSOBanner && (
<a href="https://docs.litellm.ai/docs/proxy/ui#restrict-ui-access">
Requires SSO Setup
</a>
)}
<br />
<b>Proxy Admin: </b> Can create keys, teams, users, add models, etc. <br/>
<b>Proxy Admin Viewer: </b>Can just view spend. They cannot create keys, teams or
grant users access to new models.{" "}
<b>Proxy Admin: </b> Can create keys, teams, users, add models, etc.{" "}
<br />
<b>Proxy Admin Viewer: </b>Can just view spend. They cannot create keys,
teams or grant users access to new models.{" "}
</Paragraph>
<Grid numItems={1} className="gap-2 p-2 w-full">
<Col numColSpan={1}>
<Card className="w-full mx-auto flex-auto overflow-y-auto max-h-[50vh]">
<Table>
@ -375,15 +437,24 @@ const handleInstructionsCancel = () => {
</TableCell>
<TableCell>{member["user_role"]}</TableCell>
<TableCell>
<Icon icon={PencilAltIcon} size="sm" onClick={() => setIsUpdateModalModalVisible(true)}/>
<Icon
icon={PencilAltIcon}
size="sm"
onClick={() => setIsUpdateModalModalVisible(true)}
/>
<Modal
title="Update role"
visible={isUpdateMemberModalVisible}
width={800}
footer={null}
onOk={handleMemberUpdateOk}
onCancel={handleMemberUpdateCancel}>
{modifyMemberForm(handleMemberUpdate, member["user_role"], member["user_id"])}
onCancel={handleMemberUpdateCancel}
>
{modifyMemberForm(
handleMemberUpdate,
member["user_role"],
member["user_id"]
)}
</Modal>
</TableCell>
</TableRow>
@ -407,9 +478,43 @@ const handleInstructionsCancel = () => {
width={800}
footer={null}
onOk={handleAdminOk}
onCancel={handleAdminCancel}>
onCancel={handleAdminCancel}
>
{addMemberForm(handleAdminCreate)}
</Modal>
<Modal
title="Invitation Link"
visible={isInvitationLinkModalVisible}
width={600}
footer={null}
onOk={handleAdminOk}
onCancel={handleAdminCancel}
>
{/* {JSON.stringify(invitationLinkData)} */}
<Paragraph>
Copy and send the generated link to onboard this user to the
proxy.
</Paragraph>
<div className="flex justify-between pt-5 pb-2">
<Text className="text-base">User ID</Text>
<Text>{invitationLinkData?.user_id}</Text>
</div>
<div className="flex justify-between pt-5 pb-2">
<Text>Invitation Link</Text>
<Text>
{baseUrl}/onboarding/{invitationLinkData?.id}
</Text>
</div>
<div className="flex justify-end mt-5">
<div></div>
<CopyToClipboard
text={`${baseUrl}/onboarding/${invitationLinkData?.id}`}
onCopy={() => message.success("Copied!")}
>
<Button variant="primary">Copy invitation link</Button>
</CopyToClipboard>
</div>
</Modal>
<Button
className="mb-5"
onClick={() => setIsAddMemberModalVisible(true)}
@ -441,7 +546,6 @@ const handleInstructionsCancel = () => {
onOk={handleAddSSOOk}
onCancel={handleAddSSOCancel}
>
<Form
form={form}
onFinish={handleShowInstructions}
@ -453,14 +557,24 @@ const handleInstructionsCancel = () => {
<Form.Item
label="Admin Email"
name="user_email"
rules={[{ required: true, message: "Please enter the email of the proxy admin" }]}
rules={[
{
required: true,
message: "Please enter the email of the proxy admin",
},
]}
>
<Input />
</Form.Item>
<Form.Item
label="PROXY BASE URL"
name="proxy_base_url"
rules={[{ required: true, message: "Please enter the proxy base url" }]}
rules={[
{
required: true,
message: "Please enter the proxy base url",
},
]}
>
<Input />
</Form.Item>
@ -468,7 +582,12 @@ const handleInstructionsCancel = () => {
<Form.Item
label="GOOGLE CLIENT ID"
name="google_client_id"
rules={[{ required: true, message: "Please enter the google client id" }]}
rules={[
{
required: true,
message: "Please enter the google client id",
},
]}
>
<Input.Password />
</Form.Item>
@ -476,7 +595,12 @@ const handleInstructionsCancel = () => {
<Form.Item
label="GOOGLE CLIENT SECRET"
name="google_client_secret"
rules={[{ required: true, message: "Please enter the google client secret" }]}
rules={[
{
required: true,
message: "Please enter the google client secret",
},
]}
>
<Input.Password />
</Form.Item>
@ -485,7 +609,6 @@ const handleInstructionsCancel = () => {
<Button2 htmlType="submit">Save</Button2>
</div>
</Form>
</Modal>
<Modal
title="SSO Setup Instructions"
@ -496,14 +619,13 @@ const handleInstructionsCancel = () => {
onCancel={handleInstructionsCancel}
>
<p>Follow these steps to complete the SSO setup:</p>
<Text className="mt-2">
1. DO NOT Exit this TAB
</Text>
<Text className="mt-2">1. DO NOT Exit this TAB</Text>
<Text className="mt-2">
2. Open a new tab, visit your proxy base url
</Text>
<Text className="mt-2">
3. Confirm your SSO is configured correctly and you can login on the new Tab
3. Confirm your SSO is configured correctly and you can login on
the new Tab
</Text>
<Text className="mt-2">
4. If Step 3 is successful, you can close this tab
@ -514,7 +636,10 @@ const handleInstructionsCancel = () => {
</Modal>
</div>
<Callout title="Login without SSO" color="teal">
If you need to login without sso, you can access <a href= {nonSssoUrl} target="_blank"><b>{nonSssoUrl}</b> </a>
If you need to login without sso, you can access{" "}
<a href={nonSssoUrl} target="_blank">
<b>{nonSssoUrl}</b>{" "}
</a>
</Callout>
</Grid>
</div>

View file

@ -1200,9 +1200,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
wordBreak: "break-word",
}}
>
<p style={{ fontSize: "10px" }}>
{model.model_name || "-"}
</p>
<p className="text-xs">{model.model_name || "-"}</p>
</TableCell>
<TableCell
style={{
@ -1211,9 +1209,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
wordBreak: "break-word",
}}
>
<p style={{ fontSize: "10px" }}>
{model.provider || "-"}
</p>
<p className="text-xs">{model.provider || "-"}</p>
</TableCell>
{userRole === "Admin" && (
<TableCell
@ -1229,8 +1225,8 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
maxWidth: "150px",
whiteSpace: "normal",
wordBreak: "break-word",
fontSize: "10px",
}}
className="text-xs"
title={
model && model.api_base
? model.api_base
@ -1251,7 +1247,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
wordBreak: "break-word",
}}
>
<pre style={{ fontSize: "10px" }}>
<pre className="text-xs">
{model.input_cost
? model.input_cost
: model.litellm_params.input_cost_per_token
@ -1271,7 +1267,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
wordBreak: "break-word",
}}
>
<pre style={{ fontSize: "10px" }}>
<pre className="text-xs">
{model.output_cost
? model.output_cost
: model.litellm_params.output_cost_per_token
@ -1285,7 +1281,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
</pre>
</TableCell>
<TableCell>
<p style={{ fontSize: "10px" }}>
<p className="text-xs">
{premiumUser
? formatCreatedAt(
model.model_info.created_at
@ -1294,7 +1290,7 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
</p>
</TableCell>
<TableCell>
<p style={{ fontSize: "10px" }}>
<p className="text-xs">
{premiumUser
? model.model_info.created_by || "-"
: "-"}
@ -1309,11 +1305,11 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
>
{model.model_info.db_model ? (
<Badge size="xs" className="text-white">
<p style={{ fontSize: "10px" }}>DB Model</p>
<p className="text-xs">DB Model</p>
</Badge>
) : (
<Badge size="xs" className="text-black">
<p style={{ fontSize: "10px" }}>Config Model</p>
<p className="text-xs">Config Model</p>
</Badge>
)}
</TableCell>

View file

@ -147,7 +147,7 @@ const ModelHub: React.FC<ModelHubProps> = ({
)}
</div>
<div className="grid grid-cols-2 gap-6 sm:grid-cols-3 lg:grid-cols-4">
<div className="grid grid-cols-2 gap-6 sm:grid-cols-3 lg:grid-cols-4 pr-8">
{modelHubData &&
modelHubData.map((model: ModelInfo) => (
<Card key={model.model_group} className="mt-5 mx-8">

View file

@ -208,6 +208,81 @@ export const budgetCreateCall = async (
}
};
export const invitationCreateCall = async (
accessToken: string,
userID: string // Assuming formValues is an object
) => {
try {
const url = proxyBaseUrl
? `${proxyBaseUrl}/invitation/new`
: `/invitation/new`;
const response = await fetch(url, {
method: "POST",
headers: {
Authorization: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({
user_id: userID, // Include formValues in the request body
}),
});
if (!response.ok) {
const errorData = await response.text();
message.error("Failed to create key: " + errorData, 10);
console.error("Error response from the server:", errorData);
throw new Error("Network response was not ok");
}
const data = await response.json();
console.log("API Response:", data);
return data;
// Handle success - you might want to update some state or UI based on the created key
} catch (error) {
console.error("Failed to create key:", error);
throw error;
}
};
export const invitationClaimCall = async (
accessToken: string,
formValues: Record<string, any> // Assuming formValues is an object
) => {
try {
console.log("Form Values in invitationCreateCall:", formValues); // Log the form values before making the API call
console.log("Form Values after check:", formValues);
const url = proxyBaseUrl
? `${proxyBaseUrl}/invitation/claim`
: `/invitation/claim`;
const response = await fetch(url, {
method: "POST",
headers: {
Authorization: `Bearer ${accessToken}`,
"Content-Type": "application/json",
},
body: JSON.stringify({
...formValues, // Include formValues in the request body
}),
});
if (!response.ok) {
const errorData = await response.text();
message.error("Failed to create key: " + errorData, 10);
console.error("Error response from the server:", errorData);
throw new Error("Network response was not ok");
}
const data = await response.json();
console.log("API Response:", data);
return data;
// Handle success - you might want to update some state or UI based on the created key
} catch (error) {
console.error("Failed to create key:", error);
throw error;
}
};
export const alertingSettingsCall = async (accessToken: String) => {
/**
* Get all configurable params for setting a model