diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx index d3b073ef9..a0061c5b8 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/view_key_table.tsx @@ -2,7 +2,7 @@ import React, { useEffect, useState } from "react"; import { keyDeleteCall, modelAvailableCall } from "./networking"; import { InformationCircleIcon, StatusOnlineIcon, TrashIcon, PencilAltIcon } from "@heroicons/react/outline"; -import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall } from "./networking"; +import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall, modelInfoCall } from "./networking"; import { Badge, Card, @@ -47,6 +47,15 @@ interface EditKeyModalProps { onSubmit: (data: FormData) => void; // Assuming FormData is the type of data to be submitted } + +interface ModelLimitModalProps { + visible: boolean; + onCancel: () => void; + token: ItemData; + onSubmit: (updatedMetadata: any) => void; + accessToken: string; +} + // Define the props type interface ViewKeyTableProps { userID: string; @@ -99,6 +108,7 @@ const ViewKeyTable: React.FC = ({ const [selectedToken, setSelectedToken] = useState(null); const [userModels, setUserModels] = useState([]); const initialKnownTeamIDs: Set = new Set(); + const [modelLimitModalVisible, setModelLimitModalVisible] = useState(false); const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs); @@ -125,6 +135,45 @@ const ViewKeyTable: React.FC = ({ fetchUserModels(); }, [accessToken, userID, userRole]); + const handleModelLimitClick = (token: ItemData) => { + setSelectedToken(token); + setModelLimitModalVisible(true); + }; + + const handleModelLimitSubmit = async (updatedMetadata: any) => { + if (accessToken == null || selectedToken == null) { + return; + } + + const formValues = { + ...selectedToken, + metadata: updatedMetadata, + key: selectedToken.token, + }; + + try { + let newKeyValues = await keyUpdateCall(accessToken, formValues); + console.log("Model limits updated:", newKeyValues); + + // Update the keys with the updated key + if (data) { + const updatedData = data.map((key) => + key.token === selectedToken.token ? newKeyValues : key + ); + setData(updatedData); + } + message.success("Model-specific limits updated successfully"); + } catch (error) { + console.error("Error updating model-specific limits:", error); + message.error("Failed to update model-specific limits"); + } + + setModelLimitModalVisible(false); + setSelectedToken(null); + }; + + + useEffect(() => { if (teams) { const teamIDSet: Set = new Set(); @@ -307,6 +356,165 @@ const ViewKeyTable: React.FC = ({ ); }; + + const ModelLimitModal: React.FC = ({ visible, onCancel, token, onSubmit, accessToken }) => { + const [modelLimits, setModelLimits] = useState<{ [key: string]: { tpm: number, rpm: number } }>({}); + const [availableModels, setAvailableModels] = useState([]); + const [newModelRow, setNewModelRow] = useState(null); + + useEffect(() => { + if (token.metadata) { + const tpmLimits = token.metadata.model_tpm_limit || {}; + const rpmLimits = token.metadata.model_rpm_limit || {}; + const combinedLimits: { [key: string]: { tpm: number, rpm: number } } = {}; + + Object.keys({ ...tpmLimits, ...rpmLimits }).forEach(model => { + combinedLimits[model] = { + tpm: tpmLimits[model] || 0, + rpm: rpmLimits[model] || 0 + }; + }); + + setModelLimits(combinedLimits); + } + + const fetchAvailableModels = async () => { + try { + const modelDataResponse = await modelInfoCall(accessToken, "", ""); + const allModelGroups = Array.from(new Set(modelDataResponse.data.map((model: any) => model.model_name))); + setAvailableModels(allModelGroups); + } catch (error) { + console.error("Error fetching model data:", error); + message.error("Failed to fetch available models"); + } + }; + + fetchAvailableModels(); + }, [token, accessToken]); + + const handleLimitChange = (model: string, type: 'tpm' | 'rpm', value: number | null) => { + setModelLimits(prev => ({ + ...prev, + [model]: { + ...prev[model], + [type]: value || 0 + } + })); + }; + + const handleAddLimit = () => { + setNewModelRow(''); + }; + + const handleModelSelect = (model: string) => { + if (!modelLimits[model]) { + setModelLimits(prev => ({ + ...prev, + [model]: { tpm: 0, rpm: 0 } + })); + } + setNewModelRow(null); + }; + + const handleRemoveModel = (model: string) => { + setModelLimits(prev => { + const { [model]: _, ...rest } = prev; + return rest; + }); + }; + + const handleSubmit = () => { + const updatedMetadata = { + ...token.metadata, + model_tpm_limit: Object.fromEntries(Object.entries(modelLimits).map(([model, limits]) => [model, limits.tpm])), + model_rpm_limit: Object.fromEntries(Object.entries(modelLimits).map(([model, limits]) => [model, limits.rpm])), + }; + onSubmit(updatedMetadata); + }; + + return ( + +
+ + + + Model + TPM Limit + RPM Limit + Actions + + + + {Object.entries(modelLimits).map(([model, limits]) => ( + + {model} + + handleLimitChange(model, 'tpm', value)} + /> + + + handleLimitChange(model, 'rpm', value)} + /> + + + + + + ))} + {newModelRow !== null && ( + + + + + - + - + + + + + )} + +
+ +
+
+ + +
+
+ ); + }; @@ -419,7 +627,8 @@ const ViewKeyTable: React.FC = ({ Spend (USD) Budget (USD) Models - TPM / RPM Limits + Rate Limits + Rate Limits per model @@ -546,6 +755,9 @@ const ViewKeyTable: React.FC = ({ {item.rpm_limit ? item.rpm_limit : "Unlimited"} + + + { @@ -720,6 +932,16 @@ const ViewKeyTable: React.FC = ({ onSubmit={handleEditSubmit} /> )} + +{selectedToken && ( + setModelLimitModalVisible(false)} + token={selectedToken} + onSubmit={handleModelLimitSubmit} + accessToken={accessToken} + /> + )} ); };