build(ui/litellm-dashboard): ui cleanup

litellm/proxy/_types.py

@@ -141,11 +141,33 @@ class GenerateRequestBase(LiteLLMBase):
 
 class GenerateKeyRequest(GenerateRequestBase):
     key_alias: Optional[str] = None
-    duration: Optional[str] = "1h"
+    duration: Optional[str] = None
     aliases: Optional[dict] = {}
     config: Optional[dict] = {}
 
 
+class GenerateKeyResponse(GenerateKeyRequest):
+    key: str
+    key_name: Optional[str] = None
+    expires: Optional[datetime]
+    user_id: str
+
+    @root_validator(pre=True)
+    def set_model_info(cls, values):
+        if values.get("token") is not None:
+            values.update({"key": values.get("token")})
+        dict_fields = ["metadata", "aliases", "config"]
+        for field in dict_fields:
+            value = values.get(field)
+            if value is not None and isinstance(value, str):
+                try:
+                    values[field] = json.loads(value)
+                except json.JSONDecodeError:
+                    raise ValueError(f"Field {field} should be a valid dictionary")
+
+        return values
+
+
 class UpdateKeyRequest(GenerateKeyRequest):
     # Note: the defaults of all Params here MUST BE NONE
     # else they will get overwritten
@@ -174,12 +196,6 @@ class UserAPIKeyAuth(LiteLLMBase):  # the expected response object for user api
     rpm_limit: Optional[int] = None
 
 
-class GenerateKeyResponse(LiteLLMBase):
-    key: str
-    expires: Optional[datetime]
-    user_id: str
-
-
 class DeleteKeyRequest(LiteLLMBase):
     keys: List
 

litellm/proxy/proxy_server.py

