Merge pull request #2809 from BerriAI/ui_use_token_id_in_key_gen

[UI] QA Fix Edit Key flow - return `token_id` in /key/generate respose

litellm/proxy/_types.py

@@ -306,6 +306,7 @@ class GenerateKeyResponse(GenerateKeyRequest):
     key_name: Optional[str] = None
     expires: Optional[datetime]
     user_id: Optional[str] = None
+    token_id: Optional[str] = None
 
     @root_validator(pre=True)
     def set_model_info(cls, values):

litellm/proxy/proxy_server.py

@@ -2563,7 +2563,10 @@ async def generate_key_helper_fn(
 
             ## CREATE KEY
             verbose_proxy_logger.debug("prisma_client: Creating Key= %s", key_data)
-            await prisma_client.insert_data(data=key_data, table_name="key")
+            create_key_response = await prisma_client.insert_data(
+                data=key_data, table_name="key"
+            )
+            key_data["token_id"] = getattr(create_key_response, "token", None)
         elif custom_db_client is not None:
             if table_name is None or table_name == "user":
                 ## CREATE USER (If necessary)

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

@@ -213,6 +213,11 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   
   const handleEditClick = (token: any) => {
     console.log("handleEditClick:", token);
+
+    // set token.token to token.token_id if token_id is not null
+    if (token.token_id !== null) {
+      token.token = token.token_id;
+    }
     setSelectedToken(token);
     setEditModalVisible(true);
   };