From fb150f7ce5032946af4aef739180cc3c5d172c98 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 16:52:19 -0700 Subject: [PATCH 01/14] update schema --- litellm/proxy/schema.prisma | 3 +++ schema.prisma | 3 +++ 2 files changed, 6 insertions(+) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index cf61635a0b..55bfdb2e07 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -149,6 +149,9 @@ model LiteLLM_VerificationToken { model_max_budget Json @default("{}") budget_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + key_state String? // can be "active", "inactive" + created_at DateTime @default(now()) @map("created_at") + updated_at DateTime @default(now()) @updatedAt @map("updated_at") } model LiteLLM_EndUserTable { diff --git a/schema.prisma b/schema.prisma index 8f41251041..01fe468d04 100644 --- a/schema.prisma +++ b/schema.prisma @@ -149,6 +149,9 @@ model LiteLLM_VerificationToken { model_max_budget Json @default("{}") budget_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) + key_state String? // can be "active", "inactive" + created_at DateTime @default(now()) @map("created_at") + updated_at DateTime @default(now()) @updatedAt @map("updated_at") } model LiteLLM_EndUserTable { From cbef0c0a0dd8abcd0f86203df3f494cfae5dbfc8 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 16:52:33 -0700 Subject: [PATCH 02/14] add key_state created at to token --- litellm/proxy/_types.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index d660e576d5..260f0431d4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1299,8 +1299,10 @@ class LiteLLM_VerificationToken(LiteLLMBase): model_max_budget: Dict = {} soft_budget_cooldown: bool = False litellm_budget_table: Optional[dict] = None - org_id: Optional[str] = None # org id for a given key + key_state: Optional[str] = None # can be "active", "inactive" + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None model_config = ConfigDict(protected_namespaces=()) From 5745f3d6cc8e115a38297222fd498c6b8689432d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 17:27:06 -0700 Subject: [PATCH 03/14] fix schema --- litellm/proxy/_types.py | 1 - litellm/proxy/schema.prisma | 1 - schema.prisma | 1 - 3 files changed, 3 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 260f0431d4..554c9e4e43 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1300,7 +1300,6 @@ class LiteLLM_VerificationToken(LiteLLMBase): soft_budget_cooldown: bool = False litellm_budget_table: Optional[dict] = None org_id: Optional[str] = None # org id for a given key - key_state: Optional[str] = None # can be "active", "inactive" created_at: Optional[datetime] = None updated_at: Optional[datetime] = None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 55bfdb2e07..1af0c0a346 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -149,7 +149,6 @@ model LiteLLM_VerificationToken { model_max_budget Json @default("{}") budget_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) - key_state String? // can be "active", "inactive" created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") } diff --git a/schema.prisma b/schema.prisma index 01fe468d04..86ec201f7f 100644 --- a/schema.prisma +++ b/schema.prisma @@ -149,7 +149,6 @@ model LiteLLM_VerificationToken { model_max_budget Json @default("{}") budget_id String? litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) - key_state String? // can be "active", "inactive" created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") } From 7230ee1f55b934d54bb0fa5b7fc3f275e451059c Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 17:59:44 -0700 Subject: [PATCH 04/14] add regenerate_key_fn --- .../key_management_endpoints.py | 79 +++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2e16b533c8..ebe1d9db3b 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -966,3 +966,82 @@ async def delete_verification_token(tokens: List, user_id: Optional[str] = None) verbose_proxy_logger.debug(traceback.format_exc()) raise e return deleted_tokens + + +@router.post( + "/key/{key:path}/regenerate", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], +) +@management_endpoint_wrapper +async def regenerate_key_fn( + key: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + litellm_changed_by: Optional[str] = Header( + None, + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), +) -> GenerateKeyResponse: + from litellm.proxy.proxy_server import ( + hash_token, + premium_user, + prisma_client, + user_api_key_cache, + ) + + """ + Endpoint for regenerating a key + """ + + # Check if key exists, raise exception if key is not in the DB + + ### 1. Create New copy that is duplicate of existing key + ###################################################################### + + # create duplicate of existing key + # set token = new token generated + # insert new token in DB + + # create hash of token + if prisma_client is None: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={"error": "DB not connected. prisma_client is None"}, + ) + hashed_api_key = hash_token(key) + + _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( + where={"token": hashed_api_key}, + ) + if _key_in_db is None: + raise HTTPException( + status_code=status.HTTP_404_NOT_FOUND, + detail={"error": f"Key {key} not found."}, + ) + + verbose_proxy_logger.debug("key_in_db: %s", _key_in_db) + + new_token = f"sk-{secrets.token_urlsafe(16)}" + new_token_hash = hash_token(new_token) + + # update new token in DB + updated_token = await prisma_client.db.litellm_verificationtoken.update( + where={"token": hashed_api_key}, data={"token": new_token_hash} + ) + updated_token_dict = {} + if updated_token is not None: + updated_token_dict = dict(updated_token) + + updated_token_dict["token"] = new_token + + ### 3. remove existing key entry from cache + ###################################################################### + if key: + user_api_key_cache.delete_cache(key) + + if hashed_api_key: + user_api_key_cache.delete_cache(hashed_api_key) + + return GenerateKeyResponse( + **updated_token_dict, + ) From 6b8bb69af399caefc9bdb980d163d2dc75d8f800 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:00:51 -0700 Subject: [PATCH 05/14] test test_regenerate_api_key --- litellm/tests/test_key_generate_prisma.py | 101 ++++++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index 8eedd639fc..e41a67c08b 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -56,6 +56,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_fn, generate_key_helper_fn, info_key_fn, + regenerate_key_fn, update_key_fn, ) from litellm.proxy.management_endpoints.team_endpoints import ( @@ -2935,3 +2936,103 @@ async def test_team_access_groups(prisma_client): "not allowed to call model" in e.message and "Allowed team models" in e.message ) + + +################ Unit Tests for testing regeneration of keys ########### +@pytest.mark.asyncio() +async def test_regenerate_api_key(prisma_client): + litellm.set_verbose = True + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") + await litellm.proxy.proxy_server.prisma_client.connect() + import uuid + + # generate new key + key_alias = f"test_alias_regenerate_key-{uuid.uuid4()}" + spend = 100 + max_budget = 400 + models = ["fake-openai-endpoint"] + new_key = await generate_key_fn( + data=GenerateKeyRequest( + key_alias=key_alias, spend=spend, max_budget=max_budget, models=models + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + + generated_key = new_key.key + print(generated_key) + + # assert the new key works as expected + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + async def return_body(): + return_string = f'{{"model": "fake-openai-endpoint"}}' + # return string as bytes + return return_string.encode() + + request.body = return_body + result = await user_api_key_auth(request=request, api_key=f"Bearer {generated_key}") + print(result) + + # regenerate the key + new_key = await regenerate_key_fn( + key=generated_key, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + ) + print("response from regenerate_key_fn", new_key) + + # assert the new key works as expected + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + async def return_body_2(): + return_string = f'{{"model": "fake-openai-endpoint"}}' + # return string as bytes + return return_string.encode() + + request.body = return_body_2 + result = await user_api_key_auth(request=request, api_key=f"Bearer {new_key.key}") + print(result) + + # assert the old key stops working + request = Request(scope={"type": "http"}) + request._url = URL(url="/chat/completions") + + async def return_body_3(): + return_string = f'{{"model": "fake-openai-endpoint"}}' + # return string as bytes + return return_string.encode() + + request.body = return_body_3 + try: + result = await user_api_key_auth( + request=request, api_key=f"Bearer {generated_key}" + ) + print(result) + pytest.fail(f"This should have failed!. the key has been regenerated") + except Exception as e: + print("got expected exception", e) + assert "Invalid proxy server token passed" in e.message + + # Check that the regenerated key has the same spend, max_budget, models and key_alias + assert new_key.spend == spend, f"Expected spend {spend} but got {new_key.spend}" + assert ( + new_key.max_budget == max_budget + ), f"Expected max_budget {max_budget} but got {new_key.max_budget}" + assert ( + new_key.key_alias == key_alias + ), f"Expected key_alias {key_alias} but got {new_key.key_alias}" + assert ( + new_key.models == models + ), f"Expected models {models} but got {new_key.models}" + + pass From 2615edc468042333fbe9965e88fb06113fd77d03 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:15:52 -0700 Subject: [PATCH 06/14] allow using hashed api keys on regen key --- .../proxy/management_endpoints/key_management_endpoints.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ebe1d9db3b..aef2fb3930 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1008,7 +1008,11 @@ async def regenerate_key_fn( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail={"error": "DB not connected. prisma_client is None"}, ) - hashed_api_key = hash_token(key) + + if "sk" not in key: + hashed_api_key = key + else: + hashed_api_key = hash_token(key) _key_in_db = await prisma_client.db.litellm_verificationtoken.find_unique( where={"token": hashed_api_key}, From 40c018272c6735202797662fc563bc83d5dd0acf Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:28:26 -0700 Subject: [PATCH 07/14] working regenerate key flow --- .../src/components/networking.tsx | 31 +++++++ .../src/components/view_key_table.tsx | 83 ++++++++++++++++++- 2 files changed, 112 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index f550764789..a6bd5d32cb 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -770,6 +770,37 @@ export const claimOnboardingToken = async ( throw error; } }; + +export const regenerateKeyCall = async (accessToken: string, keyToRegenerate: string) => { + try { + const url = proxyBaseUrl + ? `${proxyBaseUrl}/key/${keyToRegenerate}/regenerate` + : `/key/${keyToRegenerate}/regenerate`; + + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({}), + }); + + if (!response.ok) { + const errorData = await response.text(); + handleError(errorData); + throw new Error("Network response was not ok"); + } + + const data = await response.json(); + console.log("Regenerate key Response:", data); + return data; + } catch (error) { + console.error("Failed to regenerate key:", error); + throw error; + } +}; + let ModelListerrorShown = false; let errorTimer: NodeJS.Timeout | null = null; diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx index 44a1f6ece8..752aad37f9 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/view_key_table.tsx @@ -1,12 +1,14 @@ "use client"; import React, { useEffect, useState } from "react"; import { keyDeleteCall, modelAvailableCall } from "./networking"; -import { InformationCircleIcon, StatusOnlineIcon, TrashIcon, PencilAltIcon } from "@heroicons/react/outline"; -import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall, modelInfoCall } from "./networking"; +import { InformationCircleIcon, StatusOnlineIcon, TrashIcon, PencilAltIcon, RefreshIcon } from "@heroicons/react/outline"; +import { keySpendLogsCall, PredictedSpendLogsCall, keyUpdateCall, modelInfoCall, regenerateKeyCall } from "./networking"; import { Badge, Card, Table, + Grid, + Col, Button, TableBody, TableCell, @@ -33,6 +35,8 @@ import { Select, } from "antd"; +import { CopyToClipboard } from "react-copy-to-clipboard"; + const { Option } = Select; const isLocal = process.env.NODE_ENV === "development"; const proxyBaseUrl = isLocal ? "http://localhost:4000" : null; @@ -109,6 +113,8 @@ const ViewKeyTable: React.FC = ({ const [userModels, setUserModels] = useState([]); const initialKnownTeamIDs: Set = new Set(); const [modelLimitModalVisible, setModelLimitModalVisible] = useState(false); + const [regenerateDialogVisible, setRegenerateDialogVisible] = useState(false); + const [regeneratedKey, setRegeneratedKey] = useState(null); const [knownTeamIDs, setKnownTeamIDs] = useState(initialKnownTeamIDs); @@ -612,6 +618,18 @@ const ViewKeyTable: React.FC = ({ setKeyToDelete(null); }; + const handleRegenerateKey = async () => { + try { + const response = await regenerateKeyCall(accessToken, selectedToken.token); + setRegeneratedKey(response.key); + setRegenerateDialogVisible(false); + message.success("API Key regenerated successfully"); + } catch (error) { + console.error("Error regenerating key:", error); + message.error("Failed to regenerate API Key"); + } + }; + if (data == null) { return; } @@ -768,6 +786,7 @@ const ViewKeyTable: React.FC = ({ size="sm" /> + = ({ size="sm" onClick={() => handleEditClick(item)} /> + { + setSelectedToken(item); + setRegenerateDialogVisible(true); + }} + icon={RefreshIcon} + size="sm" + /> handleDelete(item)} icon={TrashIcon} @@ -942,6 +969,58 @@ const ViewKeyTable: React.FC = ({ accessToken={accessToken} /> )} + + {/* Regenerate Key Confirmation Dialog */} + setRegenerateDialogVisible(false)} + > +

Are you sure you want to regenerate this key?

+

Key Alias:

+ 
{selectedToken?.key_alias || 'No alias set'}
+ + + {/* Regenerated Key Display Modal */} + {regeneratedKey && ( + setRegeneratedKey(null)} + onCancel={() => setRegeneratedKey(null)} + footer={null} + > + + Save your New Key + +

+ Please save this new secret key somewhere safe and accessible. For + security reasons, you will not be able to view it again through + your LiteLLM account. If you lose this secret key, you will need to + generate a new one. +

+ + + New API Key: +
+ 
+                {regeneratedKey}
+
+
+ message.success("API Key copied to clipboard")}> + + + + + + )} ); }; From 3527e47b59fe169a261028aec8fd2e8e3b348f6d Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:33:18 -0700 Subject: [PATCH 08/14] ui regenerate an api key --- .../src/components/view_key_table.tsx | 20 ++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx index 752aad37f9..620c8ca1c1 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/view_key_table.tsx @@ -620,6 +620,10 @@ const ViewKeyTable: React.FC = ({ const handleRegenerateKey = async () => { try { + if (selectedToken == null) { + message.error("Please select a key to regenerate"); + return; + } const response = await regenerateKeyCall(accessToken, selectedToken.token); setRegeneratedKey(response.key); setRegenerateDialogVisible(false); @@ -974,8 +978,15 @@ const ViewKeyTable: React.FC = ({ setRegenerateDialogVisible(false)} + footer={[ + , + + ]} >

Are you sure you want to regenerate this key?

Key Alias:

@@ -986,9 +997,12 @@ const ViewKeyTable: React.FC = ({ {regeneratedKey && ( setRegeneratedKey(null)} onCancel={() => setRegeneratedKey(null)} - footer={null} + footer={[ + + ]} > Save your New Key From 4dc2eea58d0c11627792fef20c6fcb50a357551a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:40:51 -0700 Subject: [PATCH 09/14] update key name when regenerating a key --- .../proxy/management_endpoints/key_management_endpoints.py | 7 ++++++- litellm/tests/test_key_generate_prisma.py | 2 ++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index aef2fb3930..4688445fee 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1027,10 +1027,15 @@ async def regenerate_key_fn( new_token = f"sk-{secrets.token_urlsafe(16)}" new_token_hash = hash_token(new_token) + new_token_key_name = f"sk-...{new_token[-4:]}" # update new token in DB updated_token = await prisma_client.db.litellm_verificationtoken.update( - where={"token": hashed_api_key}, data={"token": new_token_hash} + where={"token": hashed_api_key}, + data={ + "token": new_token_hash, + "key_name": new_token_key_name, + }, ) updated_token_dict = {} if updated_token is not None: diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index e41a67c08b..49a4c95c75 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -3035,4 +3035,6 @@ async def test_regenerate_api_key(prisma_client): new_key.models == models ), f"Expected models {models} but got {new_key.models}" + assert new_key.key_name == f"sk-...{new_key.key[-4:]}" + pass From 5c2c316919743b0bb2219d059b38256dea69d0f3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:41:20 -0700 Subject: [PATCH 10/14] update ui regen key --- ui/litellm-dashboard/src/components/view_key_table.tsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx index 620c8ca1c1..41d930fc3d 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/view_key_table.tsx @@ -1005,10 +1005,10 @@ const ViewKeyTable: React.FC = ({ ]} > - Save your New Key + Regenerated Key

- Please save this new secret key somewhere safe and accessible. For + Please replace your old key with the new key generated.For security reasons, you will not be able to view it again through your LiteLLM account. If you lose this secret key, you will need to generate a new one. From 83813af0a2ec6e1dc1314cb44e7dff1b58b11e7b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:44:46 -0700 Subject: [PATCH 11/14] regenerate key --- .../src/components/view_key_table.tsx | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx index 41d930fc3d..d9e73b1094 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/view_key_table.tsx @@ -626,6 +626,17 @@ const ViewKeyTable: React.FC = ({ } const response = await regenerateKeyCall(accessToken, selectedToken.token); 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 + ); + setData(updatedData); + } + setRegenerateDialogVisible(false); message.success("API Key regenerated successfully"); } catch (error) { From b0ae0101f49d69eb8e98ebed4e5c7f2ba30dadfc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:53:28 -0700 Subject: [PATCH 12/14] 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 02ef8ebe05..b35d99cad3 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -141,6 +141,7 @@ const CreateKeyPage = () => { { >; setProxySettings: React.Dispatch>; proxySettings: any; + premiumUser: boolean; } type TeamInterface = { @@ -68,6 +69,7 @@ const UserDashboard: React.FC = ({ setKeys, setProxySettings, proxySettings, + premiumUser, }) => { const [userSpendData, setUserSpendData] = useState( null @@ -328,6 +330,7 @@ const UserDashboard: React.FC = ({ selectedTeam={selectedTeam ? selectedTeam : null} data={keys} setData={setKeys} + premiumUser={premiumUser} teams={teams} /> >; teams: any[] | null; + premiumUser: boolean; } interface ItemData { @@ -96,7 +97,8 @@ const ViewKeyTable: React.FC = ({ selectedTeam, data, setData, - teams + teams, + premiumUser }) => { const [isButtonClicked, setIsButtonClicked] = useState(false); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); @@ -619,6 +621,11 @@ const ViewKeyTable: React.FC = ({ }; 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 = ({ , - ]} > -

Are you sure you want to regenerate this key?

-

Key Alias:

- 
{selectedToken?.key_alias || 'No alias set'}
+ {premiumUser ? ( + <> +

Are you sure you want to regenerate this key?

+

Key Alias:

+ 
{selectedToken?.key_alias || 'No alias set'}
+ + ) : ( +
+

Upgrade to use this feature

+ +
+ )} {/* Regenerated Key Display Modal */} From 75cbbea0719bf646a9c042d89af88f224b08fdbe Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:54:50 -0700 Subject: [PATCH 13/14] enforce regenerating keys in enterprise tier --- .../proxy/management_endpoints/key_management_endpoints.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4688445fee..9bb07cfee3 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -993,6 +993,11 @@ async def regenerate_key_fn( Endpoint for regenerating a key """ + if premium_user is not True: + raise ValueError( + f"Regenerating Virtual Keys is an Enterprise feature, {CommonProxyErrors.not_premium_user.value}" + ) + # Check if key exists, raise exception if key is not in the DB ### 1. Create New copy that is duplicate of existing key From a043676dc48cf0ec75908b99df240c641c7817a9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:59:29 -0700 Subject: [PATCH 14/14] fix regen api key flow --- .../src/components/view_key_table.tsx | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx index e0ac435217..019a11179a 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/view_key_table.tsx @@ -1043,13 +1043,26 @@ const ViewKeyTable: React.FC = ({ Regenerated Key

- Please replace your old key with the new key generated.For + Please replace your old key with the new key generated. For security reasons, you will not be able to view it again through your LiteLLM account. If you lose this secret key, you will need to generate a new one.

+ Key Alias: +
+ 
+                {selectedToken?.key_alias || 'No alias set'}
+
+
New API Key: