From f2a7e2ee98a0dc8fd356931750d74cee8f8d68c0 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 5 Feb 2024 19:28:57 -0800 Subject: [PATCH] feat(ui): enable admin to view all valid keys created on the proxy --- litellm/proxy/proxy_server.py | 68 ++++++- litellm/proxy/utils.py | 17 +- .../src/components/navbar.tsx | 72 ++++---- .../src/components/networking.tsx | 82 ++++----- .../src/components/user_dashboard.tsx | 96 +++++----- .../src/components/view_key_spend_report.tsx | 173 +++++++++++------- .../src/components/view_key_table.tsx | 13 +- 7 files changed, 312 insertions(+), 209 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 289a36cb2b..0fe6997eec 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -322,6 +322,7 @@ async def user_api_key_auth( f"Malformed API Key passed in. Ensure Key has `Bearer ` prefix. Passed in: {passed_in_key}" ) + ### CHECK IF ADMIN ### # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead is_master_key_valid = secrets.compare_digest(api_key, master_key) if is_master_key_valid: @@ -454,6 +455,12 @@ async def user_api_key_auth( if _user is None: continue assert isinstance(_user, dict) + # check if user is admin # + if ( + _user.get("user_role", None) is not None + and _user.get("user_role") == "proxy_admin" + ): + return UserAPIKeyAuth(api_key=master_key) # Token exists, not expired now check if its in budget for the user user_max_budget = _user.get("max_budget", None) user_current_spend = _user.get("spend", None) @@ -597,10 +604,13 @@ async def user_api_key_auth( # check if user can access this route query_params = request.query_params user_id = query_params.get("user_id") + verbose_proxy_logger.debug( + f"user_id: {user_id} & valid_token.user_id: {valid_token.user_id}" + ) if user_id != valid_token.user_id: raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, - detail="user not allowed to access this key's info", + detail="key not allowed to access this user's info", ) elif route == "/user/update": raise HTTPException( @@ -1846,6 +1856,9 @@ async def startup_event(): if prisma_client is not None and master_key is not None: # add master key to db + user_id = "default_user_id" + if os.getenv("PROXY_ADMIN_ID", None) is not None: + user_id = os.getenv("PROXY_ADMIN_ID") asyncio.create_task( generate_key_helper_fn( duration=None, @@ -1854,7 +1867,8 @@ async def startup_event(): config={}, spend=0, token=master_key, - user_id="default_user_id", + user_id=user_id, + user_role="proxy_admin", ) ) @@ -3380,12 +3394,13 @@ async def auth_callback(request: Request): result = await microsoft_sso.verify_and_process(request) # User is Authe'd in - generate key for the UI to access Proxy - user_id = getattr(result, "email", None) + user_email = getattr(result, "email", None) + user_id = getattr(result, "id", None) if user_id is None: user_id = getattr(result, "first_name", "") + getattr(result, "last_name", "") response = await generate_key_helper_fn( - **{"duration": "1hr", "key_max_budget": 0, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id, "team_id": "litellm-dashboard"} # type: ignore + **{"duration": "1hr", "key_max_budget": 0, "models": [], "aliases": {}, "config": {}, "spend": 0, "user_id": user_id, "team_id": "litellm-dashboard", "user_email": user_email} # type: ignore ) key = response["token"] # type: ignore @@ -3393,10 +3408,25 @@ async def auth_callback(request: Request): litellm_dashboard_ui = "/ui/" + user_role = "app_owner" + if ( + os.getenv("PROXY_ADMIN_ID", None) is not None + and os.environ["PROXY_ADMIN_ID"] == user_id + ): + # checks if user is admin + user_role = "app_admin" + import jwt jwt_token = jwt.encode( - {"user_id": user_id, "key": key}, "secret", algorithm="HS256" + { + "user_id": user_id, + "key": key, + "user_email": user_email, + "user_role": user_role, + }, + "secret", + algorithm="HS256", ) litellm_dashboard_ui += "?userID=" + user_id + "&token=" + jwt_token @@ -3409,10 +3439,18 @@ async def auth_callback(request: Request): "/user/info", tags=["user management"], dependencies=[Depends(user_api_key_auth)] ) async def user_info( - user_id: str = fastapi.Query(..., description="User ID in the request parameters") + user_id: Optional[str] = fastapi.Query( + default=None, description="User ID in the request parameters" + ) ): """ Use this to get user information. (user row + all user key info) + + Example request + ``` + curl -X GET 'http://localhost:8000/user/info?user_id=krrish7%40berri.ai' \ + --header 'Authorization: Bearer sk-1234' + ``` """ global prisma_client try: @@ -3421,11 +3459,25 @@ async def user_info( f"Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) ## GET USER ROW ## - user_info = await prisma_client.get_data(user_id=user_id) + if user_id is not None: + user_info = await prisma_client.get_data(user_id=user_id) + else: + user_info = None ## GET ALL KEYS ## keys = await prisma_client.get_data( - user_id=user_id, table_name="key", query_type="find_all" + user_id=user_id, + table_name="key", + query_type="find_all", + expires=datetime.now(), ) + + if user_info is None: + ## make sure we still return a total spend ## + spend = 0 + for k in keys: + spend += getattr(k, "spend", 0) + user_info = {"spend": spend} + ## REMOVE HASHED TOKEN INFO before returning ## for key in keys: try: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 84b09d7265..62cbc6b4be 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -559,9 +559,20 @@ class PrismaClient: # The asterisk before `user_id_list` unpacks the list into separate arguments response = await self.db.query_raw(sql_query) elif query_type == "find_all": - response = await self.db.litellm_usertable.find_many( # type: ignore - order={"spend": "desc"}, - ) + if expires is not None: + response = await self.db.litellm_usertable.find_many( # type: ignore + order={"spend": "desc"}, + where={ # type:ignore + "OR": [ + {"expires": None}, # type:ignore + {"expires": {"gt": expires}}, # type:ignore + ], + }, + ) + else: + response = await self.db.litellm_usertable.find_many( # type: ignore + order={"spend": "desc"}, + ) return response elif table_name == "spend": verbose_proxy_logger.debug( diff --git a/ui/litellm-dashboard/src/components/navbar.tsx b/ui/litellm-dashboard/src/components/navbar.tsx index b7cb357304..946cfc4472 100644 --- a/ui/litellm-dashboard/src/components/navbar.tsx +++ b/ui/litellm-dashboard/src/components/navbar.tsx @@ -1,40 +1,50 @@ "use client"; -import Link from 'next/link'; -import Image from 'next/image' -import React, { useState } from 'react'; +import Link from "next/link"; +import Image from "next/image"; +import React, { useState } from "react"; import { useSearchParams } from "next/navigation"; -import { Button, Text, Metric,Title, TextInput, Grid, Col, Card } from "@tremor/react"; +import { + Button, + Text, + Metric, + Title, + TextInput, + Grid, + Col, + Card, +} from "@tremor/react"; // Define the props type interface NavbarProps { - userID: string | null; - userRole: string | null; + userID: string | null; + userRole: string | null; + userEmail: string | null; } -const Navbar: React.FC = ({ userID, userRole }) => { - console.log("User ID:", userID); +const Navbar: React.FC = ({ userID, userRole, userEmail }) => { + console.log("User ID:", userID); + console.log("userEmail:", userEmail); - return ( - - ) -} + return ( + + ); +}; -export default Navbar; \ No newline at end of file +export default Navbar; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 4763e475e7..5b8e422864 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -1,24 +1,24 @@ /** * Helper file for calls being made to proxy */ -import { message } from 'antd'; +import { message } from "antd"; -const proxyBaseUrl = null; -// const proxyBaseUrl = "http://localhost:4000" // http://localhost:4000 +const isLocal = process.env.NODE_ENV === "development"; +const proxyBaseUrl = isLocal ? "http://localhost:4000" : null; export const keyCreateCall = async ( accessToken: string, userID: string, - formValues: Record, // Assuming formValues is an object + formValues: Record // Assuming formValues is an object ) => { try { - console.log("Form Values in keyCreateCall:", formValues); // Log the form values before making the API call - + console.log("Form Values in keyCreateCall:", formValues); // Log the form values before making the API call + // check if formValues.description is not undefined, make it a string and add it to formValues.metadata if (formValues.description) { // add to formValues.metadata if (!formValues.metadata) { - formValues.metadata = {} + formValues.metadata = {}; } // value needs to be in "", valid JSON formValues.metadata.description = formValues.description; @@ -26,7 +26,7 @@ export const keyCreateCall = async ( delete formValues.description; formValues.metadata = JSON.stringify(formValues.metadata); } - // if formValues.metadata is not undefined, make it a valid dict + // if formValues.metadata is not undefined, make it a valid dict if (formValues.metadata) { console.log("formValues.metadata:", formValues.metadata); // if there's an exception JSON.parse, show it in the message @@ -69,15 +69,11 @@ export const keyCreateCall = async ( } }; - -export const keyDeleteCall = async ( - accessToken: String, - user_key: String -) => { +export const keyDeleteCall = async (accessToken: String, user_key: String) => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/key/delete` : `/key/delete`; - console.log("in keyDeleteCall:", user_key) - + console.log("in keyDeleteCall:", user_key); + const response = await fetch(url, { method: "POST", headers: { @@ -108,21 +104,22 @@ export const keyDeleteCall = async ( export const userInfoCall = async ( accessToken: String, - userID: String + userID: String, + userRole: String ) => { try { - const url = proxyBaseUrl ? `${proxyBaseUrl}/user/info` : `/user/info`; - console.log("in userInfoCall:", url) - const response = await fetch( - `${url}/?user_id=${userID}`, - { - method: "GET", - headers: { - Authorization: `Bearer ${accessToken}`, - "Content-Type": "application/json", - }, - } - ); + let url = proxyBaseUrl ? `${proxyBaseUrl}/user/info` : `/user/info`; + if (userRole == "App Owner") { + url = `${url}/?user_id=${userID}`; + } + message.info("Requesting user data"); + const response = await fetch(url, { + method: "GET", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); if (!response.ok) { const errorData = await response.text(); @@ -131,7 +128,7 @@ export const userInfoCall = async ( } const data = await response.json(); - console.log(data); + message.info("Received user data"); return data; // Handle success - you might want to update some state or UI based on the created key } catch (error) { @@ -140,24 +137,17 @@ export const userInfoCall = async ( } }; - -export const keySpendLogsCall = async ( - accessToken: String, - token: String -) => { +export const keySpendLogsCall = async (accessToken: String, token: String) => { try { const url = proxyBaseUrl ? `${proxyBaseUrl}/spend/logs` : `/spend/logs`; - console.log("in keySpendLogsCall:", url) - const response = await fetch( - `${url}/?api_key=${token}`, - { - method: "GET", - headers: { - Authorization: `Bearer ${accessToken}`, - "Content-Type": "application/json", - }, - } - ); + console.log("in keySpendLogsCall:", url); + const response = await fetch(`${url}/?api_key=${token}`, { + method: "GET", + headers: { + Authorization: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); if (!response.ok) { const errorData = await response.text(); message.error(errorData); @@ -171,4 +161,4 @@ export const keySpendLogsCall = async ( console.error("Failed to create key:", error); throw error; } -} +}; diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 951d0287bf..b1a06939b1 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -6,21 +6,25 @@ import CreateKey from "./create_key_button"; import ViewKeyTable from "./view_key_table"; import ViewUserSpend from "./view_user_spend"; import EnterProxyUrl from "./enter_proxy_url"; +import { message } from "antd"; import Navbar from "./navbar"; import { useSearchParams } from "next/navigation"; import { jwtDecode } from "jwt-decode"; -const proxyBaseUrl = null; -// const proxyBaseUrl = "http://localhost:4000" // http://localhost:4000 +const isLocal = process.env.NODE_ENV === "development"; +console.log("isLocal:", isLocal); +const proxyBaseUrl = isLocal ? "http://localhost:4000" : null; type UserSpendData = { spend: number; max_budget?: number | null; -} +}; const UserDashboard = () => { const [data, setData] = useState(null); // Keep the initialization of state here - const [userSpendData, setUserSpendData] = useState(null); + const [userSpendData, setUserSpendData] = useState( + null + ); // Assuming useSearchParams() hook exists and works in your setup const searchParams = useSearchParams(); @@ -30,19 +34,19 @@ const UserDashboard = () => { const token = searchParams.get("token"); const [accessToken, setAccessToken] = useState(null); const [userRole, setUserRole] = useState(null); - + const [userEmail, setUserEmail] = useState(null); function formatUserRole(userRole: string) { if (!userRole) { return "Undefined Role"; } - + console.log(`Received user role: ${userRole}`); switch (userRole.toLowerCase()) { case "app_owner": return "App Owner"; case "demo_app_owner": - return "AppOwner"; - case "admin": + return "App Owner"; + case "app_admin": return "Admin"; case "app_user": return "App User"; @@ -53,7 +57,7 @@ const UserDashboard = () => { // Moved useEffect inside the component and used a condition to run fetch only if the params are available useEffect(() => { - if (token){ + if (token) { const decoded = jwtDecode(token) as { [key: string]: any }; if (decoded) { // cast decoded to dictionary @@ -71,17 +75,19 @@ const UserDashboard = () => { } else { console.log("User role not defined"); } + + if (decoded.user_email) { + setUserEmail(decoded.user_email); + } else { + console.log(`User Email is not set ${decoded}`); + } } } - if (userID && accessToken && !data) { + if (userID && accessToken && userRole && !data) { const fetchData = async () => { try { - const response = await userInfoCall( - accessToken, - userID - ); - console.log("Response:", response); - setUserSpendData(response["user_info"]) + const response = await userInfoCall(accessToken, userID, userRole); + setUserSpendData(response["user_info"]); setData(response["keys"]); // Assuming this is the correct path to your data } catch (error) { console.error("There was an error fetching the data", error); @@ -93,53 +99,45 @@ const UserDashboard = () => { }, [userID, token, accessToken, data]); if (userID == null || token == null) { - - // Now you can construct the full URL - const url = proxyBaseUrl ? `${proxyBaseUrl}/sso/key/generate` : `/sso/key/generate`; + const url = proxyBaseUrl + ? `${proxyBaseUrl}/sso/key/generate` + : `/sso/key/generate`; console.log("Full URL:", url); window.location.href = url; return null; - } - else if (accessToken == null) { + } else if (accessToken == null) { return null; } if (userRole == null) { - setUserRole("App Owner") + setUserRole("App Owner"); } - + return (
- + - - - - - - + + + + + +
- ); }; -export default UserDashboard; \ No newline at end of file +export default UserDashboard; diff --git a/ui/litellm-dashboard/src/components/view_key_spend_report.tsx b/ui/litellm-dashboard/src/components/view_key_spend_report.tsx index 40961325ec..e90401e5bf 100644 --- a/ui/litellm-dashboard/src/components/view_key_spend_report.tsx +++ b/ui/litellm-dashboard/src/components/view_key_spend_report.tsx @@ -1,8 +1,26 @@ "use client"; import React, { useState, useEffect } from "react"; -import { Button as Button2, Modal, Form, Input, InputNumber, Select, message } from "antd"; -import { Button, Text, Card, Table, BarChart, Title, Subtitle, BarList, Metric } from "@tremor/react"; +import { + Button as Button2, + Modal, + Form, + Input, + InputNumber, + Select, + message, +} from "antd"; +import { + Button, + Text, + Card, + Table, + BarChart, + Title, + Subtitle, + BarList, + Metric, +} from "@tremor/react"; import { keySpendLogsCall } from "./networking"; interface ViewKeySpendReportProps { @@ -14,18 +32,30 @@ interface ViewKeySpendReportProps { } type ResponseValueType = { - startTime: string; // Assuming startTime is a string, adjust it if it's of a different type - spend: number; // Assuming spend is a number, adjust it if it's of a different type - user: string; // Assuming user is a string, adjust it if it's of a different type - }; + startTime: string; // Assuming startTime is a string, adjust it if it's of a different type + spend: number; // Assuming spend is a number, adjust it if it's of a different type + user: string; // Assuming user is a string, adjust it if it's of a different type +}; -const ViewKeySpendReport: React.FC = ({ token, accessToken, keySpend, keyBudget, keyName }) => { +const ViewKeySpendReport: React.FC = ({ + token, + accessToken, + keySpend, + keyBudget, + keyName, +}) => { const [isModalVisible, setIsModalVisible] = useState(false); - const [data, setData] = useState<{ day: string; spend: number; }[] | null>(null); - const [userData, setUserData] = useState<{ name: string; value: number; }[] | null>(null); + const [data, setData] = useState<{ day: string; spend: number }[] | null>( + null + ); + const [userData, setUserData] = useState< + { name: string; value: number }[] | null + >(null); const showModal = () => { + console.log("Show Modal triggered"); setIsModalVisible(true); + fetchData(); }; const handleOk = () => { @@ -41,68 +71,79 @@ const ViewKeySpendReport: React.FC = ({ token, accessTo try { if (accessToken == null || token == null) { return; - } - const response = await keySpendLogsCall(accessToken=accessToken, token=token); + } + console.log(`accessToken: ${accessToken}; token: ${token}`); + const response = await keySpendLogsCall( + (accessToken = accessToken), + (token = token) + ); console.log("Response:", response); // loop through response // get spend, startTime for each element, place in new array - - const pricePerDay: Record = (Object.values(response) as ResponseValueType[]).reduce((acc: Record, value) => { + const pricePerDay: Record = ( + Object.values(response) as ResponseValueType[] + ).reduce((acc: Record, value) => { const startTime = new Date(value.startTime); - const day = new Intl.DateTimeFormat('en-US', { day: '2-digit', month: 'short' }).format(startTime); - + const day = new Intl.DateTimeFormat("en-US", { + day: "2-digit", + month: "short", + }).format(startTime); + acc[day] = (acc[day] || 0) + value.spend; - + return acc; }, {}); - - + // sort pricePerDay by day // Convert object to array of key-value pairs - const pricePerDayArray = Object.entries(pricePerDay); + const pricePerDayArray = Object.entries(pricePerDay); - // Sort the array based on the date (key) - pricePerDayArray.sort(([aKey], [bKey]) => { - const dateA = new Date(aKey); - const dateB = new Date(bKey); - return dateA.getTime() - dateB.getTime(); - }); - - // Convert the sorted array back to an object - const sortedPricePerDay = Object.fromEntries(pricePerDayArray); + // Sort the array based on the date (key) + pricePerDayArray.sort(([aKey], [bKey]) => { + const dateA = new Date(aKey); + const dateB = new Date(bKey); + return dateA.getTime() - dateB.getTime(); + }); + // Convert the sorted array back to an object + const sortedPricePerDay = Object.fromEntries(pricePerDayArray); console.log(sortedPricePerDay); - - const pricePerUser: Record = (Object.values(response) as ResponseValueType[]).reduce((acc: Record, value) => { + + const pricePerUser: Record = ( + Object.values(response) as ResponseValueType[] + ).reduce((acc: Record, value) => { const user = value.user; acc[user] = (acc[user] || 0) + value.spend; - + return acc; }, {}); - - + console.log(pricePerDay); console.log(pricePerUser); const arrayBarChart = []; - // [ - // { - // "day": "02 Feb", - // "spend": pricePerDay["02 Feb"], - // } - // ] + // [ + // { + // "day": "02 Feb", + // "spend": pricePerDay["02 Feb"], + // } + // ] for (const [key, value] of Object.entries(sortedPricePerDay)) { arrayBarChart.push({ day: key, spend: value }); } - // get 5 most expensive users - const sortedUsers = Object.entries(pricePerUser).sort((a, b) => b[1] - a[1]); + const sortedUsers = Object.entries(pricePerUser).sort( + (a, b) => b[1] - a[1] + ); const top5Users = sortedUsers.slice(0, 5); - const userChart = top5Users.map(([key, value]) => ({ name: key, value: value })); - + const userChart = top5Users.map(([key, value]) => ({ + name: key, + value: value, + })); + setData(arrayBarChart); setUserData(userChart); console.log("arrayBarChart:", arrayBarChart); @@ -112,11 +153,10 @@ const ViewKeySpendReport: React.FC = ({ token, accessTo } }; - useEffect(() => { - // Fetch data only when the token changes - fetchData(); - }, [token]); // Dependency array containing the 'token' variable - + // useEffect(() => { + // // Fetch data only when the token changes + // fetchData(); + // }, [token]); // Dependency array containing the 'token' variable if (!token) { return null; @@ -134,33 +174,28 @@ const ViewKeySpendReport: React.FC = ({ token, accessTo onCancel={handleCancel} footer={null} > - Key Name: {keyName} + Key Name: {keyName} Monthly Spend ${keySpend} - {data && ( + {data && ( - )} - - Top 5 Users Spend (USD) - - {userData && ( - - )} - - + )} + + Top 5 Users Spend (USD) + + {userData && ( + + )} + ); diff --git a/ui/litellm-dashboard/src/components/view_key_table.tsx b/ui/litellm-dashboard/src/components/view_key_table.tsx index 8522a6bb16..4813bbe4e3 100644 --- a/ui/litellm-dashboard/src/components/view_key_table.tsx +++ b/ui/litellm-dashboard/src/components/view_key_table.tsx @@ -1,5 +1,5 @@ "use client"; -import React, { useEffect } from "react"; +import React, { useEffect, useState } from "react"; import { keyDeleteCall } from "./networking"; import { StatusOnlineIcon, TrashIcon } from "@heroicons/react/outline"; import { @@ -32,6 +32,8 @@ const ViewKeyTable: React.FC = ({ data, setData, }) => { + const [isButtonClicked, setIsButtonClicked] = useState(false); + const handleDelete = async (token: String) => { if (data == null) { return; @@ -116,8 +118,13 @@ const ViewKeyTable: React.FC = ({ /> - - + );