feat(proxy_server.py): support key lists for /key/info

litellm/proxy/_types.py

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

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.
 
@@ -2958,6 +2958,7 @@ async def delete_key_fn(data: DeleteKeyRequest):
     "/key/info", tags=["key management"], dependencies=[Depends(user_api_key_auth)]
 )
 async def info_key_fn(
+    data: Optional[KeyRequest] = None,
     key: Optional[str] = fastapi.Query(
         default=None, description="Key in the request parameters"
     ),
@@ -2966,16 +2967,23 @@ async def info_key_fn(
     """
     Retrieve information about a key.
     Parameters:
-        key: Optional[str] = Query parameter representing the key in the request
+        key: Optional[str | list] = Query 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:
+    **1 key**
     ```
     curl -X GET "http://0.0.0.0:8000/key/info?key=sk-02Wr4IAlN3NvPXvL5JVvDA" \
 -H "Authorization: Bearer sk-1234"
     ```
+    **multiple keys**
+    ```
+    curl -X GET "http://0.0.0.0:8000/key/info" \
+-H "Authorization: Bearer sk-1234" \
+-d {"keys": ["sk-1", "sk-2", "sk-3"]}
+    ```
 
     Example Curl - if no key is passed, it will use the Key Passed in Authorization Header
     ```
@@ -2989,16 +2997,24 @@ async def info_key_fn(
             raise Exception(
                 f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys"
             )
-        if key == None:
-            key = user_api_key_dict.api_key
-        key_info = await prisma_client.get_data(token=key)
-        ## REMOVE HASHED TOKEN INFO BEFORE RETURNING ##
-        try:
-            key_info = key_info.model_dump()  # noqa
-        except:
-            # if using pydantic v1
-            key_info = key_info.dict()
-        key_info.pop("token")
+        if key is None and data is None:
+            key_info = await prisma_client.get_data(token=user_api_key_dict.api_key)
+
+        if key is not None:  # single key
+            key_info = await prisma_client.get_data(token=key)
+        elif data is not None:  # list of keys
+            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}
         return {"key": key, "info": key_info}
     except Exception as e:
         if isinstance(e, HTTPException):

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 token.startswith("sk-"):
-                        hashed_token = self.hash_token(token=token)
-                    verbose_proxy_logger.debug(
-                        f"PrismaClient: find_unique for token: {hashed_token}"
-                    )
+                    if isinstance(token, str):
+                        hashed_token = token
+                        if token.startswith("sk-"):
+                            hashed_token = self.hash_token(token=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(new_token)
+                            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

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

@@ -187,6 +187,12 @@ const UsagePage: React.FC<UsagePageProps> = ({
       };
       fetchData();
     }
+    if (topKeys.length > 0) {
+      /**
+       * get the key alias / secret names for the created keys
+       * Use that for the key name instead of token hash for easier association
+       */
+    }
   }, [accessToken, token, userRole, userID, startTime, endTime]);
 
   topUsers.forEach((obj) => {