diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ca5d4b05b1e..9b7ad75de0b 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -187,7 +187,7 @@ class UpdateKeyRequest(GenerateKeyRequest): metadata: Optional[dict] = None -class DeleteKeyRequest(LiteLLMBase): +class KeyRequest(LiteLLMBase): keys: List diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 33eaae47251..53162d29c73 100644 --- a/litellm/proxy/proxy_server.py +++ b/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): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 35b6472577a..41646149b61 100644 --- a/litellm/proxy/utils.py +++ b/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 diff --git a/ui/litellm-dashboard/src/components/usage.tsx b/ui/litellm-dashboard/src/components/usage.tsx index 4b80332a78c..7b113e8125c 100644 --- a/ui/litellm-dashboard/src/components/usage.tsx +++ b/ui/litellm-dashboard/src/components/usage.tsx @@ -187,6 +187,12 @@ const UsagePage: React.FC = ({ }; 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) => {