@@ -1386,12 +1386,7 @@ async def generate_key_helper_fn(
     except Exception as e:
         traceback.print_exc()
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
-    return {
-        "token": token,
-        "expires": expires,
-        "user_id": user_id,
-        "max_budget": max_budget,
-    }
+    return key_data
 
 
 async def delete_verification_token(tokens: List):
@@ -2399,12 +2394,9 @@ async def generate_key_fn(
             data_json["key_budget_duration"] = data_json.pop("budget_duration", None)
 
         response = await generate_key_helper_fn(**data_json)
-        return GenerateKeyResponse(
-            key=response["token"],
-            expires=response["expires"],
-            user_id=response["user_id"],
-        )
+        return GenerateKeyResponse(**response)
     except Exception as e:
+        traceback.print_exc()
         if isinstance(e, HTTPException):
             raise ProxyException(
                 message=getattr(e, "detail", f"Authentication Error({str(e)})"),

ui/litellm-dashboard/package.json

@@ -14,6 +14,7 @@
     "@heroicons/react": "^1.0.6",
     "@remixicon/react": "^4.1.1",
     "@tremor/react": "^3.13.3",
+    "antd": "^5.13.2",
     "next": "14.1.0",
     "react": "^18",
     "react-dom": "^18"

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

@@ -1,39 +1,91 @@
 "use client";
 
-import React, { use } from "react";
-import { Button, TextInput } from "@tremor/react";
-
+import React, { useState, useEffect, useRef } from "react";
+import { Button, TextInput, Grid, Col } from "@tremor/react";
+import { message } from "antd";
 import { Card, Metric, Text } from "@tremor/react";
-import { createKeyCall } from "./networking";
+import { keyCreateCall } from "./networking";
 // Define the props type
 interface CreateKeyProps {
   userID: string;
   accessToken: string;
   proxyBaseUrl: string;
+  data: any[] | null;
+  setData: React.Dispatch<React.SetStateAction<any[] | null>>;
 }
 
+import { Modal, Button as Button2 } from "antd";
+
 const CreateKey: React.FC<CreateKeyProps> = ({
   userID,
   accessToken,
   proxyBaseUrl,
+  data,
+  setData,
 }) => {
-  const handleClick = () => {
-    console.log("Hello World");
+  const [isModalVisible, setIsModalVisible] = useState(false);
+  const [apiKey, setApiKey] = useState(null);
+  const handleOk = () => {
+    // Handle the OK action
+    console.log("OK Clicked");
+    setIsModalVisible(false);
+  };
+
+  const handleCancel = () => {
+    // Handle the cancel action or closing the modal
+    console.log("Modal closed");
+    setIsModalVisible(false);
+    setApiKey(null);
+  };
+
+  const handleCreate = async () => {
+    if (data == null) {
+      return;
+    }
+    try {
+      message.info("Making API Call");
+      setIsModalVisible(true);
+      const response = await keyCreateCall(proxyBaseUrl, accessToken, userID);
+      // Successfully completed the deletion. Update the state to trigger a rerender.
+      setData([...data, response]);
+      setApiKey(response["key"]);
+      message.success("API Key Created");
+    } catch (error) {
+      console.error("Error deleting the key:", error);
+      // Handle any error situations, such as displaying an error message to the user.
+    }
   };
 
   return (
-    <Button
-      className="mx-auto"
-      onClick={() =>
-        createKeyCall(
-          (proxyBaseUrl = proxyBaseUrl),
-          (accessToken = accessToken),
-          (userID = userID)
-        )
-      }
-    >
-      + Create New Key
-    </Button>
+    <div>
+      <Button className="mx-auto" onClick={handleCreate}>
+        + Create New Key
+      </Button>
+      <Modal
+        title="Save your key"
+        open={isModalVisible}
+        onOk={handleOk}
+        onCancel={handleCancel}
+      >
+        <Grid numItems={1} className="gap-2 w-full">
+          <Col numColSpan={1}>
+            <p>
+              Please save this secret key somewhere safe and accessible. For
+              security reasons, <b>you won't be able to view it again</b>{" "}
+              through your LiteLLM account. If you lose this secret key, you'll
+              need to generate a new one.
+            </p>
+          </Col>
+          <Col numColSpan={1}>
+            {apiKey != null ? (
+              <Text>API Key: {apiKey}</Text>
+            ) : (
+              <Text>Key being created, this might take 30s</Text>
+            )}
+          </Col>
+        </Grid>
+      </Modal>
+    </div>
   );
 };
 

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

@@ -2,7 +2,7 @@
  * Helper file for calls being made to proxy
  */
 
-export const createKeyCall = async (
+export const keyCreateCall = async (
   proxyBaseUrl: String,
   accessToken: String,
   userID: String
@@ -27,9 +27,42 @@ export const createKeyCall = async (
 
     const data = await response.json();
     console.log(data);
+    return data;
     // Handle success - you might want to update some state or UI based on the created key
   } catch (error) {
     console.error("Failed to create key:", error);
+    throw error;
+  }
+};
+
+export const keyDeleteCall = async (
+  proxyBaseUrl: String,
+  accessToken: String,
+  user_key: String
+) => {
+  try {
+    const response = await fetch(`${proxyBaseUrl}/key/delete`, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        keys: [user_key],
+      }),
+    });
+
+    if (!response.ok) {
+      throw new Error("Network response was not ok");
+    }
+
+    const data = await response.json();
+    console.log(data);
+    return data;
+    // Handle success - you might want to update some state or UI based on the created key
+  } catch (error) {
+    console.error("Failed to create key:", error);
+    throw error;
   }
 };
 

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

@@ -1,32 +1,41 @@
 "use client";
-import React from "react";
+import React, { useState, useEffect } from "react";
+import { userInfoCall } from "./networking";
 import { Grid, Col, Card, Text } from "@tremor/react";
 import CreateKey from "./create_key_button";
 import ViewKeyTable from "./view_key_table";
 import EnterProxyUrl from "./enter_proxy_url";
 import { useSearchParams } from "next/navigation";
 
-
-
-export default function UserDashboard() {
+const UserDashboard = () => {
+  const [data, setData] = useState<null | any[]>(null); // Keep the initialization of state here
+  // Assuming useSearchParams() hook exists and works in your setup
   const searchParams = useSearchParams();
   const userID = searchParams.get("userID");
   const accessToken = searchParams.get("accessToken");
   const proxyBaseUrl = searchParams.get("proxyBaseUrl");
 
-  const handleProxyUrlChange = (url: string) => {
-    // Do something with the entered proxy URL, e.g., save it in the state
-    console.log('Entered Proxy URL:', url);
-  };
+  // Moved useEffect inside the component and used a condition to run fetch only if the params are available
+  useEffect(() => {
+    if (userID && accessToken && proxyBaseUrl && !data) {
+      const fetchData = async () => {
+        try {
+          const response = await userInfoCall(
+            proxyBaseUrl,
+            accessToken,
+            userID
+          );
+          setData(response["keys"]); // Assuming this is the correct path to your data
+        } catch (error) {
+          console.error("There was an error fetching the data", error);
+          // Optionally, update your UI to reflect the error state here as well
+        }
+      };
+      fetchData();
+    }
+  }, [userID, accessToken, proxyBaseUrl, data]);
 
-  if (proxyBaseUrl == null) {
-    return (
-      <div>
-        <EnterProxyUrl onUrlChange={handleProxyUrlChange} />
-      </div>
-    );
-  }
-  if (userID == null || accessToken == null || proxyBaseUrl == null) {
+  if (!userID || !accessToken || !proxyBaseUrl || !data) {
     return (
       <Card
         className="max-w-xs mx-auto"
@@ -39,19 +48,25 @@ export default function UserDashboard() {
   }
 
   return (
-    <Grid numItems={1} className="gap-0 p-10 h-[75vh]">
+    <Grid numItems={1} className="gap-0 p-10 h-[75vh] w-full">
       <Col numColSpan={1}>
         <ViewKeyTable
           userID={userID}
           accessToken={accessToken}
           proxyBaseUrl={proxyBaseUrl}
+          data={data}
+          setData={setData}
         />
         <CreateKey
           userID={userID}
           accessToken={accessToken}
           proxyBaseUrl={proxyBaseUrl}
+          data={data}
+          setData={setData}
         />
       </Col>
     </Grid>
   );
-}
+};
+
+export default UserDashboard;

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

@@ -1,7 +1,7 @@
 "use client";
-import React, { useEffect, useState } from "react";
-import { userInfoCall } from "./networking";
-import { StatusOnlineIcon } from "@heroicons/react/outline";
+import React, { useEffect } from "react";
+import { keyDeleteCall } from "./networking";
+import { StatusOnlineIcon, TrashIcon } from "@heroicons/react/outline";
 import {
   Badge,
   Card,
@@ -13,6 +13,7 @@ import {
   TableRow,
   Text,
   Title,
+  Icon,
 } from "@tremor/react";
 
 // Define the props type
@@ -20,96 +21,75 @@ interface ViewKeyTableProps {
   userID: string;
   accessToken: string;
   proxyBaseUrl: string;
+  data: any[] | null;
+  setData: React.Dispatch<React.SetStateAction<any[] | null>>;
 }
 
-const data = [
-  {
-    key_alias: "my test key",
-    key_name: "sk-...hd74",
-    spend: 23.0,
-    expires: "active",
-    token: "23902dwojd90",
-  },
-  {
-    key_alias: "my test key",
-    key_name: "sk-...hd74",
-    spend: 23.0,
-    expires: "active",
-    token: "23902dwojd90",
-  },
-  {
-    key_alias: "my test key",
-    key_name: "sk-...hd74",
-    spend: 23.0,
-    expires: "active",
-    token: "23902dwojd90",
-  },
-  {
-    key_alias: "my test key",
-    key_name: "sk-...hd74",
-    spend: 23.0,
-    expires: "active",
-    token: "23902dwojd90",
-  },
-];
-
 const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
   userID,
   accessToken,
   proxyBaseUrl,
+  data,
+  setData,
 }) => {
-  const [data, setData] = useState(null); // State to store the data from the API
-
-  useEffect(() => {
-    const fetchData = async () => {
-      try {
-        const response = await userInfoCall(
-          (proxyBaseUrl = proxyBaseUrl),
-          (accessToken = accessToken),
-          (userID = userID)
-        );
-        setData(response["keys"]); // Update state with the fetched data
-      } catch (error) {
-        console.error("There was an error fetching the data", error);
-        // Optionally, update your UI to reflect the error state
-      }
-    };
-
-    fetchData(); // Call the async function to fetch data
-  }, []); // Empty dependency array
+  const handleDelete = async (token: String) => {
+    if (data == null) {
+      return;
+    }
+    try {
+      await keyDeleteCall(proxyBaseUrl, accessToken, token);
+      // Successfully completed the deletion. Update the state to trigger a rerender.
+      const filteredData = data.filter((item) => item.token !== token);
+      setData(filteredData);
+    } catch (error) {
+      console.error("Error deleting the key:", error);
+      // Handle any error situations, such as displaying an error message to the user.
+    }
+  };
 
   if (data == null) {
     return;
   }
+  console.log("RERENDER TRIGGERED");
   return (
-    <Card className="flex-auto overflow-y-auto max-h-[50vh] mb-4">
+    <Card className="w-full mx-auto flex-auto overflow-y-auto max-h-[50vh] mb-4">
       <Title>API Keys</Title>
       <Table className="mt-5">
         <TableHead>
           <TableRow>
-            <TableHeaderCell>Alias</TableHeaderCell>
             <TableHeaderCell>Secret Key</TableHeaderCell>
             <TableHeaderCell>Spend</TableHeaderCell>
-            <TableHeaderCell>Status</TableHeaderCell>
+            <TableHeaderCell>Expires</TableHeaderCell>
           </TableRow>
         </TableHead>
         <TableBody>
-          {data.map((item) => (
-            <TableRow key={item.token}>
-              <TableCell>{item.key_alias}</TableCell>
-              <TableCell>
-                <Text>{item.key_name}</Text>
-              </TableCell>
-              <TableCell>
-                <Text>{item.spend}</Text>
-              </TableCell>
-              <TableCell>
-                <Badge color="emerald" icon={StatusOnlineIcon}>
-                  {item.expires}
-                </Badge>
-              </TableCell>
-            </TableRow>
-          ))}
+          {data.map((item) => {
+            console.log(item);
+            return (
+              <TableRow key={item.token}>
+                <TableCell>
+                  <Text>{item.key_name}</Text>
+                </TableCell>
+                <TableCell>
+                  <Text>{item.spend}</Text>
+                </TableCell>
+                <TableCell>
+                  {item.expires != null ? (
+                    <Text>{item.expires}</Text>
+                  ) : (
+                    <Text>Never expires</Text>
+                  )}
+                </TableCell>
+                <TableCell>
+                  <Icon
+                    onClick={() => handleDelete(item.token)}
+                    icon={TrashIcon}
+                    size="xs"
+                  />
+                </TableCell>
+              </TableRow>
+            );
+          })}
         </TableBody>
       </Table>
     </Card>