From 325538f8a4975c2251094d6a94ba41b36e3a62ee Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 28 Aug 2024 12:27:32 -0700
Subject: [PATCH] fix(key_management_endpoints.py): expose 'key' param, for
 setting your own key value

---
 litellm/proxy/_types.py                                  | 1 +
 .../management_endpoints/key_management_endpoints.py     | 9 ++++++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index f2bef739c1d..dd038d80bd4 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -587,6 +587,7 @@ class GenerateRequestBase(LiteLLMBase):
 
 class GenerateKeyRequest(GenerateRequestBase):
     key_alias: Optional[str] = None
+    key: Optional[str] = None
     duration: Optional[str] = None
     aliases: Optional[dict] = {}
     config: Optional[dict] = {}
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 9bb07cfee34..00e17400c31 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -55,6 +55,7 @@ async def generate_key_fn(
     Parameters:
     - duration: Optional[str] - Specify the length of time the token is valid for. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d").
     - key_alias: Optional[str] - User defined key alias
+    - key: Optional[str] - User defined key value. If not set, a 16-digit unique sk-key is created for you.
     - team_id: Optional[str] - The team id of the key
     - user_id: Optional[str] - The user id of the key
     - models: Optional[list] - Model_name's a user is allowed to call. (if empty, key is allowed to call all models)
@@ -728,6 +729,9 @@ async def generate_key_helper_fn(
     max_budget: Optional[float] = None,  # max_budget is used to Budget Per user
     budget_duration: Optional[str] = None,  # max_budget is used to Budget Per user
     token: Optional[str] = None,
+    key: Optional[
+        str
+    ] = None,  # dev-friendly alt param for 'token'. Exposed on `/key/generate` for setting key value yourself.
     user_id: Optional[str] = None,
     team_id: Optional[str] = None,
     user_email: Optional[str] = None,
@@ -763,7 +767,10 @@ async def generate_key_helper_fn(
         )
 
     if token is None:
-        token = f"sk-{secrets.token_urlsafe(16)}"
+        if key is not None:
+            token = key
+        else:
+            token = f"sk-{secrets.token_urlsafe(16)}"
 
     if duration is None:  # allow tokens that never expire
         expires = None