diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9b7ad75de0..0d7355dad7 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -188,7 +188,7 @@ class UpdateKeyRequest(GenerateKeyRequest): class KeyRequest(LiteLLMBase): - keys: List + keys: List[str] class NewUserRequest(GenerateKeyRequest): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 53162d29c7..52420611ea 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2954,11 +2954,75 @@ async def delete_key_fn(data: KeyRequest): ) +@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. + 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)] ) async def info_key_fn( - data: Optional[KeyRequest] = None, key: Optional[str] = fastapi.Query( default=None, description="Key in the request parameters" ), @@ -2997,24 +3061,11 @@ 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 is None and data is None: + if key 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 41646149b6..bb2625dced 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -556,7 +556,7 @@ class PrismaClient: new_token = self.hash_token(token=t) hashed_tokens.append(new_token) else: - hashed_tokens.append(new_token) + hashed_tokens.append(t) where_filter["token"]["in"] = hashed_tokens response = await self.db.litellm_verificationtoken.find_many( order={"spend": "desc"}, where=where_filter # type: ignore diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 7026f394a7..06a4360535 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -200,6 +200,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 7b113e8125..76c8d7fddc 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 { @@ -79,7 +79,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; } @@ -168,18 +168,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 @@ -187,19 +198,8 @@ 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) => { - Object.values(obj).forEach((value) => { - console.log(value); - }); - }); return (
@@ -227,7 +227,7 @@ const UsagePage: React.FC = ({ index="key" categories={["spend"]} colors={["blue"]} - yAxisWidth={200} + yAxisWidth={80} tickGap={5} layout="vertical" showXAxis={false}