Merge pull request #1807 from BerriAI/litellm_view_key_spend

[UI] View Key Spend Reports 🤠

litellm/proxy/proxy_server.py

@@ -600,6 +600,24 @@ async def user_api_key_auth(
                     pass
                 elif allow_user_auth == True and route == "/key/delete":
                     pass
+                elif route == "/spend/logs":
+                    # check if user can access this route
+                    # user can only access this route if
+                    # - api_key they need logs for has the same user_id as the one used for auth
+                    query_params = request.query_params
+                    api_key = query_params.get(
+                        "api_key"
+                    )  # UI, will only pass hashed tokens
+                    token_info = await prisma_client.get_data(
+                        token=api_key, table_name="key", query_type="find_unique"
+                    )
+                    if secrets.compare_digest(token_info.user_id, valid_token.user_id):
+                        pass
+                    else:
+                        raise HTTPException(
+                            status_code=status.HTTP_403_FORBIDDEN,
+                            detail="user not allowed to access this key's info",
+                        )
                 else:
                     raise Exception(
                         f"Only master key can be used to generate, delete, update or get info for new keys/users. Value of allow_user_auth={allow_user_auth}"

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

@@ -11,7 +11,7 @@ const { Option } = Select;
 
 interface CreateKeyProps {
   userID: string;
-  userRole: string;
+  userRole: string | null;
   accessToken: string;
   data: any[] | null;
   setData: React.Dispatch<React.SetStateAction<any[] | null>>;

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

@@ -139,3 +139,36 @@ export const userInfoCall = async (
     throw error;
   }
 };
+
+
+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",
+        },
+      }
+    );
+    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;
+  }
+}

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

@@ -13,12 +13,19 @@ import { jwtDecode } from "jwt-decode";
 // const proxyBaseUrl = null;
 const proxyBaseUrl = "http://localhost:4000" // http://localhost:4000
 
+type UserSpendData = {
+  spend: number;
+  max_budget?: number | null;
+}
+
 const UserDashboard = () => {
   const [data, setData] = useState<null | any[]>(null); // Keep the initialization of state here
-  const [userSpendData, setUserSpendData] = useState<null | any[]>(null);
+  const [userSpendData, setUserSpendData] = useState<UserSpendData | null>(null);
+
   // Assuming useSearchParams() hook exists and works in your setup
   const searchParams = useSearchParams();
   const userID = searchParams.get("userID");
+  const viewSpend = searchParams.get("viewSpend");
 
   const token = searchParams.get("token");
   const [accessToken, setAccessToken] = useState<string | null>(null);
@@ -99,6 +106,9 @@ const UserDashboard = () => {
     return null;
   }
 
+  if (userRole == null) {
+    setUserRole("App Owner")
+  }
   
   return (
     <div>
@@ -114,7 +124,6 @@ const UserDashboard = () => {
         />
         <ViewKeyTable
           userID={userID}
-          userRole={userRole}
           accessToken={accessToken}
           data={data}
           setData={setData}

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

@@ -0,0 +1,167 @@
+"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 { keySpendLogsCall } from "./networking";
+
+interface ViewKeySpendReportProps {
+  token: string;
+  accessToken: string;
+  keySpend: number;
+  keyBudget: number;
+  keyName: string;
+}
+
+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
+  };
+
+const ViewKeySpendReport: React.FC<ViewKeySpendReportProps> = ({ 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 showModal = () => {
+    setIsModalVisible(true);
+  };
+
+  const handleOk = () => {
+    setIsModalVisible(false);
+  };
+
+  const handleCancel = () => {
+    setIsModalVisible(false);
+  };
+
+  if (!token) {
+    return null;
+  }
+
+  // call keySpendLogsCall and set the data
+  const fetchData = async () => {
+    try {
+      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<string, number> = (Object.values(response) as ResponseValueType[]).reduce((acc: Record<string, number>, value) => {
+        const startTime = new Date(value.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);
+
+        // 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<string, number> = (Object.values(response) as ResponseValueType[]).reduce((acc: Record<string, number>, 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"],
+        // }
+        // ]
+      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 top5Users = sortedUsers.slice(0, 5);
+      const userChart = top5Users.map(([key, value]) => ({ name: key, value: value }));
+      
+      setData(arrayBarChart);
+      setUserData(userChart);
+      console.log("arrayBarChart:", arrayBarChart);
+    } catch (error) {
+      console.error("There was an error fetching the data", error);
+      // Optionally, update your UI to reflect the error state here as well
+    }
+  };
+
+  useEffect(() => {
+    // Fetch data only when the token changes
+    fetchData();
+  }, [token]);  // Dependency array containing the 'token' variable
+
+
+
+  return (
+    <div>
+      <Button className="mx-auto" onClick={showModal}>
+        View Spend Report
+      </Button>
+      <Modal
+        visible={isModalVisible}
+        width={1000}
+        onOk={handleOk}
+        onCancel={handleCancel}
+        footer={null}
+      >
+        <Title style={{ textAlign: 'left' }}>Key Name: {keyName}</Title>
+
+        <Metric>Monthly Spend ${keySpend}</Metric>
+
+        <Card className="mt-6 mb-6">
+        {data && (
+            <BarChart
+                className="mt-6"
+                data={data}
+                colors={["green"]}
+                index="day"
+                categories={["spend"]}
+                yAxisWidth={48}
+            />
+        )}
+    </Card>
+    <Title className="mt-6">Top 5 Users Spend (USD)</Title>
+    <Card className="mb-6">
+        {userData && (
+            <BarList
+            className="mt-6"
+            data={userData}
+            color="teal"
+            />
+        )}
+    </Card>
+
+      </Modal>
+    </div>
+  );
+};
+
+export default ViewKeySpendReport;

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

@@ -6,6 +6,7 @@ import {
   Badge,
   Card,
   Table,
+  Button,
   TableBody,
   TableCell,
   TableHead,
@@ -15,6 +16,7 @@ import {
   Title,
   Icon,
 } from "@tremor/react";
+import ViewKeySpendReport from "./view_key_spend_report";
 
 // Define the props type
 interface ViewKeyTableProps {
@@ -110,9 +112,13 @@ const ViewKeyTable: React.FC<ViewKeyTableProps> = ({
                   <Icon
                     onClick={() => handleDelete(item.token)}
                     icon={TrashIcon}
-                    size="xs"
+                    size="sm"
                   />
                 </TableCell>
+                <TableCell>
+                  <ViewKeySpendReport token={item.token} accessToken={accessToken} keySpend={item.spend} keyBudget={item.max_budget} keyName={item.key_name} />
+                
+                </TableCell>
               </TableRow>
             );
           })}