Merge pull request #4185 from BerriAI/litellm_reset_budget

Reset budget option for teams

ui/litellm-dashboard/src/components/budgets/budget_modal.tsx

@@ -122,6 +122,7 @@ const BudgetModal: React.FC<BudgetModalProps> = ({
               >
                 <Select defaultValue={null} placeholder="n/a">
                   <Select.Option value="24h">daily</Select.Option>
+                  <Select.Option value="7d">weekly</Select.Option>
                   <Select.Option value="30d">monthly</Select.Option>
                 </Select>
               </Form.Item>

ui/litellm-dashboard/src/components/budgets/budget_panel.tsx

@@ -6,6 +6,7 @@
 import React, { useState, useEffect } from "react";
 import BudgetSettings from "./budget_settings";
 import BudgetModal from "./budget_modal";
+import EditBudgetModal from "./edit_budget_modal";
 import {
   Table,
   TableBody,
@@ -43,7 +44,7 @@ interface BudgetSettingsPageProps {
   accessToken: string | null;
 }
 
-interface budgetItem {
+export interface budgetItem {
   budget_id: string;
   max_budget: string | null;
   rpm_limit: number | null;
@@ -52,6 +53,8 @@ interface budgetItem {
 
 const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
   const [isModalVisible, setIsModalVisible] = useState(false);
+  const [isEditModalVisible, setIsEditModalVisible] = useState(false);
+  const [selectedBudget, setSelectedBudget] = useState<budgetItem | null>(null);
   const [budgetList, setBudgetList] = useState<budgetItem[]>([]);
   useEffect(() => {
     if (!accessToken) {
@@ -62,6 +65,15 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
     });
   }, [accessToken]);
 
+
+  const handleEditCall = async (budget_id: string, index: number) => {
+    if (accessToken == null) {
+      return;
+    }
+    setSelectedBudget(budgetList[index])
+    setIsEditModalVisible(true)
+  };
+  
   const handleDeleteCall = async (budget_id: string, index: number) => {
     if (accessToken == null) {
       return;
@@ -94,6 +106,15 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
         setIsModalVisible={setIsModalVisible}
         setBudgetList={setBudgetList}
       />
+      {
+        selectedBudget && <EditBudgetModal
+          accessToken={accessToken}
+          isModalVisible={isEditModalVisible}
+          setIsModalVisible={setIsEditModalVisible}
+          setBudgetList={setBudgetList}
+          existingBudget={selectedBudget}
+        />
+      }
       <Card>
         <Text>Create a budget to assign to customers.</Text>
         <Table>
@@ -119,6 +140,11 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
                 <TableCell>
                   {value.rpm_limit ? value.rpm_limit : "n/a"}
                 </TableCell>
+                <Icon
+                  icon={PencilAltIcon}
+                  size="sm"
+                  onClick={() => handleEditCall(value.budget_id, index)}
+                />
                 <Icon
                   icon={TrashIcon}
                   size="sm"

ui/litellm-dashboard/src/components/budgets/edit_budget_modal.tsx

@@ -0,0 +1,145 @@
+import React from "react";
+import {
+  Button,
+  TextInput,
+  Grid,
+  Col,
+  Accordion,
+  AccordionHeader,
+  AccordionBody,
+} from "@tremor/react";
+import {
+  Button as Button2,
+  Modal,
+  Form,
+  Input,
+  InputNumber,
+  Select,
+  message,
+} from "antd";
+import { budgetCreateCall } from "../networking";
+import { budgetItem } from "./budget_panel";
+
+interface BudgetModalProps {
+  isModalVisible: boolean;
+  accessToken: string | null;
+  setIsModalVisible: React.Dispatch<React.SetStateAction<boolean>>;
+  setBudgetList: React.Dispatch<React.SetStateAction<any[]>>;
+  existingBudget: budgetItem
+}
+const EditBudgetModal: React.FC<BudgetModalProps> = ({
+  isModalVisible,
+  accessToken,
+  setIsModalVisible,
+  setBudgetList,
+  existingBudget
+}) => {
+  const [form] = Form.useForm();
+  const handleOk = () => {
+    setIsModalVisible(false);
+    form.resetFields();
+  };
+
+  const handleCancel = () => {
+    setIsModalVisible(false);
+    form.resetFields();
+  };
+
+  const handleCreate = async (formValues: Record<string, any>) => {
+    if (accessToken == null || accessToken == undefined) {
+      return;
+    }
+    try {
+      message.info("Making API Call");
+      // setIsModalVisible(true);
+      const response = await budgetCreateCall(accessToken, formValues);
+      console.log("key create Response:", response);
+      setBudgetList((prevData) =>
+        prevData ? [...prevData, response] : [response]
+      ); // Check if prevData is null
+      message.success("API Key Created");
+      form.resetFields();
+    } catch (error) {
+      console.error("Error creating the key:", error);
+      message.error(`Error creating the key: ${error}`, 20);
+    }
+  };
+
+  return (
+    <Modal
+      title="Edit Budget"
+      visible={isModalVisible}
+      width={800}
+      footer={null}
+      onOk={handleOk}
+      onCancel={handleCancel}
+    >
+      <Form
+        form={form}
+        onFinish={handleCreate}
+        labelCol={{ span: 8 }}
+        wrapperCol={{ span: 16 }}
+        labelAlign="left"
+        initialValues={existingBudget}
+      >
+        <>
+          <Form.Item
+            label="Budget ID"
+            name="budget_id"
+            rules={[
+              {
+                required: true,
+                message: "Please input a human-friendly name for the budget",
+              },
+            ]}
+            help="A human-friendly name for the budget"
+          >
+            <TextInput placeholder="" />
+          </Form.Item>
+          <Form.Item
+            label="Max Tokens per minute"
+            name="tpm_limit"
+            help="Default is model limit."
+          >
+            <InputNumber step={1} precision={2} width={200} />
+          </Form.Item>
+          <Form.Item
+            label="Max Requests per minute"
+            name="rpm_limit"
+            help="Default is model limit."
+          >
+            <InputNumber step={1} precision={2} width={200} />
+          </Form.Item>
+
+          <Accordion className="mt-20 mb-8">
+            <AccordionHeader>
+              <b>Optional Settings</b>
+            </AccordionHeader>
+            <AccordionBody>
+              <Form.Item label="Max Budget (USD)" name="max_budget">
+                <InputNumber step={0.01} precision={2} width={200} />
+              </Form.Item>
+              <Form.Item
+                className="mt-8"
+                label="Reset Budget"
+                name="budget_duration"
+              >
+                <Select defaultValue={null} placeholder="n/a">
+                  <Select.Option value="24h">daily</Select.Option>
+                  <Select.Option value="7d">weekly</Select.Option>
+                  <Select.Option value="30d">monthly</Select.Option>
+                </Select>
+              </Form.Item>
+            </AccordionBody>
+          </Accordion>
+        </>
+
+        <div style={{ textAlign: "right", marginTop: "10px" }}>
+          <Button2 htmlType="submit">Edit Budget</Button2>
+        </div>
+      </Form>
+    </Modal>
+  );
+};
+
+export default EditBudgetModal;

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

@@ -261,6 +261,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({
                 >
                   <Select defaultValue={null} placeholder="n/a">
                     <Select.Option value="24h">daily</Select.Option>
+                    <Select.Option value="7d">weekly</Select.Option>
                     <Select.Option value="30d">monthly</Select.Option>
                   </Select>
                 </Form.Item>

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

@@ -151,6 +151,17 @@ const Team: React.FC<TeamProps> = ({
             <Form.Item label="Max Budget (USD)" name="max_budget">
               <InputNumber step={0.01} precision={2} width={200} />
             </Form.Item>
+            <Form.Item
+              className="mt-8"
+              label="Reset Budget"
+              name="budget_duration"
+            >
+              <Select2 defaultValue={null} placeholder="n/a">
+                <Select2.Option value="24h">daily</Select2.Option>
+                <Select2.Option value="7d">weekly</Select2.Option>
+                <Select2.Option value="30d">monthly</Select2.Option>
+              </Select2>
+            </Form.Item>
             <Form.Item label="Tokens per minute Limit (TPM)" name="tpm_limit">
               <InputNumber step={1} width={400} />
             </Form.Item>
@@ -635,6 +646,17 @@ const Team: React.FC<TeamProps> = ({
                 <Form.Item label="Max Budget (USD)" name="max_budget">
                   <InputNumber step={0.01} precision={2} width={200} />
                 </Form.Item>
+                <Form.Item
+                  className="mt-8"
+                  label="Reset Budget"
+                  name="budget_duration"
+                >
+                  <Select2 defaultValue={null} placeholder="n/a">
+                    <Select2.Option value="24h">daily</Select2.Option>
+                    <Select2.Option value="7d">weekly</Select2.Option>
+                    <Select2.Option value="30d">monthly</Select2.Option>
+                  </Select2>
+                </Form.Item>
                 <Form.Item
                   label="Tokens per minute Limit (TPM)"
                   name="tpm_limit"