Merge pull request #18576 from BerriAI/litellm_ui_budget_2

[Fix] UI - Change Budget page to Have Tabs

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

@@ -1,5 +1,6 @@
 import * as networking from "../networking";
 import { fireEvent, render, waitFor, screen } from "@testing-library/react";
+import { act } from "@testing-library/react";
 import { afterEach, describe, expect, it, vi } from "vitest";
 import BudgetPanel from "./budget_panel";
 
@@ -24,11 +25,11 @@ describe("Budget Panel", () => {
       },
     ]);
 
-    const { getByText } = render(<BudgetPanel accessToken="token-123" />);
+    render(<BudgetPanel accessToken="token-123" />);
 
     await waitFor(() => {
-      expect(getByText("Create a budget to assign to customers.")).toBeInTheDocument();
-      expect(getByText("budget-1")).toBeInTheDocument();
+      expect(screen.getByText("Create a budget to assign to customers.")).toBeInTheDocument();
+      expect(screen.getByText("budget-1")).toBeInTheDocument();
     });
   });
 
@@ -43,23 +44,102 @@ describe("Budget Panel", () => {
       },
     ]);
 
-    const { getByText, container } = render(<BudgetPanel accessToken="token-123" />);
+    render(<BudgetPanel accessToken="token-123" />);
 
     await waitFor(() => {
-      expect(getByText("budget-to-delete")).toBeInTheDocument();
+      expect(screen.getByText("budget-to-delete")).toBeInTheDocument();
     });
 
-    // Find the first table row in tbody and click the second icon (trash/delete)
-    const bodyRows = container.querySelectorAll("tbody tr");
-    expect(bodyRows.length).toBeGreaterThan(0);
-    const firstRow = bodyRows[0];
-    const rowClickableIcons = firstRow.querySelectorAll(".cursor-pointer");
-    expect(rowClickableIcons.length).toBeGreaterThan(1);
+    const deleteButton = screen.getByTestId("delete-budget-button");
 
-    fireEvent.click(rowClickableIcons[1]);
+    act(() => {
+      fireEvent.click(deleteButton);
+    });
 
     await waitFor(() => {
       expect(screen.getByText("Delete Budget?")).toBeInTheDocument();
     });
   });
+
+  it("should successfully delete a budget", async () => {
+    vi.mocked(networking.getBudgetList).mockResolvedValue([
+      {
+        budget_id: "budget-to-delete",
+        max_budget: "200",
+        rpm_limit: 20,
+        tpm_limit: 2000,
+        updated_at: "2024-01-02T00:00:00Z",
+      },
+    ]);
+    vi.mocked(networking.budgetDeleteCall).mockResolvedValue(undefined);
+
+    render(<BudgetPanel accessToken="token-123" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("budget-to-delete")).toBeInTheDocument();
+    });
+
+    // Open delete modal
+    const deleteButton = screen.getByTestId("delete-budget-button");
+    act(() => {
+      fireEvent.click(deleteButton);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Delete Budget?")).toBeInTheDocument();
+    });
+
+    // Confirm delete
+    const confirmButton = screen.getByRole("button", { name: /delete/i });
+    act(() => {
+      fireEvent.click(confirmButton);
+    });
+
+    await waitFor(() => {
+      expect(networking.budgetDeleteCall).toHaveBeenCalledWith("token-123", "budget-to-delete");
+      expect(networking.getBudgetList).toHaveBeenCalledTimes(2); // Initial load + refresh after delete
+    });
+  });
+
+  it("should handle delete error", async () => {
+    vi.mocked(networking.getBudgetList).mockResolvedValue([
+      {
+        budget_id: "budget-to-delete",
+        max_budget: "200",
+        rpm_limit: 20,
+        tpm_limit: 2000,
+        updated_at: "2024-01-02T00:00:00Z",
+      },
+    ]);
+    vi.mocked(networking.budgetDeleteCall).mockRejectedValue(new Error("Delete failed"));
+
+    render(<BudgetPanel accessToken="token-123" />);
+
+    await waitFor(() => {
+      expect(screen.getByText("budget-to-delete")).toBeInTheDocument();
+    });
+
+    // Open delete modal
+    const deleteButton = screen.getByTestId("delete-budget-button");
+    act(() => {
+      fireEvent.click(deleteButton);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Delete Budget?")).toBeInTheDocument();
+    });
+
+    // Confirm delete
+    const confirmButton = screen.getByRole("button", { name: /delete/i });
+    act(() => {
+      fireEvent.click(confirmButton);
+    });
+
+    await waitFor(() => {
+      expect(networking.budgetDeleteCall).toHaveBeenCalledWith("token-123", "budget-to-delete");
+    });
+
+    // Modal should still be open (error handling)
+    expect(screen.getByText("Delete Budget?")).toBeInTheDocument();
+  });
 });

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

