litellm/ui/litellm-dashboard/src/components/view_key_table.tsx
2024-04-02 15:27:50 -07:00

627 lines
20 KiB
TypeScript

"use client";
import React, { useEffect, useState } from "react";
import { keyDeleteCall } from "./networking";
import { InformationCircleIcon, StatusOnlineIcon, TrashIcon, PencilAltIcon } from "@heroicons/react/outline";
import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall } from "./networking";
import {
Badge,
Card,
Table,
Button,
TableBody,
TableCell,
TableHead,
TableHeaderCell,
TableRow,
Dialog,
DialogPanel,
Text,
Title,
Icon,
BarChart,
} from "@tremor/react";
import {
Button as Button2,
Modal,
Form,
Input,
Select as Select2,
InputNumber,
message,
Select,
} from "antd";
import ViewKeySpendReport from "./view_key_spend_report";
const { Option } = Select;
interface EditKeyModalProps {
visible: boolean;
onCancel: () => void;
token: any; // Assuming TeamType is a type representing your team object
onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted
}
// Define the props type
interface ViewKeyTableProps {
userID: string;
accessToken: string;
selectedTeam: any | null;
userModels: string[];
data: any[] | null;
setData: React.Dispatch<React.SetStateAction<any[] | null>>;
}
interface ItemData {
key_alias: string | null;
key_name: string;
spend: string;
max_budget: string | null;
models: string[];
tpm_limit: string | null;
rpm_limit: string | null;
token: string;
id: number;
team_id: string;
metadata: any;
expires: any;
// Add any other properties that exist in the item data
}
const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
userID,
accessToken,
selectedTeam,
userModels,
data,
setData,
}) => {
const [isButtonClicked, setIsButtonClicked] = useState(false);
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
const [keyToDelete, setKeyToDelete] = useState<string | null>(null);
const [openDialogId, setOpenDialogId] = React.useState<null | number>(null);
const [selectedItem, setSelectedItem] = useState<ItemData | null>(null);
const [spendData, setSpendData] = useState<{ day: string; spend: number }[] | null>(
null
);
const [predictedSpendString, setPredictedSpendString] = useState("");
const [editModalVisible, setEditModalVisible] = useState(false);
const [selectedToken, setSelectedToken] = useState<ItemData | null>(null);
const EditKeyModal: React.FC<EditKeyModalProps> = ({ visible, onCancel, token, onSubmit }) => {
const [form] = Form.useForm();
// check token.models length == 0
if (token.models.length == 0 && selectedTeam) {
token.models = selectedTeam.models;
}
const handleOk = () => {
form
.validateFields()
.then((values) => {
// const updatedValues = {...values, team_id: team.team_id};
// onSubmit(updatedValues);
form.resetFields();
})
.catch((error) => {
console.error("Validation failed:", error);
});
};
return (
<Modal
title="Edit Key"
visible={visible}
width={800}
footer={null}
onOk={handleOk}
onCancel={onCancel}
>
<Form
form={form}
onFinish={handleEditSubmit}
initialValues={token} // Pass initial values here
labelCol={{ span: 8 }}
wrapperCol={{ span: 16 }}
labelAlign="left"
>
<>
<Form.Item
label="Key Name"
name="key_alias"
rules={[{ required: true, message: 'Please input a key name' }]}
help="required"
>
<Input />
</Form.Item>
<Form.Item label="Models" name="models">
<Select
mode="multiple"
placeholder="Select models"
style={{ width: "100%" }}
>
{selectedTeam && selectedTeam.models ? (
selectedTeam.models.map((model: string) => (
<Option key={model} value={model}>
{model}
</Option>
))
) : (
userModels.map((model: string) => (
<Option key={model} value={model}>
{model}
</Option>
))
)}
</Select>
</Form.Item>
<Form.Item
className="mt-8"
label="Max Budget (USD)"
name="max_budget"
help={`Budget cannot exceed team max budget: $${selectedTeam?.max_budget !== null && selectedTeam?.max_budget !== undefined ? selectedTeam?.max_budget : 'unlimited'}`}
rules={[
{
validator: async (_, value) => {
if (value && selectedTeam && selectedTeam.max_budget !== null && value > selectedTeam.max_budget) {
throw new Error(`Budget cannot exceed team max budget: $${selectedTeam.max_budget}`);
}
},
},
]}
>
<InputNumber step={0.01} precision={2} width={200} />
</Form.Item>
<Form.Item
label="token"
name="token"
hidden={true}
></Form.Item>
</>
<div style={{ textAlign: "right", marginTop: "10px" }}>
<Button2 htmlType="submit">Edit Key</Button2>
</div>
</Form>
</Modal>
);
};
const handleEditClick = (token: any) => {
console.log("handleEditClick:", token);
setSelectedToken(token);
setEditModalVisible(true);
};
const handleEditCancel = () => {
setEditModalVisible(false);
setSelectedToken(null);
};
const handleEditSubmit = async (formValues: Record<string, any>) => {
// Call API to update team with teamId and values
if (accessToken == null) {
return;
}
const currentKey = formValues.token;
formValues.key = currentKey;
console.log("handleEditSubmit:", formValues);
let newKeyValues = await keyUpdateCall(accessToken, formValues);
console.log("handleEditSubmit: newKeyValues", newKeyValues);
// Update the keys with the update key
if (data) {
const updatedData = data.map((key) =>
key.token === currentKey ? newKeyValues : key
);
setData(updatedData);
}
message.success("Key updated successfully");
setEditModalVisible(false);
setSelectedToken(null);
};
// call keySpendLogsCall and set the data
const fetchData = async (item: ItemData | null) => {
try {
if (accessToken == null || item == null) {
return;
}
console.log(`accessToken: ${accessToken}; token: ${item.token}`);
const response = await keySpendLogsCall(accessToken, item.token);
console.log("Response:", response);
setSpendData(response);
// predict spend based on response
try {
const predictedSpend = await PredictedSpendLogsCall(accessToken, response);
console.log("Response2:", predictedSpend);
// append predictedSpend to data
const combinedData = [...response, ...predictedSpend.response];
setSpendData(combinedData);
setPredictedSpendString(predictedSpend.predicted_spend)
console.log("Combined Data:", combinedData);
} catch (error) {
console.error("There was an error fetching the predicted data", error);
}
// setPredictedSpend(predictedSpend);
} catch (error) {
console.error("There was an error fetching the data", error);
}
};
useEffect(() => {
fetchData(selectedItem);
}, [selectedItem]);
const handleDelete = async (token: string) => {
if (data == null) {
return;
}
// Set the key to delete and open the confirmation modal
setKeyToDelete(token);
localStorage.removeItem("userData" + userID);
setIsDeleteModalOpen(true);
};
const confirmDelete = async () => {
if (keyToDelete == null || data == null) {
return;
}
try {
await keyDeleteCall(accessToken, keyToDelete);
// Successfully completed the deletion. Update the state to trigger a rerender.
const filteredData = data.filter((item) => item.token !== keyToDelete);
setData(filteredData);
} catch (error) {
console.error("Error deleting the key:", error);
// Handle any error situations, such as displaying an error message to the user.
}
// Close the confirmation modal and reset the keyToDelete
setIsDeleteModalOpen(false);
setKeyToDelete(null);
};
const cancelDelete = () => {
// Close the confirmation modal and reset the keyToDelete
setIsDeleteModalOpen(false);
setKeyToDelete(null);
};
if (data == null) {
return;
}
// useEffect(() => {
// if (openDialogId !== null && selectedItem !== null) {
// fetchData(selectedItem);
// }
// }, [openDialogId, selectedItem]);
console.log("RERENDER TRIGGERED");
return (
<div>
<Card className="w-full mx-auto flex-auto overflow-y-auto max-h-[50vh] mb-4 mt-2">
<Table className="mt-5">
<TableHead>
<TableRow>
<TableHeaderCell>Key Alias</TableHeaderCell>
<TableHeaderCell>Secret Key</TableHeaderCell>
<TableHeaderCell>Spend (USD)</TableHeaderCell>
<TableHeaderCell>Budget (USD)</TableHeaderCell>
{/* <TableHeaderCell>Spend Report</TableHeaderCell> */}
{/* <TableHeaderCell>Team</TableHeaderCell> */}
{/* <TableHeaderCell>Metadata</TableHeaderCell> */}
<TableHeaderCell>Models</TableHeaderCell>
<TableHeaderCell>TPM / RPM Limits</TableHeaderCell>
{/* <TableHeaderCell>Expires</TableHeaderCell> */}
</TableRow>
</TableHead>
<TableBody>
{data.map((item) => {
console.log(item);
// skip item if item.team_id == "litellm-dashboard"
if (item.team_id === "litellm-dashboard") {
return null;
}
if (selectedTeam) {
if (item.team_id != selectedTeam.team_id) {
return null;
}
}
return (
<TableRow key={item.token}>
<TableCell style={{ maxWidth: "2px", whiteSpace: "pre-wrap", overflow: "hidden" }}>
{item.key_alias != null ? (
<Text>{item.key_alias}</Text>
) : (
<Text>Not Set</Text>
)}
</TableCell>
<TableCell>
<Text>{item.key_name}</Text>
</TableCell>
<TableCell>
<Text>
{(() => {
try {
return parseFloat(item.spend).toFixed(4);
} catch (error) {
return item.spend;
}
})()}
</Text>
</TableCell>
<TableCell>
{item.max_budget != null ? (
<Text>{item.max_budget}</Text>
) : (
<Text>Unlimited</Text>
)}
</TableCell>
{/* <TableCell style={{ maxWidth: '2px' }}>
<ViewKeySpendReport
token={item.token}
accessToken={accessToken}
keySpend={item.spend}
keyBudget={item.max_budget}
keyName={item.key_name}
/>
</TableCell> */}
{/* <TableCell style={{ maxWidth: "4px", whiteSpace: "pre-wrap", overflow: "hidden" }}>
<Text>{item.team_alias && item.team_alias != "None" ? item.team_alias : item.team_id}</Text>
</TableCell> */}
{/* <TableCell style={{ maxWidth: "4px", whiteSpace: "pre-wrap", overflow: "hidden" }}>
<Text>{JSON.stringify(item.metadata).slice(0, 400)}</Text>
</TableCell> */}
<TableCell>
{Array.isArray(item.models) ? (
<div style={{ display: "flex", flexDirection: "column" }}>
{item.models.length === 0 ? (
<>
{selectedTeam && selectedTeam.models && selectedTeam.models.length > 0 ? (
selectedTeam.models.map((model: string, index: number) => (
<Badge key={index} size={"xs"} className="mb-1" color="blue">
<Text>{model.length > 30 ? `${model.slice(0, 30)}...` : model}</Text>
</Badge>
))
) : (
// If selected team is None or selected team's models are empty, show all models
<Badge size={"xs"} className="mb-1" color="blue">
<Text>all-models</Text>
</Badge>
)}
</>
) : (
item.models.map((model: string, index: number) => (
<Badge key={index} size={"xs"} className="mb-1" color="blue">
<Text>{model.length > 30 ? `${model.slice(0, 30)}...` : model}</Text>
</Badge>
))
)}
</div>
) : null}
</TableCell>
<TableCell>
<Text>
TPM: {item.tpm_limit ? item.tpm_limit : "Unlimited"}{" "}
<br></br> RPM:{" "}
{item.rpm_limit ? item.rpm_limit : "Unlimited"}
</Text>
</TableCell>
<TableCell>
<Icon
onClick={() => {
setSelectedItem(item);
setOpenDialogId(item.id);
}}
icon={InformationCircleIcon}
size="sm"
/>
<Dialog
open={openDialogId !== null}
onClose={() => {
setOpenDialogId(null);
setSelectedItem(null);
}}
>
<DialogPanel>
{selectedItem && (
<>
<div className="grid grid-cols-1 gap-6 sm:grid-cols-2 lg:grid-cols-3">
<Card>
<p className="text-tremor-default font-medium text-tremor-content dark:text-dark-tremor-content">
Spend
</p>
<div className="mt-2 flex items-baseline space-x-2.5">
<p className="text-tremor font-semibold text-tremor-content-strong dark:text-dark-tremor-content-strong">
{(() => {
try {
return parseFloat(selectedItem.spend).toFixed(4);
} catch (error) {
return selectedItem.spend;
}
})()}
</p>
</div>
</Card>
<Card key={item.name}>
<p className="text-tremor-default font-medium text-tremor-content dark:text-dark-tremor-content">
Budget
</p>
<div className="mt-2 flex items-baseline space-x-2.5">
<p className="text-tremor font-semibold text-tremor-content-strong dark:text-dark-tremor-content-strong">
{selectedItem.max_budget != null ? (
<>{selectedItem.max_budget}</>
) : (
<>Unlimited</>
)}
</p>
</div>
</Card>
<Card key={item.name}>
<p className="text-tremor-default font-medium text-tremor-content dark:text-dark-tremor-content">
Expires
</p>
<div className="mt-2 flex items-baseline space-x-2.5">
<p className="text-tremor-default font-small text-tremor-content-strong dark:text-dark-tremor-content-strong">
{selectedItem.expires != null ? (
<>
{new Date(selectedItem.expires).toLocaleString(undefined, {
day: 'numeric',
month: 'long',
year: 'numeric',
hour: 'numeric',
minute: 'numeric',
second: 'numeric'
})}
</>
) : (
<>Never</>
)}
</p>
</div>
</Card>
</div>
<Card className="mt-6 mb-6">
{spendData && (
<BarChart
className="mt-6"
data={spendData}
colors={["blue", "amber"]}
index="date"
categories={["spend", "predicted_spend"]}
yAxisWidth={80}
/>
)}
</Card>
<Title>Metadata</Title>
<Text>{JSON.stringify(selectedItem.metadata)}</Text>
<Button
variant="light"
className="mx-auto flex items-center"
onClick={() => {
setOpenDialogId(null);
setSelectedItem(null);
}}
>
Close
</Button>
</>
)}
</DialogPanel>
</Dialog>
<Icon
icon={PencilAltIcon}
size="sm"
onClick={() => handleEditClick(item)}
/>
<Icon
onClick={() => handleDelete(item.token)}
icon={TrashIcon}
size="sm"
/>
</TableCell>
</TableRow>
);
})}
</TableBody>
</Table>
{isDeleteModalOpen && (
<div className="fixed z-10 inset-0 overflow-y-auto">
<div className="flex items-end justify-center min-h-screen pt-4 px-4 pb-20 text-center sm:block sm:p-0">
<div
className="fixed inset-0 transition-opacity"
aria-hidden="true"
>
<div className="absolute inset-0 bg-gray-500 opacity-75"></div>
</div>
{/* Modal Panel */}
<span
className="hidden sm:inline-block sm:align-middle sm:h-screen"
aria-hidden="true"
>
&#8203;
</span>
{/* Confirmation Modal Content */}
<div className="inline-block align-bottom bg-white rounded-lg text-left overflow-hidden shadow-xl transform transition-all sm:my-8 sm:align-middle sm:max-w-lg sm:w-full">
<div className="bg-white px-4 pt-5 pb-4 sm:p-6 sm:pb-4">
<div className="sm:flex sm:items-start">
<div className="mt-3 text-center sm:mt-0 sm:ml-4 sm:text-left">
<h3 className="text-lg leading-6 font-medium text-gray-900">
Delete Key
</h3>
<div className="mt-2">
<p className="text-sm text-gray-500">
Are you sure you want to delete this key ?
</p>
</div>
</div>
</div>
</div>
<div className="bg-gray-50 px-4 py-3 sm:px-6 sm:flex sm:flex-row-reverse">
<Button onClick={confirmDelete} color="red" className="ml-2">
Delete
</Button>
<Button onClick={cancelDelete}>Cancel</Button>
</div>
</div>
</div>
</div>
)}
</Card>
{selectedToken && (
<EditKeyModal
visible={editModalVisible}
onCancel={handleEditCancel}
token={selectedToken}
onSubmit={handleEditSubmit}
/>
)}
</div>
);
};
export default ViewKeyTable;