From b0ae0101f49d69eb8e98ebed4e5c7f2ba30dadfc Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 26 Aug 2024 18:53:28 -0700
Subject: [PATCH] make regenerating api keys enterprise

---
 ui/litellm-dashboard/src/app/page.tsx         |  2 ++
 .../src/components/user_dashboard.tsx         |  3 ++
 .../src/components/view_key_table.tsx         | 36 +++++++++++++++----
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx
index 02ef8ebe0..b35d99cad 100644
--- a/ui/litellm-dashboard/src/app/page.tsx
+++ b/ui/litellm-dashboard/src/app/page.tsx
@@ -141,6 +141,7 @@ const CreateKeyPage = () => {
           <UserDashboard
               userID={userID}
               userRole={userRole}
+              premiumUser={premiumUser}
               teams={teams}
               keys={keys}
               setUserRole={setUserRole}
@@ -175,6 +176,7 @@ const CreateKeyPage = () => {
             <UserDashboard
               userID={userID}
               userRole={userRole}
+              premiumUser={premiumUser}
               teams={teams}
               keys={keys}
               setUserRole={setUserRole}
diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx
index 3d1d4ea60..660e1a6d9 100644
--- a/ui/litellm-dashboard/src/components/user_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx
@@ -48,6 +48,7 @@ interface UserDashboardProps {
   setKeys: React.Dispatch<React.SetStateAction<Object[] | null>>;
   setProxySettings: React.Dispatch<React.SetStateAction<any>>;
   proxySettings: any;
+  premiumUser: boolean;
 }
 
 type TeamInterface = {
@@ -68,6 +69,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
   setKeys,
   setProxySettings,
   proxySettings,
+  premiumUser,
 }) => {
   const [userSpendData, setUserSpendData] = useState<UserSpendData | null>(
     null
@@ -328,6 +330,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
             selectedTeam={selectedTeam ? selectedTeam : null}
             data={keys}
             setData={setKeys}
+            premiumUser={premiumUser}
             teams={teams}
           />
           <CreateKey
diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx
index d9e73b109..e0ac43521 100644
--- a/ui/litellm-dashboard/src/components/view_key_table.tsx
+++ b/ui/litellm-dashboard/src/components/view_key_table.tsx
@@ -69,6 +69,7 @@ interface ViewKeyTableProps {
   data: any[] | null;
   setData: React.Dispatch<React.SetStateAction<any[] | null>>;
   teams: any[] | null;
+  premiumUser: boolean;
 }
 
 interface ItemData {
@@ -96,7 +97,8 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   selectedTeam,
   data,
   setData,
-  teams
+  teams,
+  premiumUser
 }) => {
   const [isButtonClicked, setIsButtonClicked] = useState(false);
   const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false);
@@ -619,6 +621,11 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   };
 
   const handleRegenerateKey = async () => {
+    if (!premiumUser) {
+      message.error("Regenerate API Key is an Enterprise feature. Please upgrade to use this feature.");
+      return;
+    }
+
     try {
       if (selectedToken == null) {
         message.error("Please select a key to regenerate");
@@ -994,14 +1001,31 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
         <Button key="cancel" onClick={() => setRegenerateDialogVisible(false)} className="mr-2">
           Cancel
         </Button>,
-        <Button key="regenerate" onClick={handleRegenerateKey}>
-          Regenerate
+        <Button
+          key="regenerate"
+          onClick={handleRegenerateKey}
+          disabled={!premiumUser}
+        >
+          {premiumUser ? "Regenerate" : "Upgrade to Regenerate"}
         </Button>
       ]}
     >
-      <p>Are you sure you want to regenerate this key?</p>
-      <p>Key Alias:</p>
-      <pre>{selectedToken?.key_alias || 'No alias set'}</pre>
+      {premiumUser ? (
+        <>
+          <p>Are you sure you want to regenerate this key?</p>
+          <p>Key Alias:</p>
+          <pre>{selectedToken?.key_alias || 'No alias set'}</pre>
+        </>
+      ) : (
+        <div>
+          <p className="mb-2 text-gray-500 italic text-[12px]">Upgrade to use this feature</p>
+          <Button variant="primary" className="mb-2">
+            <a href="https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat" target="_blank">
+              Get Free Trial
+            </a>
+          </Button>
+        </div>
+      )}
     </Modal>
 
     {/* Regenerated Key Display Modal */}