@@ -27,6 +27,7 @@ import NotificationsManager from "../molecules/notifications_manager";
 import { budgetDeleteCall, getBudgetList } from "../networking";
 import BudgetModal from "./budget_modal";
 import EditBudgetModal from "./edit_budget_modal";
+import { CREATE_END_USER_CURL_COMMAND, CHAT_COMPLETIONS_CURL_COMMAND, OPENAI_SDK_PYTHON_CODE } from "./constants";
 
 interface BudgetSettingsPageProps {
   accessToken: string | null;
@@ -110,139 +111,111 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ accessToken }) => {
       <Button size="sm" variant="primary" className="mb-2" onClick={() => setIsCreateModelVisible(true)}>
         + Create Budget
       </Button>
-      <BudgetModal
-        accessToken={accessToken}
-        isModalVisible={isCreateModelVisible}
-        setIsModalVisible={setIsCreateModelVisible}
-        setBudgetList={setBudgetList}
-      />
-      {selectedBudget && (
-        <EditBudgetModal
-          accessToken={accessToken}
-          isModalVisible={isEditModalVisible}
-          setIsModalVisible={setIsEditModalVisible}
-          setBudgetList={setBudgetList}
-          existingBudget={selectedBudget}
-          handleUpdateCall={handleUpdateCall}
-        />
-      )}
-      <Card>
-        <Text>Create a budget to assign to customers.</Text>
-        <Table>
-          <TableHead>
-            <TableRow>
-              <TableHeaderCell>Budget ID</TableHeaderCell>
-              <TableHeaderCell>Max Budget</TableHeaderCell>
-              <TableHeaderCell>TPM</TableHeaderCell>
-              <TableHeaderCell>RPM</TableHeaderCell>
-            </TableRow>
-          </TableHead>
+      <TabGroup>
+        <TabList>
+          <Tab>Budgets</Tab>
+          <Tab>Examples</Tab>
+        </TabList>
+        <TabPanels>
+          <TabPanel>
+            <div className="mt-6">
+              <BudgetModal
+                accessToken={accessToken}
+                isModalVisible={isCreateModelVisible}
+                setIsModalVisible={setIsCreateModelVisible}
+                setBudgetList={setBudgetList}
+              />
+              {selectedBudget && (
+                <EditBudgetModal
+                  accessToken={accessToken}
+                  isModalVisible={isEditModalVisible}
+                  setIsModalVisible={setIsEditModalVisible}
+                  setBudgetList={setBudgetList}
+                  existingBudget={selectedBudget}
+                  handleUpdateCall={handleUpdateCall}
+                />
+              )}
+              <Card>
+                <Text>Create a budget to assign to customers.</Text>
+                <Table>
+                  <TableHead>
+                    <TableRow>
+                      <TableHeaderCell>Budget ID</TableHeaderCell>
+                      <TableHeaderCell>Max Budget</TableHeaderCell>
+                      <TableHeaderCell>TPM</TableHeaderCell>
+                      <TableHeaderCell>RPM</TableHeaderCell>
+                    </TableRow>
+                  </TableHead>
 
-          <TableBody>
-            {budgetList
-              .slice() // Creates a shallow copy to avoid mutating the original array
-              .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) // Sort by updated_at in descending order
-              .map((value: budgetItem, index: number) => (
-                <TableRow key={index}>
-                  <TableCell>{value.budget_id}</TableCell>
-                  <TableCell>{value.max_budget ? value.max_budget : "n/a"}</TableCell>
-                  <TableCell>{value.tpm_limit ? value.tpm_limit : "n/a"}</TableCell>
-                  <TableCell>{value.rpm_limit ? value.rpm_limit : "n/a"}</TableCell>
-                  <TableIconActionButton
-                    variant="Edit"
-                    tooltipText="Edit budget"
-                    onClick={() => handleEditCall(value)}
-                  />
-                  <TableIconActionButton
-                    variant="Delete"
-                    tooltipText="Delete budget"
-                    onClick={() => handleDeleteClick(value)}
-                  />
-                </TableRow>
-              ))}
-          </TableBody>
-        </Table>
-      </Card>
-      <DeleteResourceModal
-        isOpen={isDeleteModalVisible}
-        title="Delete Budget?"
-        message="Are you sure you want to delete this budget? This action cannot be undone."
-        resourceInformationTitle="Budget Information"
-        resourceInformation={[
-          { label: "Budget ID", value: selectedBudget?.budget_id, code: true },
-          { label: "Max Budget", value: selectedBudget?.max_budget },
-          { label: "TPM", value: selectedBudget?.tpm_limit },
-          { label: "RPM", value: selectedBudget?.rpm_limit },
-        ]}
-        onCancel={handleDeleteCancel}
-        onOk={handleDeleteConfirm}
-        confirmLoading={isDeleting}
-      />
-      <div className="mt-5">
-        <Text className="text-base">How to use budget id</Text>
-        <TabGroup>
-          <TabList>
-            <Tab>Assign Budget to Customer</Tab>
-            <Tab>Test it (Curl)</Tab>
-
-            <Tab>Test it (OpenAI SDK)</Tab>
-          </TabList>
-          <TabPanels>
-            <TabPanel>
-              <SyntaxHighlighter language="bash">
-                {`
-curl -X POST --location '<your_proxy_base_url>/end_user/new' \
-
--H 'Authorization: Bearer <your-master-key>' \
-
--H 'Content-Type: application/json' \
-
--d '{"user_id": "my-customer-id', "budget_id": "<BUDGET_ID>"}' # 👈 KEY CHANGE
-
-            `}
-              </SyntaxHighlighter>
-            </TabPanel>
-            <TabPanel>
-              <SyntaxHighlighter language="bash">
-                {`
-curl -X POST --location '<your_proxy_base_url>/chat/completions' \
-
--H 'Authorization: Bearer <your-master-key>' \
-
--H 'Content-Type: application/json' \
-
--d '{
-  "model": "gpt-3.5-turbo', 
-  "messages":[{"role": "user", "content": "Hey, how's it going?"}],
-  "user": "my-customer-id"
-}' # 👈 KEY CHANGE
-
-            `}
-              </SyntaxHighlighter>
-            </TabPanel>
-            <TabPanel>
-              <SyntaxHighlighter language="python">
-                {`from openai import OpenAI
-client = OpenAI(
-  base_url="<your_proxy_base_url>",
-  api_key="<your_proxy_key>"
-)
-
-completion = client.chat.completions.create(
-  model="gpt-3.5-turbo",
-  messages=[
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": "Hello!"}
-  ],
-  user="my-customer-id"
-)
-
-print(completion.choices[0].message)`}
-              </SyntaxHighlighter>
-            </TabPanel>
-          </TabPanels>
-        </TabGroup>
-      </div>
+                  <TableBody>
+                    {budgetList
+                      .slice() // Creates a shallow copy to avoid mutating the original array
+                      .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) // Sort by updated_at in descending order
+                      .map((value: budgetItem, index: number) => (
+                        <TableRow key={index}>
+                          <TableCell>{value.budget_id}</TableCell>
+                          <TableCell>{value.max_budget ? value.max_budget : "n/a"}</TableCell>
+                          <TableCell>{value.tpm_limit ? value.tpm_limit : "n/a"}</TableCell>
+                          <TableCell>{value.rpm_limit ? value.rpm_limit : "n/a"}</TableCell>
+                          <TableIconActionButton
+                            variant="Edit"
+                            tooltipText="Edit budget"
+                            onClick={() => handleEditCall(value)}
+                            dataTestId="edit-budget-button"
+                          />
+                          <TableIconActionButton
+                            variant="Delete"
+                            tooltipText="Delete budget"
+                            onClick={() => handleDeleteClick(value)}
+                            dataTestId="delete-budget-button"
+                          />
+                        </TableRow>
+                      ))}
+                  </TableBody>
+                </Table>
+              </Card>
+              <DeleteResourceModal
+                isOpen={isDeleteModalVisible}
+                title="Delete Budget?"
+                message="Are you sure you want to delete this budget? This action cannot be undone."
+                resourceInformationTitle="Budget Information"
+                resourceInformation={[
+                  { label: "Budget ID", value: selectedBudget?.budget_id, code: true },
+                  { label: "Max Budget", value: selectedBudget?.max_budget },
+                  { label: "TPM", value: selectedBudget?.tpm_limit },
+                  { label: "RPM", value: selectedBudget?.rpm_limit },
+                ]}
+                onCancel={handleDeleteCancel}
+                onOk={handleDeleteConfirm}
+                confirmLoading={isDeleting}
+              />
+            </div>
+          </TabPanel>
+          <TabPanel>
+            <div className="mt-6">
+              <Text className="text-base">How to use budget id</Text>
+              <TabGroup>
+                <TabList>
+                  <Tab>Assign Budget to Customer</Tab>
+                  <Tab>Test it (Curl)</Tab>
+                  <Tab>Test it (OpenAI SDK)</Tab>
+                </TabList>
+                <TabPanels>
+                  <TabPanel>
+                    <SyntaxHighlighter language="bash">{CREATE_END_USER_CURL_COMMAND}</SyntaxHighlighter>
+                  </TabPanel>
+                  <TabPanel>
+                    <SyntaxHighlighter language="bash">{CHAT_COMPLETIONS_CURL_COMMAND}</SyntaxHighlighter>
+                  </TabPanel>
+                  <TabPanel>
+                    <SyntaxHighlighter language="python">{OPENAI_SDK_PYTHON_CODE}</SyntaxHighlighter>
+                  </TabPanel>
+                </TabPanels>
+              </TabGroup>
+            </div>
+          </TabPanel>
+        </TabPanels>
+      </TabGroup>
     </div>
   );
 };

