Merge pull request #1926 from BerriAI/litellm_ui_further_improvements

Enable viewing key alias instead of hashed tokens

litellm/proxy/_types.py

@@ -187,8 +187,8 @@ class UpdateKeyRequest(GenerateKeyRequest):
     metadata: Optional[dict] = None
 
 
-class DeleteKeyRequest(LiteLLMBase):
+class KeyRequest(LiteLLMBase):
-    keys: List
+    keys: List[str]
 
 
 class NewUserRequest(GenerateKeyRequest):

litellm/proxy/proxy_server.py

@@ -2894,7 +2894,7 @@ async def update_key_fn(request: Request, data: UpdateKeyRequest):
 @router.post(
     "/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
 )
-async def delete_key_fn(data: DeleteKeyRequest):
+async def delete_key_fn(data: KeyRequest):
     """
     Delete a key from the key management system.
 
@@ -2954,6 +2954,73 @@ async def delete_key_fn(data: DeleteKeyRequest):
         )
 
 
+@router.post(
+    "/v2/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
+)
+async def info_key_fn_v2(
+    data: Optional[KeyRequest] = None,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Retrieve information about a list of keys.
+
+    **New endpoint**. Currently admin only.
+    Parameters:
+        keys: Optional[list] = body parameter representing the key(s) in the request
+        user_api_key_dict: UserAPIKeyAuth = Dependency representing the user's API key
+    Returns:
+        Dict containing the key and its associated information
+    
+    Example Curl:
+    ```
+    curl -X GET "http://0.0.0.0:8000/key/info" \
+-H "Authorization: Bearer sk-1234" \
+-d {"keys": ["sk-1", "sk-2", "sk-3"]}
+    ```
+    """
+    global prisma_client
+    try:
+        if prisma_client is None:
+            raise Exception(
+                f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
+            )
+        if data is None:
+            raise HTTPException(
+                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+                detail={"message": "Malformed request. No keys passed in."},
+            )
+
+        key_info = await prisma_client.get_data(
+            token=data.keys, table_name="key", query_type="find_all"
+        )
+        filtered_key_info = []
+        for k in key_info:
+            try:
+                k = k.model_dump()  # noqa
+            except:
+                # if using pydantic v1
+                k = k.dict()
+            filtered_key_info.append(k)
+        return {"key": data.keys, "info": filtered_key_info}
+
+    except Exception as e:
+        if isinstance(e, HTTPException):
+            raise ProxyException(
+                message=getattr(e, "detail", f"Authentication Error({str(e)})"),
+                type="auth_error",
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST),
+            )
+        elif isinstance(e, ProxyException):
+            raise e
+        raise ProxyException(
+            message="Authentication Error, " + str(e),
+            type="auth_error",
+            param=getattr(e, "param", "None"),
+            code=status.HTTP_400_BAD_REQUEST,
+        )
+
+
 @router.get(
     "/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
 )

litellm/proxy/utils.py

@@ -483,7 +483,7 @@ class PrismaClient:
     )
     async def get_data(
         self,
-        token: Optional[str] = None,
+        token: Optional[Union[str, list]] = None,
         user_id: Optional[str] = None,
         user_id_list: Optional[list] = None,
         key_val: Optional[dict] = None,
@@ -497,12 +497,13 @@ class PrismaClient:
             if token is not None or (table_name is not None and table_name == "key"):
                 # check if plain text or hash
                 if token is not None:
-                    hashed_token = token
+                    if isinstance(token, str):
-                    if token.startswith("sk-"):
+                        hashed_token = token
-                        hashed_token = self.hash_token(token=token)
+                        if token.startswith("sk-"):
-                    verbose_proxy_logger.debug(
+                            hashed_token = self.hash_token(token=token)
-                        f"PrismaClient: find_unique for token: {hashed_token}"
+                        verbose_proxy_logger.debug(
-                    )
+                            f"PrismaClient: find_unique for token: {hashed_token}"
+                        )
                 if query_type == "find_unique":
                     response = await self.db.litellm_verificationtoken.find_unique(
                         where={"token": hashed_token}
@@ -540,8 +541,25 @@ class PrismaClient:
                             if isinstance(r.expires, datetime):
                                 r.expires = r.expires.isoformat()
                 elif query_type == "find_all":
+                    where_filter: dict = {}
+                    if token is not None:
+                        where_filter["token"] = {}
+                        if isinstance(token, str):
+                            if token.startswith("sk-"):
+                                token = self.hash_token(token=token)
+                            where_filter["token"]["in"] = [token]
+                        elif isinstance(token, list):
+                            hashed_tokens = []
+                            for t in token:
+                                assert isinstance(t, str)
+                                if t.startswith("sk-"):
+                                    new_token = self.hash_token(token=t)
+                                    hashed_tokens.append(new_token)
+                                else:
+                                    hashed_tokens.append(t)
+                            where_filter["token"]["in"] = hashed_tokens
                     response = await self.db.litellm_verificationtoken.find_many(
-                        order={"spend": "desc"},
+                        order={"spend": "desc"}, where=where_filter  # type: ignore
                     )
                 if response is not None:
                     return response

litellm/tests/test_key_generate_prisma.py

@@ -59,7 +59,7 @@ from litellm.proxy._types import (
     NewUserRequest,
     GenerateKeyRequest,
     DynamoDBArgs,
-    DeleteKeyRequest,
+    KeyRequest,
     UpdateKeyRequest,
     GenerateKeyRequest,
 )
@@ -740,7 +740,7 @@ def test_delete_key(prisma_client):
             generated_key = key.key
             bearer_token = "Bearer " + generated_key
 
-            delete_key_request = DeleteKeyRequest(keys=[generated_key])
+            delete_key_request = KeyRequest(keys=[generated_key])
 
             # delete the key
             result_delete_key = await delete_key_fn(data=delete_key_request)
@@ -778,7 +778,7 @@ def test_delete_key_auth(prisma_client):
             generated_key = key.key
             bearer_token = "Bearer " + generated_key
 
-            delete_key_request = DeleteKeyRequest(keys=[generated_key])
+            delete_key_request = KeyRequest(keys=[generated_key])
 
             # delete the key
             result_delete_key = await delete_key_fn(data=delete_key_request)
@@ -839,7 +839,7 @@ def test_generate_and_call_key_info(prisma_client):
             }
 
             # cleanup - delete key
-            delete_key_request = DeleteKeyRequest(keys=[generated_key])
+            delete_key_request = KeyRequest(keys=[generated_key])
 
             # delete the key
             await delete_key_fn(data=delete_key_request)
@@ -908,7 +908,7 @@ def test_generate_and_update_key(prisma_client):
             assert result["info"]["models"] == ["ada", "babbage", "curie", "davinci"]
 
             # cleanup - delete key
-            delete_key_request = DeleteKeyRequest(keys=[generated_key])
+            delete_key_request = KeyRequest(keys=[generated_key])
 
             # delete the key
             await delete_key_fn(data=delete_key_request)

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

@@ -202,6 +202,36 @@ export const userSpendLogsCall = async (
   }
 };
 
