Merge pull request #5566 from BerriAI/litellm_ui_regen_keys

[Feat] Allow setting duration time when regenerating key

ui/litellm-dashboard/src/components/networking.tsx

@@ -771,7 +771,7 @@ export const claimOnboardingToken = async (
   }
 };
 
-export const regenerateKeyCall = async (accessToken: string, keyToRegenerate: string) => {
+export const regenerateKeyCall = async (accessToken: string, keyToRegenerate: string, formData: any) => {
   try {
     const url = proxyBaseUrl
       ? `${proxyBaseUrl}/key/${keyToRegenerate}/regenerate`
@@ -783,7 +783,7 @@ export const regenerateKeyCall = async (accessToken: string, keyToRegenerate: st
         [globalLitellmHeaderName]: `Bearer ${accessToken}`,
         "Content-Type": "application/json",
       },
-      body: JSON.stringify({}),
+      body: JSON.stringify(formData),
     });
 
     if (!response.ok) {

ui/litellm-dashboard/src/components/view_key_table.tsx

@@ -1,6 +1,7 @@
 "use client";
 import React, { useEffect, useState } from "react";
 import { keyDeleteCall, modelAvailableCall } from "./networking";
+import { add } from 'date-fns';
 import { InformationCircleIcon, StatusOnlineIcon, TrashIcon, PencilAltIcon, RefreshIcon } from "@heroicons/react/outline";
 import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall, modelInfoCall, regenerateKeyCall } from "./networking";
 import {
@@ -22,6 +23,7 @@ import {
   Subtitle,
   Icon,
   BarChart,
+  TextInput,
 } from "@tremor/react";
 import { Select as Select3, SelectItem, MultiSelect, MultiSelectItem } from "@tremor/react";
 import {
@@ -33,7 +35,8 @@ import {
   InputNumber,
   message,
   Select,
-  Tooltip
+  Tooltip,
+  DatePicker,
 } from "antd";
 
 import { CopyToClipboard } from "react-copy-to-clipboard";
@@ -120,9 +123,63 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   const [modelLimitModalVisible, setModelLimitModalVisible] = useState(false);
   const [regenerateDialogVisible, setRegenerateDialogVisible] = useState(false);
   const [regeneratedKey, setRegeneratedKey] = useState<string | null>(null);
+  const [regenerateFormData, setRegenerateFormData] = useState<any>(null);
+  const [regenerateForm] = Form.useForm();
+  const [newExpiryTime, setNewExpiryTime] = useState<string | null>(null);
 
   const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs);
 
+
+  useEffect(() => {
+    const calculateNewExpiryTime = (duration: string | undefined) => {
+      if (!duration) {
+        return null;
+      }
+  
+      try {
+        const now = new Date();
+        let newExpiry: Date;
+  
+        if (duration.endsWith('s')) {
+          newExpiry = add(now, { seconds: parseInt(duration) });
+        } else if (duration.endsWith('h')) {
+          newExpiry = add(now, { hours: parseInt(duration) });
+        } else if (duration.endsWith('d')) {
+          newExpiry = add(now, { days: parseInt(duration) });
+        } else {
+          throw new Error('Invalid duration format');
+        }
+  
+        return newExpiry.toLocaleString('en-US', {
+          year: 'numeric',
+          month: 'numeric',
+          day: 'numeric',
+          hour: 'numeric',
+          minute: 'numeric',
+          second: 'numeric',
+          hour12: true
+        });
+      } catch (error) {
+        return null;
+      }
+    };
+
+    console.log("in calculateNewExpiryTime for selectedToken", selectedToken);
+        
+  
+    // When a new duration is entered
+    if (regenerateFormData?.duration) {
+      setNewExpiryTime(calculateNewExpiryTime(regenerateFormData.duration));
+    } else {
+      setNewExpiryTime(null);
+    }
+
+    console.log("calculateNewExpiryTime:", newExpiryTime);
+  }, [selectedToken, regenerateFormData?.duration]);
+  
+
+  
+
   useEffect(() => {
     const fetchUserModels = async () => {
       try {
@@ -146,6 +203,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
     fetchUserModels();
   }, [accessToken, userID, userRole]);
 
+
   const handleModelLimitClick = (token: ItemData) => {
     setSelectedToken(token);
     setModelLimitModalVisible(true);
@@ -678,31 +736,53 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
     setKeyToDelete(null);
   };
 
+  const handleRegenerateClick = (token: any) => {
+    setSelectedToken(token);
+    setNewExpiryTime(null);
+    regenerateForm.setFieldsValue({
+      key_alias: token.key_alias,
+      max_budget: token.max_budget,
+      tpm_limit: token.tpm_limit,
+      rpm_limit: token.rpm_limit,
+      duration: token.duration || '',
+    });
+    setRegenerateDialogVisible(true);
+  };
+
+  const handleRegenerateFormChange = (field: string, value: any) => {
+    setRegenerateFormData((prev: any) => ({
+      ...prev,
+      [field]: value,
+    }));
+  };
+
   const handleRegenerateKey = async () => {
     if (!premiumUser) {
       message.error("Regenerate API Key is an Enterprise feature. Please upgrade to use this feature.");
       return;
     }
 
+    if (selectedToken == null) {
+      return;
+    }
+
     try {
-      if (selectedToken == null) {
-        message.error("Please select a key to regenerate");
-        return;
-      }
-      const response = await regenerateKeyCall(accessToken, selectedToken.token);
+      const formValues = await regenerateForm.validateFields();
+      const response = await regenerateKeyCall(accessToken, selectedToken.token, formValues);
       setRegeneratedKey(response.key);
 
       // Update the data state with the new key_name
       if (data) {
         const updatedData = data.map(item => 
-          item.token === selectedToken.token 
-            ? { ...item, key_name: response.key_name } 
+          item.token === selectedToken?.token 
+            ? { ...item, key_name: response.key_name, ...formValues } 
             : item
         );
         setData(updatedData);
       }
 
       setRegenerateDialogVisible(false);
+      regenerateForm.resetFields();
       message.success("API Key regenerated successfully");
     } catch (error) {
       console.error("Error regenerating key:", error);
@@ -997,10 +1077,7 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
                     onClick={() => handleEditClick(item)}
                   />
                   <Icon
-                      onClick={() => {
-                        setSelectedToken(item);
-                        setRegenerateDialogVisible(true);
-                      }}
+                      onClick={() => handleRegenerateClick(item)}
                       icon={RefreshIcon}
                       size="sm"
                     />
@@ -1080,13 +1157,19 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
         />
       )}
 
-    {/* Regenerate Key Confirmation Dialog */}
+    {/* Regenerate Key Form Modal */}
     <Modal
       title="Regenerate API Key"
       visible={regenerateDialogVisible}
-      onCancel={() => setRegenerateDialogVisible(false)}
+      onCancel={() => {
+        setRegenerateDialogVisible(false);
+        regenerateForm.resetFields();
+      }}
       footer={[
-        <Button key="cancel" onClick={() => setRegenerateDialogVisible(false)} className="mr-2">
+        <Button key="cancel" onClick={() => {
+          setRegenerateDialogVisible(false);
+          regenerateForm.resetFields();
+        }} className="mr-2">
           Cancel
         </Button>,
         <Button
@@ -1099,11 +1182,49 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
       ]}
     >
       {premiumUser ? (
-        <>
-          <p>Are you sure you want to regenerate this key?</p>
-          <p>Key Alias:</p>
-          <pre>{selectedToken?.key_alias || 'No alias set'}</pre>
-        </>
+        <Form 
+        form={regenerateForm} 
+        layout="vertical"
+        onValuesChange={(changedValues, allValues) => {
+          if ('duration' in changedValues) {
+            handleRegenerateFormChange('duration', changedValues.duration);
+          }
+        }}
+      >
+          <Form.Item name="key_alias" label="Key Alias">
+            <TextInput disabled={true} />
+          </Form.Item>
+          <Form.Item name="max_budget" label="Max Budget (USD)">
+            <InputNumber step={0.01} precision={2} style={{ width: '100%' }} />
+          </Form.Item>
+          <Form.Item name="tpm_limit" label="TPM Limit">
+            <InputNumber style={{ width: '100%' }} />
+          </Form.Item>
+          <Form.Item name="rpm_limit" label="RPM Limit">
+            <InputNumber style={{ width: '100%' }} />
+          </Form.Item>
+          <Form.Item
+            name="duration"
+            label="Expire Key (eg: 30s, 30h, 30d)"
+            className="mt-8"
+          >
+            <TextInput placeholder="" />
+          </Form.Item>
+          <div className="mt-2 text-sm text-gray-500">
+            Current expiry: {
+              selectedToken?.expires != null ? (
+                new Date(selectedToken.expires).toLocaleString()
+              ) : (
+                'Never'
+              )
+            }
+          </div>
+          {newExpiryTime && (
+            <div className="mt-2 text-sm text-green-600">
+              New expiry: {newExpiryTime}
+            </div>
+          )}
+        </Form>
       ) : (
         <div>
           <p className="mb-2 text-gray-500 italic text-[12px]">Upgrade to use this feature</p>