From 4dc2eea58d0c11627792fef20c6fcb50a357551a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 26 Aug 2024 18:40:51 -0700 Subject: [PATCH] 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 aef2fb393..4688445fe 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 e41a67c08..49a4c95c7 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