refactor pt2

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

@@ -11,6 +11,7 @@ import CredentialsPanel from "@/components/model_add/credentials";
 import { getCallbacksCall, setCallbacksCall } from "@/components/networking";
 import { Providers, getPlaceholder, getProviderModels } from "@/components/provider_info_helpers";
 import { getDisplayModelName } from "@/components/view_model/model_name_display";
+import { transformModelData } from "./utils/modelDataTransformer";
 import { all_admin_roles, internalUserRoles, isProxyAdminRole, isUserTeamAdminForAnyTeam } from "@/utils/roles";
 import { RefreshIcon } from "@heroicons/react/outline";
 import { useQueryClient } from "@tanstack/react-query";
@@ -44,12 +45,7 @@ interface GlobalRetryPolicyObject {
   [retryPolicyKey: string]: number;
 }
 
-const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
-  modelData = { data: [] },
-  setModelData,
-  premiumUser,
-  teams,
-}) => {
+const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({ premiumUser, teams }) => {
   const { accessToken, token, userRole, userId: userID } = useAuthorized();
   const [addModelForm] = Form.useForm();
   const [lastRefreshed, setLastRefreshed] = useState("");
@@ -68,10 +64,10 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
 
   const queryClient = useQueryClient();
   const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo();
-  const { data: modelCostMapData } = useModelCostMap();
-  const { data: credentialsResponse } = useCredentials();
+  const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap();
+  const { data: credentialsResponse, isLoading: isLoadingCredentials } = useCredentials();
   const credentialsList = credentialsResponse?.credentials || [];
-  const { data: uiSettings } = useUISettings();
+  const { data: uiSettings, isLoading: isLoadingUISettings } = useUISettings();
 
   const availableModelGroups = useMemo(() => {
     if (!modelDataResponse?.data) return [];
@@ -110,75 +106,9 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
   };
 
   const processedModelData = useMemo(() => {
-    return modelDataResponse?.data?.map((model: any) => {
-      return {};
-    });
-  }, [modelDataResponse?.data]);
-
-  // loop through model data and edit each row
-  for (let i = 0; i < modelData.data.length; i++) {
-    let curr_model = modelData.data[i];
-    let litellm_model_name = curr_model?.litellm_params?.model;
-    let custom_llm_provider = curr_model?.litellm_params?.custom_llm_provider;
-    let model_info = curr_model?.model_info;
-
-    let provider = "";
-    let input_cost = "Undefined";
-    let output_cost = "Undefined";
-    let max_tokens = "Undefined";
-    let max_input_tokens = "Undefined";
-    let cleanedLitellmParams = {};
-
-    // Check if litellm_model_name is null or undefined
-    if (litellm_model_name) {
-      // Split litellm_model_name based on "/"
-      let splitModel = litellm_model_name.split("/");
-
-      // Get the first element in the split
-      let firstElement = splitModel[0];
-
-      // If there is only one element, default provider to openai
-      provider = custom_llm_provider;
-      if (!provider) {
-        provider = splitModel.length === 1 ? getProviderFromModel(litellm_model_name) : firstElement;
-      }
-    } else {
-      // litellm_model_name is null or undefined, default provider to openai
-      provider = "-";
-    }
-
-    if (model_info) {
-      input_cost = model_info?.input_cost_per_token;
-      output_cost = model_info?.output_cost_per_token;
-      max_tokens = model_info?.max_tokens;
-      max_input_tokens = model_info?.max_input_tokens;
-    }
-
-    if (curr_model?.litellm_params) {
-      cleanedLitellmParams = Object.fromEntries(
-        Object.entries(curr_model?.litellm_params).filter(([key]) => key !== "model" && key !== "api_base"),
-      );
-    }
-
-    modelData.data[i].provider = provider;
-    modelData.data[i].input_cost = input_cost;
-    modelData.data[i].output_cost = output_cost;
-    modelData.data[i].litellm_model_name = litellm_model_name;
-
-    // Convert Cost in terms of Cost per 1M tokens
-    if (modelData.data[i].input_cost) {
-      modelData.data[i].input_cost = (Number(modelData.data[i].input_cost) * 1000000).toFixed(2);
-    }
-
-    if (modelData.data[i].output_cost) {
-      modelData.data[i].output_cost = (Number(modelData.data[i].output_cost) * 1000000).toFixed(2);
-    }
-
-    modelData.data[i].max_tokens = max_tokens;
-    modelData.data[i].max_input_tokens = max_input_tokens;
-    modelData.data[i].api_base = curr_model?.litellm_params?.api_base;
-    modelData.data[i].cleanedLitellmParams = cleanedLitellmParams;
-  }
+    if (!modelDataResponse?.data) return { data: [] };
+    return transformModelData(modelDataResponse, getProviderFromModel);
+  }, [modelDataResponse?.data, getProviderFromModel]);
 
   const isProxyAdmin = userRole && isProxyAdminRole(userRole);
   const isInternalUser = userRole && internalUserRoles.includes(userRole);
@@ -260,8 +190,6 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
     }
     const fetchData = async () => {
       try {
-        setModelData(modelDataResponse);
-
         const routerSettingsInfo = await getCallbacksCall(accessToken, userID, userRole);
         let router_settings = routerSettingsInfo.router_settings;
 
@@ -284,9 +212,7 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
     }
   }, [accessToken, token, userRole, userID, modelDataResponse]);
 
