Merge pull request #17465 from BerriAI/litellm_ui_model_page_perf_2

[Infra] UI - Change Credentials to use React Query

ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.ts

@@ -0,0 +1,13 @@
+import { credentialListCall, CredentialsResponse } from "@/components/networking";
+import { useQuery } from "@tanstack/react-query";
+import { createQueryKeys } from "../common/queryKeysFactory";
+
+const credentialsKeys = createQueryKeys("credentials");
+
+export const useCredentials = (accessToken: string | null) => {
+  return useQuery<CredentialsResponse>({
+    queryKey: credentialsKeys.list({}),
+    queryFn: async () => await credentialListCall(accessToken!),
+    enabled: Boolean(accessToken),
+  });
+};

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx

@@ -1,47 +1,45 @@
-import React, { useState, useEffect, useRef } from "react";
-import { Text, Grid, Col } from "@tremor/react";
 import { useQueryClient } from "@tanstack/react-query";
-import { CredentialItem, credentialListCall, CredentialsResponse } from "@/components/networking";
+import { Col, Grid, Text } from "@tremor/react";
+import React, { useEffect, useRef, useState } from "react";
 
 import { handleAddModelSubmit } from "@/components/add_model/handle_add_model_submit";
 
+import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials";
+import { useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels";
+import { Team } from "@/components/key_team_helpers/key_list";
 import CredentialsPanel from "@/components/model_add/credentials";
-import { getDisplayModelName } from "@/components/view_model/model_name_display";
-import { TabPanel, TabPanels, TabGroup, TabList, Tab, Icon } from "@tremor/react";
-import { DateRangePickerValue } from "@tremor/react";
 import {
-  modelCostMap,
-  modelMetricsCall,
-  streamingModelMetricsCall,
-  modelExceptionsCall,
-  modelMetricsSlowResponsesCall,
-  getCallbacksCall,
-  setCallbacksCall,
-  modelSettingsCall,
   adminGlobalActivityExceptions,
   adminGlobalActivityExceptionsPerDeployment,
   allEndUsersCall,
+  getCallbacksCall,
+  modelCostMap,
+  modelExceptionsCall,
+  modelMetricsCall,
+  modelMetricsSlowResponsesCall,
+  modelSettingsCall,
+  setCallbacksCall,
+  streamingModelMetricsCall,
 } from "@/components/networking";
-import { useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels";
-import { Form } from "antd";
-import { Typography } from "antd";
-import { RefreshIcon } from "@heroicons/react/outline";
-import type { UploadProps } from "antd";
-import { Team } from "@/components/key_team_helpers/key_list";
-import TeamInfoView from "../../../components/team/team_info";
 import { Providers, getPlaceholder, getProviderModels } from "@/components/provider_info_helpers";
-import ModelInfoView from "../../../components/model_info_view";
+import { getDisplayModelName } from "@/components/view_model/model_name_display";
+import { RefreshIcon } from "@heroicons/react/outline";
+import { DateRangePickerValue, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react";
+import type { UploadProps } from "antd";
+import { Form, Typography } from "antd";
 import AddModelTab from "../../../components/add_model/add_model_tab";
+import ModelInfoView from "../../../components/model_info_view";
+import TeamInfoView from "../../../components/team/team_info";
 
-import HealthCheckComponent from "../../../components/model_dashboard/HealthCheckComponent";
-import PassThroughSettings from "../../../components/pass_through_settings";
-import ModelGroupAliasSettings from "../../../components/model_group_alias_settings";
-import { all_admin_roles } from "@/utils/roles";
-import NotificationsManager from "../../../components/molecules/notifications_manager";
 import AllModelsTab from "@/app/(dashboard)/models-and-endpoints/components/AllModelsTab";
-import PriceDataManagementTab from "@/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab";
-import ModelRetrySettingsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab";
 import ModelAnalyticsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab";
+import ModelRetrySettingsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab";
+import PriceDataManagementTab from "@/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab";
+import { all_admin_roles } from "@/utils/roles";
+import HealthCheckComponent from "../../../components/model_dashboard/HealthCheckComponent";
+import ModelGroupAliasSettings from "../../../components/model_group_alias_settings";
+import NotificationsManager from "../../../components/molecules/notifications_manager";
+import PassThroughSettings from "../../../components/pass_through_settings";
 
 interface ModelDashboardProps {
   accessToken: string | null;
@@ -134,8 +132,6 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
 
   const [allEndUsers, setAllEndUsers] = useState<any[]>([]);
 
-  const [credentialsList, setCredentialsList] = useState<CredentialItem[]>([]);
-
   // Model Group Alias state
   const [modelGroupAlias, setModelGroupAlias] = useState<{ [key: string]: string }>({});
 
@@ -160,21 +156,14 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
     isLoading: isLoadingModels,
     refetch: refetchModels,
   } = useModelsInfo(accessToken, userID, userRole);
+  const { data: credentialsResponse } = useCredentials(accessToken);
+  const credentialsList = credentialsResponse?.credentials || [];
 
   const setProviderModelsFn = (provider: Providers) => {
     const _providerModels = getProviderModels(provider, modelMap);
     setProviderModels(_providerModels);
   };
 
-  const fetchCredentials = async (accessToken: string) => {
-    try {
-      const response: CredentialsResponse = await credentialListCall(accessToken);
-      setCredentialsList(response.credentials);
-    } catch (error) {
-      console.error("Error fetching credentials:", error);
-    }
-  };
-
   useEffect(() => {
     const handleClickOutside = (event: MouseEvent) => {
       if (dropdownRef.current && !dropdownRef.current.contains(event.target as Node)) {
@@ -686,12 +675,7 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
                   />
                 </TabPanel>
                 <TabPanel>
-                  <CredentialsPanel
-                    accessToken={accessToken}
-                    uploadProps={uploadProps}
-                    credentialList={credentialsList}
-                    fetchCredentials={fetchCredentials}
-                  />
+                  <CredentialsPanel uploadProps={uploadProps} />
                 </TabPanel>
                 <TabPanel>
                   <PassThroughSettings

ui/litellm-dashboard/src/components/model_add/credentials.test.tsx

@@ -1,46 +1,51 @@
 import { CredentialItem } from "@/components/networking";
-import { render, screen, waitFor } from "@testing-library/react";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { act, fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { UploadProps } from "antd/es/upload";
 import { describe, expect, it, vi } from "vitest";
 import CredentialsPanel from "./credentials";
 
 const DEFAULT_UPLOAD_PROPS = {} as UploadProps;
 
+const mockUseAuthorized = vi.fn();
+const mockUseCredentials = vi.fn();
+
+vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
+  default: () => mockUseAuthorized(),
+}));
+
+vi.mock("@/app/(dashboard)/hooks/credentials/useCredentials", () => ({
+  useCredentials: () => mockUseCredentials(),
+}));
+
+const createQueryClient = () =>
+  new QueryClient({
+    defaultOptions: {
+      queries: {
+        retry: false,
+        gcTime: 0,
+      },
+    },
+  });
+
 describe("CredentialsPanel", () => {
   it("should render", () => {
-    const fetchCredentials = vi.fn(() => Promise.resolve());
+    mockUseAuthorized.mockReturnValue({ accessToken: "test-token" });
+    mockUseCredentials.mockReturnValue({
+      data: { credentials: [] },
+      refetch: vi.fn(),
+    });
 
     render(
-      <CredentialsPanel
-        accessToken="test-token"
-        uploadProps={DEFAULT_UPLOAD_PROPS}
-        credentialList={[]}
-        fetchCredentials={fetchCredentials}
-      />,
+      <QueryClientProvider client={createQueryClient()}>
+        <CredentialsPanel uploadProps={DEFAULT_UPLOAD_PROPS} />
+      </QueryClientProvider>,
     );
 
     expect(screen.getByRole("button", { name: /add credential/i })).toBeInTheDocument();
   });
 
-  it("should call fetchCredentials when accessToken exists", async () => {
-    const fetchCredentials = vi.fn(() => Promise.resolve());
-
-    render(
-      <CredentialsPanel
-        accessToken="test-token"
-        uploadProps={DEFAULT_UPLOAD_PROPS}
-        credentialList={[]}
-        fetchCredentials={fetchCredentials}
-      />,
-    );
-
-    await waitFor(() => {
-      expect(fetchCredentials).toHaveBeenCalledWith("test-token");
-    });
-  });
-
   it("should display provided credentials", () => {
-    const fetchCredentials = vi.fn(() => Promise.resolve());
     const credentials: CredentialItem[] = [
       {
         credential_name: "openai-key",
@@ -49,30 +54,58 @@ describe("CredentialsPanel", () => {
       },
     ];
 
+    mockUseAuthorized.mockReturnValue({ accessToken: "test-token" });
+    mockUseCredentials.mockReturnValue({
+      data: { credentials },
+      refetch: vi.fn(),
+    });
+
     render(
-      <CredentialsPanel
-        accessToken="another-token"
-        uploadProps={DEFAULT_UPLOAD_PROPS}
-        credentialList={credentials}
-        fetchCredentials={fetchCredentials}
-      />,
+      <QueryClientProvider client={createQueryClient()}>
+        <CredentialsPanel uploadProps={DEFAULT_UPLOAD_PROPS} />
+      </QueryClientProvider>,
     );
 
     expect(screen.getByText("openai-key")).toBeInTheDocument();
   });
 
   it("should display empty state when no credentials are provided", () => {
-    const fetchCredentials = vi.fn(() => Promise.resolve());
+    mockUseAuthorized.mockReturnValue({ accessToken: "test-token" });
+    mockUseCredentials.mockReturnValue({
+      data: { credentials: [] },
+      refetch: vi.fn(),
+    });
 
     render(
-      <CredentialsPanel
-        accessToken="test-token"
-        uploadProps={DEFAULT_UPLOAD_PROPS}
-        credentialList={[]}
-        fetchCredentials={fetchCredentials}
-      />,
+      <QueryClientProvider client={createQueryClient()}>
+        <CredentialsPanel uploadProps={DEFAULT_UPLOAD_PROPS} />
+      </QueryClientProvider>,
     );
 
     expect(screen.getByText("No credentials configured")).toBeInTheDocument();
   });
+
+  it("should open add modal when add button is clicked", async () => {
+    mockUseAuthorized.mockReturnValue({ accessToken: "test-token" });
+    mockUseCredentials.mockReturnValue({
+      data: { credentials: [] },
+      refetch: vi.fn(),
+    });
+
+    render(
+      <QueryClientProvider client={createQueryClient()}>
+        <CredentialsPanel uploadProps={DEFAULT_UPLOAD_PROPS} />
+      </QueryClientProvider>,
+    );
+
+    const addButton = screen.getByRole("button", { name: /add credential/i });
+
+    act(() => {
+      fireEvent.click(addButton);
+    });
+
+    await waitFor(() => {
+      expect(screen.getByText("Add New Credential")).toBeInTheDocument();
+    });
+  });
 });

ui/litellm-dashboard/src/components/model_add/credentials.tsx

@@ -19,24 +19,22 @@ import {
 } from "@tremor/react";
 import { Form } from "antd";
 import { UploadProps } from "antd/es/upload";
-import { useEffect, useState } from "react";
+import { useState } from "react";
 import DeleteResourceModal from "../common_components/DeleteResourceModal";
 import NotificationsManager from "../molecules/notifications_manager";
 import AddCredentialsTab from "./AddCredentialModal";
 import EditCredentialsModal from "./EditCredentialModal";
+import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials";
+import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 interface CredentialsPanelProps {
-  accessToken: string | null;
   uploadProps: UploadProps;
-  credentialList: CredentialItem[];
-  fetchCredentials: (accessToken: string) => Promise<void>;
 }
 
-const CredentialsPanel: React.FC<CredentialsPanelProps> = ({
-  accessToken,
-  uploadProps,
-  credentialList,
-  fetchCredentials,
-}) => {
+const CredentialsPanel: React.FC<CredentialsPanelProps> = ({ uploadProps }) => {
+  const { accessToken } = useAuthorized();
+  const { data: credentialsResponse, refetch: refetchCredentials } = useCredentials(accessToken);
+  const credentialList = credentialsResponse?.credentials || [];
+
   const [isAddModalOpen, setIsAddModalOpen] = useState(false);
   const [isUpdateModalOpen, setIsUpdateModalOpen] = useState(false);
   const [selectedCredential, setSelectedCredential] = useState<CredentialItem | null>(null);
@@ -66,7 +64,7 @@ const CredentialsPanel: React.FC<CredentialsPanelProps> = ({
     await credentialUpdateCall(accessToken, values.credential_name, newCredential);
     NotificationsManager.success("Credential updated successfully");
     setIsUpdateModalOpen(false);
-    await fetchCredentials(accessToken);
+    await refetchCredentials();
   };
 
   const handleAddCredential = async (values: any) => {
@@ -90,16 +88,9 @@ const CredentialsPanel: React.FC<CredentialsPanelProps> = ({
     await credentialCreateCall(accessToken, newCredential);
     NotificationsManager.success("Credential added successfully");
     setIsAddModalOpen(false);
-    await fetchCredentials(accessToken);
+    await refetchCredentials();
   };
 
-  useEffect(() => {
-    if (!accessToken) {
-      return;
-    }
-    fetchCredentials(accessToken);
-  }, [accessToken]);
-
   const renderProviderBadge = (provider: string) => {
     const providerColors: Record<string, string> = {
       openai: "blue",
@@ -124,7 +115,7 @@ const CredentialsPanel: React.FC<CredentialsPanelProps> = ({
     try {
       await credentialDeleteCall(accessToken, credentialToDelete.credential_name);
       NotificationsManager.success("Credential deleted successfully");
-      await fetchCredentials(accessToken);
+      await refetchCredentials();
     } catch (error) {
       NotificationsManager.error("Failed to delete credential");
     } finally {

ui/litellm-dashboard/src/components/templates/model_dashboard.tsx

@@ -1351,12 +1351,7 @@ const OldModelDashboard: React.FC<ModelDashboardProps> = ({
                   />
                 </TabPanel>
                 <TabPanel>
-                  <CredentialsPanel
-                    accessToken={accessToken}
-                    uploadProps={uploadProps}
-                    credentialList={credentialsList}
-                    fetchCredentials={fetchCredentials}
-                  />
+                  <CredentialsPanel uploadProps={uploadProps} />
                 </TabPanel>
                 <TabPanel>
                   <PassThroughSettings