diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ca5d4b05b..0d7355dad 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -187,8 +187,8 @@ class UpdateKeyRequest(GenerateKeyRequest): metadata: Optional[dict] = None -class DeleteKeyRequest(LiteLLMBase): - keys: List +class KeyRequest(LiteLLMBase): + keys: List[str] class NewUserRequest(GenerateKeyRequest): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 33eaae472..c5c843177 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. @@ -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)] ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 35b647257..bb2625dce 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(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 diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index d4f405b7b..89d4f4d3e 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/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) diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 56c5f6db1..4521ca446 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/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`; diff --git a/ui/litellm-dashboard/src/components/usage.tsx b/ui/litellm-dashboard/src/components/usage.tsx index 91cd57bcf..84c942eec 100644 --- a/ui/litellm-dashboard/src/components/usage.tsx +++ b/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 = ({ if (accessToken && token && userRole && userID) { const fetchData = async () => { try { - const response = await userSpendLogsCall( + await userSpendLogsCall( accessToken, token, userRole, userID, startTime, endTime - ); - - setTopKeys(getTopKeys(response)); - setTopUsers(getTopUsers(response)); - setKeySpendData(response); + ).then(async (response) => { + const topKeysResponse = await keyInfoCall( + accessToken, + getTopKeys(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 = ({ } }, [accessToken, token, userRole, userID, startTime, endTime]); - topUsers.forEach((obj) => { - Object.values(obj).forEach((value) => { - console.log(value); - }); - }); return (
@@ -217,7 +223,7 @@ const UsagePage: React.FC = ({ index="key" categories={["spend"]} colors={["blue"]} - yAxisWidth={200} + yAxisWidth={80} tickGap={5} layout="vertical" showXAxis={false}