-  if (!modelData || isLoadingModels) {
-    return <div>Loading...</div>;
-  }
+  const isLoading = isLoadingModels || isLoadingModelCostMap || isLoadingCredentials || isLoadingUISettings;
 
   if (userRole && userRole == "Admin Viewer") {
     const { Title, Paragraph } = Typography;
@@ -348,13 +274,13 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
               )}
             </div>
           </div>
-          {selectedModelId ? (
+          {selectedModelId && !isLoading ? (
             <ModelInfoView
               modelId={selectedModelId}
               onClose={() => {
                 setSelectedModelId(null);
               }}
-              modelData={modelData.data.find((model: any) => model.model_info.id === selectedModelId)}
+              modelData={processedModelData.data.find((model: any) => model.model_info.id === selectedModelId)}
               accessToken={accessToken}
               userID={userID}
               userRole={userRole}
@@ -426,14 +352,14 @@ const ModelsAndEndpointsView: React.FC<ModelDashboardProps> = ({
                     accessToken={accessToken}
                     userRole={userRole}
                     userID={userID}
-                    modelData={modelData}
+                    modelData={processedModelData}
                     premiumUser={premiumUser}
                   />
                 </TabPanel>
                 <TabPanel>
                   <HealthCheckComponent
                     accessToken={accessToken}
-                    modelData={modelData}
+                    modelData={processedModelData}
                     all_models_on_proxy={allModelsOnProxy}
                     getDisplayModelName={getDisplayModelName}
                     setSelectedModelId={setSelectedModelId}

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx

@@ -10,6 +10,28 @@ vi.mock("../../hooks/models/useModels", () => ({
   useModelsInfo: () => mockUseModelsInfo(),
 }));
 
+// Mock the useModelCostMap hook
+const mockUseModelCostMap = vi.fn(() => ({
+  data: {
+    "gpt-4": { litellm_provider: "openai" },
+    "gpt-3.5-turbo": { litellm_provider: "openai" },
+    "gpt-4-accessible": { litellm_provider: "openai" },
+    "gpt-3.5-turbo-blocked": { litellm_provider: "openai" },
+    "gpt-4-sales": { litellm_provider: "openai" },
+    "gpt-4-engineering": { litellm_provider: "openai" },
+    "gpt-4-personal": { litellm_provider: "openai" },
+    "gpt-4-team-only": { litellm_provider: "openai" },
+    "gpt-4-config": { litellm_provider: "openai" },
+    "gpt-4-db": { litellm_provider: "openai" },
+  },
+  isLoading: false,
+  error: null,
+})) as any;
+
+vi.mock("../../hooks/models/useModelCostMap", () => ({
+  useModelCostMap: () => mockUseModelCostMap(),
+}));
+
 // Mock the useTeams hook (react-query implementation)
 const mockUseTeams = vi.fn(() => ({
   data: [],
@@ -22,6 +44,13 @@ vi.mock("../../hooks/teams/useTeams", () => ({
   useTeams: () => mockUseTeams(),
 }));
 
+// Helper function to create model cost map mock return value
+const createModelCostMapMock = (data: Record<string, any>) => ({
+  data,
+  isLoading: false,
+  error: null,
+});
+
 describe("AllModelsTab", () => {
   const mockSetSelectedModelGroup = vi.fn();
   const mockSetSelectedModelId = vi.fn();
@@ -64,6 +93,8 @@ describe("AllModelsTab", () => {
       refetch: vi.fn(),
     });
 
+    mockUseModelCostMap.mockReturnValueOnce(createModelCostMapMock({}));
+
     render(<AllModelsTab {...defaultProps} />);
     expect(screen.getByText("Current Team:")).toBeInTheDocument();
   });
@@ -92,6 +123,13 @@ describe("AllModelsTab", () => {
       refetch: vi.fn(),
     });
 
+    mockUseModelCostMap.mockReturnValueOnce(
+      createModelCostMapMock({
+        "gpt-4-accessible": { litellm_provider: "openai" },
+        "gpt-3.5-turbo-blocked": { litellm_provider: "openai" },
+      }),
+    );
+
     const modelData = {
       data: [
         {
@@ -146,6 +184,13 @@ describe("AllModelsTab", () => {
       refetch: vi.fn(),
     });
 
+    mockUseModelCostMap.mockReturnValueOnce(
+      createModelCostMapMock({
+        "gpt-4-sales": { litellm_provider: "openai" },
+        "gpt-4-engineering": { litellm_provider: "openai" },
+      }),
+    );
+
     const modelData = {
       data: [
         {
@@ -184,6 +229,13 @@ describe("AllModelsTab", () => {
       refetch: vi.fn(),
     });
 
+    mockUseModelCostMap.mockReturnValueOnce(
+      createModelCostMapMock({
+        "gpt-4-personal": { litellm_provider: "openai" },
+        "gpt-4-team-only": { litellm_provider: "openai" },
+      }),
+    );
+
     const modelData = {
       data: [
         {
@@ -224,6 +276,13 @@ describe("AllModelsTab", () => {
       refetch: vi.fn(),
     });
 
+    mockUseModelCostMap.mockReturnValueOnce(
+      createModelCostMapMock({
+        "gpt-4-config": { litellm_provider: "openai" },
+        "gpt-4-db": { litellm_provider: "openai" },
+      }),
+    );
+
     const modelData = {
       data: [
         {
@@ -277,6 +336,12 @@ describe("AllModelsTab", () => {
       refetch: vi.fn(),
     });
 
+    mockUseModelCostMap.mockReturnValueOnce(
+      createModelCostMapMock({
+        "gpt-4-config": { litellm_provider: "openai" },
+      }),
+    );
+
     const modelData = {
       data: [
         {

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

@@ -1,9 +1,11 @@
 import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
+import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap";
 import { Team } from "@/components/key_team_helpers/key_list";
 import { ModelDataTable } from "@/components/model_dashboard/table";
 import { columns } from "@/components/molecules/models/columns";
 import { getDisplayModelName } from "@/components/view_model/model_name_display";
+import { transformModelData } from "../utils/modelDataTransformer";
 import { InfoCircleOutlined } from "@ant-design/icons";
 import { PaginationState } from "@tanstack/react-table";
 import { Grid, Select, SelectItem, TabPanel, Text } from "@tremor/react";
@@ -29,10 +31,25 @@ const AllModelsTab = ({
   setSelectedModelId,
   setSelectedTeamId,
 }: AllModelsTabProps) => {
-  const { data: modelData } = useModelsInfo();
+  const { data: rawModelData } = useModelsInfo();
+  const { data: modelCostMapData } = useModelCostMap();
   const { userId, userRole, premiumUser } = useAuthorized();
   const { data: teams } = useTeams();
 
+  const getProviderFromModel = (model: string) => {
+    if (modelCostMapData !== null && modelCostMapData !== undefined) {
+      if (typeof modelCostMapData == "object" && model in modelCostMapData) {
+        return modelCostMapData[model]["litellm_provider"];
+      }
+    }
+    return "openai";
+  };
+
+  const modelData = useMemo(() => {
+    if (!rawModelData) return { data: [] };
+    return transformModelData(rawModelData, getProviderFromModel);
+  }, [rawModelData, modelCostMapData]);
+
   const [modelNameSearch, setModelNameSearch] = useState<string>("");
   const [modelViewMode, setModelViewMode] = useState<ModelViewMode>("current_team");
   const [currentTeam, setCurrentTeam] = useState<Team | "personal">("personal");

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.test.ts

@@ -0,0 +1,53 @@
+import { transformModelData } from "./modelDataTransformer";
+import { describe, it, expect } from "vitest";
+describe("transformModelData", () => {
+  const mockGetProviderFromModel = (model: string) => {
+    if (model.includes("gpt")) return "openai";
+    if (model.includes("claude")) return "anthropic";
+    return "openai";
+  };
+
+  it("should transform raw model data correctly", () => {
+    const rawData = {
+      data: [
+        {
+          model_name: "gpt-4",
+          litellm_params: {
+            model: "gpt-4",
+            api_base: "https://api.openai.com",
+            api_key: "sk-123",
+          },
+          model_info: {
+            input_cost_per_token: 0.0000015,
+            output_cost_per_token: 0.000002,
+            max_tokens: 8192,
+            max_input_tokens: 128000,
+          },
+        },
+      ],
+    };
+
+    const result = transformModelData(rawData, mockGetProviderFromModel);
+
+    expect(result.data[0]).toHaveProperty("provider", "openai");
+    expect(result.data[0]).toHaveProperty("input_cost", "1.50");
+    expect(result.data[0]).toHaveProperty("output_cost", "2.00");
+    expect(result.data[0]).toHaveProperty("max_tokens", 8192);
+    expect(result.data[0]).toHaveProperty("max_input_tokens", 128000);
+    expect(result.data[0]).toHaveProperty("api_base", "https://api.openai.com");
+    expect(result.data[0]).toHaveProperty("litellm_model_name", "gpt-4");
+    expect(result.data[0]).toHaveProperty("cleanedLitellmParams");
+    expect(result.data[0].cleanedLitellmParams).not.toHaveProperty("model");
+    expect(result.data[0].cleanedLitellmParams).not.toHaveProperty("api_base");
+  });
+
+  it("should handle empty data", () => {
+    const result = transformModelData({ data: [] }, mockGetProviderFromModel);
+    expect(result).toEqual({ data: [] });
+  });
+
+  it("should handle null/undefined data", () => {
+    const result = transformModelData(null, mockGetProviderFromModel);
+    expect(result).toEqual({ data: [] });
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.ts

@@ -0,0 +1,76 @@
+/**
+ * Utility function to transform raw model data into the format expected by UI components
+ * This creates a new transformed data object without mutating the original
+ */
+export const transformModelData = (rawModelData: any, getProviderFromModel: (model: string) => string) => {
+  if (!rawModelData?.data) return { data: [] };
+
+  // Deep copy the data to avoid mutating the original
+  const transformedData = JSON.parse(JSON.stringify(rawModelData.data));
+
+  for (let i = 0; i < transformedData.length; i++) {
+    let curr_model = transformedData[i];
+    let litellm_model_name = curr_model?.litellm_params?.model;
+    let custom_llm_provider = curr_model?.litellm_params?.custom_llm_provider;
+    let model_info = curr_model?.model_info;
+
+    let provider = "";
+    let input_cost = "Undefined";
+    let output_cost = "Undefined";
+    let max_tokens = "Undefined";
+    let max_input_tokens = "Undefined";
+    let cleanedLitellmParams = {};
+
+    // Check if litellm_model_name is null or undefined
+    if (litellm_model_name) {
+      // Split litellm_model_name based on "/"
+      let splitModel = litellm_model_name.split("/");
+
+      // Get the first element in the split
+      let firstElement = splitModel[0];
+
+      // If there is only one element, default provider to openai
+      provider = custom_llm_provider;
+      if (!provider) {
+        provider = splitModel.length === 1 ? getProviderFromModel(litellm_model_name) : firstElement;
+      }
+    } else {
+      // litellm_model_name is null or undefined, default provider to openai
+      provider = "-";
+    }
+
+    if (model_info) {
+      input_cost = model_info?.input_cost_per_token;
+      output_cost = model_info?.output_cost_per_token;
+      max_tokens = model_info?.max_tokens;
+      max_input_tokens = model_info?.max_input_tokens;
+    }
+
+    if (curr_model?.litellm_params) {
+      cleanedLitellmParams = Object.fromEntries(
+        Object.entries(curr_model?.litellm_params).filter(([key]) => key !== "model" && key !== "api_base"),
+      );
+    }
+
+    transformedData[i].provider = provider;
+    transformedData[i].input_cost = input_cost;
+    transformedData[i].output_cost = output_cost;
+    transformedData[i].litellm_model_name = litellm_model_name;
+
+    // Convert Cost in terms of Cost per 1M tokens
+    if (transformedData[i].input_cost) {
+      transformedData[i].input_cost = (Number(transformedData[i].input_cost) * 1000000).toFixed(2);
+    }
+
+    if (transformedData[i].output_cost) {
+      transformedData[i].output_cost = (Number(transformedData[i].output_cost) * 1000000).toFixed(2);
+    }
+
+    transformedData[i].max_tokens = max_tokens;
+    transformedData[i].max_input_tokens = max_input_tokens;
+    transformedData[i].api_base = curr_model?.litellm_params?.api_base;
+    transformedData[i].cleanedLitellmParams = cleanedLitellmParams;
+  }
+
+  return { data: transformedData };
+};