ui/litellm-dashboard/src/components/budgets/constants.ts

@@ -0,0 +1,42 @@
+export const CREATE_END_USER_CURL_COMMAND = `
+curl -X POST --location '<your_proxy_base_url>/end_user/new' \\
+
+-H 'Authorization: Bearer <your-master-key>' \\
+
+-H 'Content-Type: application/json' \\
+
+-d '{"user_id": "my-customer-id', "budget_id": "<BUDGET_ID>"}' # 👈 KEY CHANGE
+
+`;
+
+export const CHAT_COMPLETIONS_CURL_COMMAND = `
+curl -X POST --location '<your_proxy_base_url>/chat/completions' \\
+
+-H 'Authorization: Bearer <your-master-key>' \\
+
+-H 'Content-Type: application/json' \\
+
+-d '{
+  "model": "gpt-3.5-turbo',
+  "messages":[{"role": "user", "content": "Hey, how's it going?"}],
+  "user": "my-customer-id"
+}' # 👈 KEY CHANGE
+
+`;
+
+export const OPENAI_SDK_PYTHON_CODE = `from openai import OpenAI
+client = OpenAI(
+  base_url="<your_proxy_base_url>",
+  api_key="<your_proxy_key>"
+)
+
+completion = client.chat.completions.create(
+  model="gpt-3.5-turbo",
+  messages=[
+    {"role": "system", "content": "You are a helpful assistant."},
+    {"role": "user", "content": "Hello!"}
+  ],
+  user="my-customer-id"
+)
+
+print(completion.choices[0].message)`;