+export const keyInfoCall = async (accessToken: String, keys: String[]) => {
+  try {
+    let url = proxyBaseUrl ? `${proxyBaseUrl}/v2/key/info` : `/key/info`;
+
+    const response = await fetch(url, {
+      method: "POST",
+      headers: {
+        Authorization: `Bearer ${accessToken}`,
+        "Content-Type": "application/json",
+      },
+      body: JSON.stringify({
+        keys: keys,
+      }),
+    });
+
+    if (!response.ok) {
+      const errorData = await response.text();
+      message.error(errorData);
+      throw new Error("Network response was not ok");
+    }
+
+    const data = await response.json();
+    console.log(data);
+    return data;
+  } catch (error) {
+    console.error("Failed to create key:", error);
+    throw error;
+  }
+};
+
 export const spendUsersCall = async (accessToken: String, userID: String) => {
   try {
     const url = proxyBaseUrl ? `${proxyBaseUrl}/spend/users` : `/spend/users`;

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

@@ -2,7 +2,7 @@ import { BarChart, Card, Title } from "@tremor/react";
 
 import React, { useState, useEffect } from "react";
 import { Grid, Col, Text, LineChart } from "@tremor/react";
-import { userSpendLogsCall } from "./networking";
+import { userSpendLogsCall, keyInfoCall } from "./networking";
 import { start } from "repl";
 
 interface UsagePageProps {
@@ -75,7 +75,7 @@ function getTopKeys(data: Array<{ [key: string]: unknown }>): any[] {
 
   spendKeys.sort((a, b) => Number(b.spend) - Number(a.spend));
 
-  const topKeys = spendKeys.slice(0, 5);
+  const topKeys = spendKeys.slice(0, 5).map((k) => k.key);
   console.log(`topKeys: ${Object.keys(topKeys[0])}`);
   return topKeys;
 }
@@ -164,18 +164,29 @@ const UsagePage: React.FC<UsagePageProps> = ({
     if (accessToken && token && userRole && userID) {
       const fetchData = async () => {
         try {
-          const response = await userSpendLogsCall(
+          await userSpendLogsCall(
             accessToken,
             token,
             userRole,
             userID,
             startTime,
             endTime
-          );
+          ).then(async (response) => {
-
+            const topKeysResponse = await keyInfoCall(
-          setTopKeys(getTopKeys(response));
+              accessToken,
-          setTopUsers(getTopUsers(response));
+              getTopKeys(response)
-          setKeySpendData(response);
+            );
+            const filtered_keys = topKeysResponse["info"].map((k: any) => ({
+              key: (k["key_name"] || k["key_alias"] || k["token"]).substring(
+                0,
+                7
+              ),
+              spend: k["spend"],
+            }));
+            setTopKeys(filtered_keys);
+            setTopUsers(getTopUsers(response));
+            setKeySpendData(response);
+          });
         } catch (error) {
           console.error("There was an error fetching the data", error);
           // Optionally, update your UI to reflect the error state here as well
@@ -185,11 +196,6 @@ const UsagePage: React.FC<UsagePageProps> = ({
     }
   }, [accessToken, token, userRole, userID, startTime, endTime]);
 
-  topUsers.forEach((obj) => {
-    Object.values(obj).forEach((value) => {
-      console.log(value);
-    });
-  });
   return (
     <div style={{ width: "100%" }}>
       <Grid numItems={2} className="gap-2 p-10 h-[75vh] w-full">
@@ -217,7 +223,7 @@ const UsagePage: React.FC<UsagePageProps> = ({
               index="key"
               categories={["spend"]}
               colors={["blue"]}
-              yAxisWidth={200}
+              yAxisWidth={80}
               tickGap={5}
               layout="vertical"
               showXAxis={false}