diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts new file mode 100644 index 00000000000..6ff784ebd90 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.test.ts @@ -0,0 +1,147 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useStoreModelInDB } from "./useStoreModelInDB"; + +vi.mock("@/components/networking", () => ({ + getProxyBaseUrl: vi.fn(() => ""), + getGlobalLitellmHeaderName: vi.fn(() => "Authorization"), +})); + +describe("useStoreModelInDB", () => { + let queryClient: QueryClient; + let fetchSpy: ReturnType; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + mutations: { retry: false }, + }, + }); + + fetchSpy = vi.fn(); + global.fetch = fetchSpy; + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should send correct request body to /config/field/update", async () => { + fetchSpy.mockResolvedValue({ + ok: true, + json: async () => ({ message: "Success" }), + }); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: true }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(fetchSpy).toHaveBeenCalledWith( + "/config/field/update", + expect.objectContaining({ + method: "POST", + body: JSON.stringify({ + field_name: "store_model_in_db", + field_value: true, + config_type: "general_settings", + }), + }) + ); + }); + + it("should handle setting store_model_in_db to false", async () => { + fetchSpy.mockResolvedValue({ + ok: true, + json: async () => ({ message: "Success" }), + }); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: false }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(fetchSpy).toHaveBeenCalledWith( + "/config/field/update", + expect.objectContaining({ + body: JSON.stringify({ + field_name: "store_model_in_db", + field_value: false, + config_type: "general_settings", + }), + }) + ); + }); + + it("should throw error when access token is missing", async () => { + vi.spyOn( + await import("../useAuthorized"), + "default" + ).mockReturnValue({ + accessToken: null, + userRole: null, + userId: null, + token: null, + userEmail: null, + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + } as any); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: true }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error?.message).toBe("Access token is required"); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + it("should handle API error response", async () => { + fetchSpy.mockResolvedValue({ + ok: false, + json: async () => ({ detail: "Unauthorized" }), + }); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: true }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error?.message).toBe("Unauthorized"); + }); + + it("should use fallback error message when API returns empty error", async () => { + fetchSpy.mockResolvedValue({ + ok: false, + json: async () => ({}), + }); + + const { result } = renderHook(() => useStoreModelInDB(), { wrapper }); + + result.current.mutate({ store_model_in_db: true }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error?.message).toBe("Failed to update model storage settings"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts new file mode 100644 index 00000000000..e6efbd724cd --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB.ts @@ -0,0 +1,59 @@ +import { useMutation, UseMutationResult } from "@tanstack/react-query"; +import { getProxyBaseUrl, getGlobalLitellmHeaderName } from "@/components/networking"; +import useAuthorized from "../useAuthorized"; + +export interface StoreModelInDBParams { + store_model_in_db: boolean; +} + +export interface StoreModelInDBResponse { + message: string; +} + +const performStoreModelInDB = async ( + accessToken: string, + params: StoreModelInDBParams +): Promise => { + const proxyBaseUrl = getProxyBaseUrl(); + const url = proxyBaseUrl ? `${proxyBaseUrl}/config/field/update` : `/config/field/update`; + + const response = await fetch(url, { + method: "POST", + headers: { + [getGlobalLitellmHeaderName()]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ + field_name: "store_model_in_db", + field_value: params.store_model_in_db, + config_type: "general_settings", + }), + }); + + if (!response.ok) { + const errorData = await response.json().catch(() => ({})); + const errorMessage = + errorData?.error?.message || errorData?.message || errorData?.detail || "Failed to update model storage settings"; + throw new Error(errorMessage); + } + + const data = await response.json(); + return data; +}; + +export const useStoreModelInDB = (): UseMutationResult< + StoreModelInDBResponse, + Error, + StoreModelInDBParams +> => { + const { accessToken } = useAuthorized(); + + return useMutation({ + mutationFn: async (params: StoreModelInDBParams) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return await performStoreModelInDB(accessToken, params); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index e252f273316..36948630d8a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -5,10 +5,11 @@ import { Team } from "@/components/key_team_helpers/key_list"; import { AllModelsDataTable } from "@/components/model_dashboard/all_models_table"; import { columns } from "@/components/molecules/models/columns"; import { getDisplayModelName } from "@/components/view_model/model_name_display"; -import { InfoCircleOutlined } from "@ant-design/icons"; +import { InfoCircleOutlined, SettingOutlined } from "@ant-design/icons"; import { PaginationState, SortingState } from "@tanstack/react-table"; import { Grid, TabPanel } from "@tremor/react"; -import { Badge, Select, Skeleton, Space, Typography } from "antd"; +import { Badge, Button, Select, Skeleton, Space, Typography } from "antd"; +import ModelSettingsModal from "@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal"; import debounce from "lodash/debounce"; import { useEffect, useMemo, useState } from "react"; import { useModelsInfo } from "../../hooks/models/useModels"; @@ -51,6 +52,7 @@ const AllModelsTab = ({ pageSize: 50, }); const [sorting, setSorting] = useState([]); + const [isModelSettingsModalVisible, setIsModelSettingsModalVisible] = useState(false); // Debounce search input const debouncedUpdateSearch = useMemo( @@ -326,62 +328,71 @@ const AllModelsTab = ({
{/* Search and Filter Controls */} -
- {/* Model Name Search */} -
- setModelNameSearch(e.target.value)} - /> - - +
+ {/* Model Name Search */} +
+ setModelNameSearch(e.target.value)} /> - + + + +
+ + {/* Filter Button */} + + + {/* Reset Filters Button */} +
- {/* Filter Button */} - - - {/* Reset Filters Button */} - + {/* Model Settings Button */} +
{/* Additional Filters */} @@ -505,6 +516,11 @@ const AllModelsTab = ({
+ setIsModelSettingsModalVisible(false)} + onSuccess={() => setIsModelSettingsModalVisible(false)} + /> ); }; diff --git a/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.test.tsx new file mode 100644 index 00000000000..d4103a1b38e --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.test.tsx @@ -0,0 +1,306 @@ +import { useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"; +import { useStoreModelInDB } from "@/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; +import ModelSettingsModal from "./ModelSettingsModal"; + +vi.mock("@/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB"); +vi.mock("@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"); +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn(), +})); + +const mockUseStoreModelInDB = vi.mocked(useStoreModelInDB); +const mockUseProxyConfig = vi.mocked(useProxyConfig); +const mockNotificationsManager = vi.mocked(NotificationsManager); +const mockParseErrorMessage = vi.mocked(parseErrorMessage); + +describe("ModelSettingsModal", () => { + const mockOnCancel = vi.fn(); + const mockOnSuccess = vi.fn(); + const mockMutateAsync = vi.fn(); + const mockRefetch = vi.fn(); + + const defaultProps = { + isVisible: true, + onCancel: mockOnCancel, + onSuccess: mockOnSuccess, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockUseStoreModelInDB.mockReturnValue({ + mutateAsync: mockMutateAsync, + isPending: false, + } as any); + mockUseProxyConfig.mockReturnValue({ + data: [], + isLoading: false, + refetch: mockRefetch, + } as any); + mockParseErrorMessage.mockImplementation((error: any) => error?.message || String(error)); + }); + + it("should render the modal", () => { + renderWithProviders(); + expect(screen.getByRole("dialog")).toBeInTheDocument(); + expect(screen.getByText("Model Settings")).toBeInTheDocument(); + }); + + it("should render form field with initial values", () => { + renderWithProviders(); + expect(screen.getByText("Store Model in DB")).toBeInTheDocument(); + }); + + it("should render cancel and save buttons", () => { + renderWithProviders(); + expect(screen.getByRole("button", { name: "Cancel" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Save Settings" })).toBeInTheDocument(); + }); + + it("should call onCancel when cancel button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + await user.click(cancelButton); + + expect(mockOnCancel).toHaveBeenCalledTimes(1); + }); + + it("should call onCancel when modal close button is clicked", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const closeButton = screen.getByRole("button", { name: /close/i }); + await user.click(closeButton); + + expect(mockOnCancel).toHaveBeenCalledTimes(1); + }); + + it("should toggle store model switch", async () => { + const user = userEvent.setup(); + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + expect(switchElement).not.toBeChecked(); + + await user.click(switchElement); + + await waitFor(() => { + expect(switchElement).toBeChecked(); + }); + }); + + it("should submit form with store_model_in_db enabled", async () => { + const user = userEvent.setup(); + mockMutateAsync.mockImplementation(async (params, options) => { + await Promise.resolve(); + options?.onSuccess?.(); + return { message: "Success" }; + }); + + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + await user.click(switchElement); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockMutateAsync).toHaveBeenCalledWith( + { store_model_in_db: true }, + expect.any(Object) + ); + }); + }); + + it("should submit form with store_model_in_db disabled", async () => { + const user = userEvent.setup(); + mockMutateAsync.mockImplementation(async (params, options) => { + await Promise.resolve(); + options?.onSuccess?.(); + return { message: "Success" }; + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockMutateAsync).toHaveBeenCalledWith( + { store_model_in_db: false }, + expect.any(Object) + ); + }); + }); + + it("should show success notification and call onSuccess on successful submission", async () => { + const user = userEvent.setup(); + mockMutateAsync.mockImplementation(async (params, options) => { + await Promise.resolve(); + options?.onSuccess?.(); + return { message: "Success" }; + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalledWith("Model storage settings updated successfully"); + expect(mockRefetch).toHaveBeenCalled(); + expect(mockOnSuccess).toHaveBeenCalledTimes(1); + }); + }); + + it("should show error notification when submission fails", async () => { + const user = userEvent.setup(); + const error = new Error("Network error"); + mockMutateAsync.mockRejectedValue(error); + mockParseErrorMessage.mockReturnValue("Network error"); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to save model storage settings: Network error"); + }); + }); + + it("should show error notification from onError callback", async () => { + const user = userEvent.setup(); + const error = new Error("Backend error"); + mockMutateAsync.mockImplementation((params, options) => { + options?.onError?.(error); + return Promise.reject(error); + }); + mockParseErrorMessage.mockReturnValue("Backend error"); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.fromBackend).toHaveBeenCalledWith("Failed to save model storage settings: Backend error"); + }); + }); + + it("should disable cancel button when pending", () => { + mockUseStoreModelInDB.mockReturnValue({ + mutateAsync: mockMutateAsync, + isPending: true, + } as any); + + renderWithProviders(); + + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + expect(cancelButton).toBeDisabled(); + }); + + it("should disable cancel button when loading config", () => { + mockUseProxyConfig.mockReturnValue({ + data: undefined, + isLoading: true, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + expect(cancelButton).toBeDisabled(); + }); + + it("should show loading state on save button when pending", () => { + mockUseStoreModelInDB.mockReturnValue({ + mutateAsync: mockMutateAsync, + isPending: true, + } as any); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: /Saving/i }); + expect(saveButton).toBeInTheDocument(); + expect(saveButton.className).toContain("ant-btn-loading"); + }); + + it("should not render modal when isVisible is false", () => { + renderWithProviders(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + }); + + it("should call refetch when modal opens", () => { + renderWithProviders(); + expect(mockRefetch).toHaveBeenCalledTimes(1); + }); + + it("should render form with initial values from config data", () => { + mockUseProxyConfig.mockReturnValue({ + data: [ + { + field_name: "store_model_in_db", + field_type: "bool", + field_description: "Store model in DB", + field_value: true, + stored_in_db: true, + field_default_value: false, + }, + ], + isLoading: false, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + const switchElement = screen.getByRole("switch"); + expect(switchElement).toBeChecked(); + }); + + it("should show skeleton loader when config is loading", () => { + mockUseProxyConfig.mockReturnValue({ + data: undefined, + isLoading: true, + refetch: mockRefetch, + } as any); + + renderWithProviders(); + + expect(screen.queryByRole("switch")).not.toBeInTheDocument(); + const skeletons = document.querySelectorAll(".ant-skeleton"); + expect(skeletons.length).toBeGreaterThan(0); + }); + + it("should not call onSuccess when it is not provided", async () => { + const user = userEvent.setup(); + mockMutateAsync.mockImplementation(async (params, options) => { + await Promise.resolve(); + options?.onSuccess?.(); + return { message: "Success" }; + }); + + renderWithProviders(); + + const saveButton = screen.getByRole("button", { name: "Save Settings" }); + await user.click(saveButton); + + await waitFor(() => { + expect(mockNotificationsManager.success).toHaveBeenCalled(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.tsx b/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.tsx new file mode 100644 index 00000000000..106f82585dd --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_dashboard/ModelSettingsModal/ModelSettingsModal.tsx @@ -0,0 +1,104 @@ +"use client"; + +import { ConfigType, useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig"; +import { StoreModelInDBParams, useStoreModelInDB } from "@/app/(dashboard)/hooks/storeModelInDB/useStoreModelInDB"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { Button, Form, Modal, Skeleton, Space, Switch, Typography } from "antd"; +import React, { useEffect, useMemo } from "react"; + +interface ModelSettingsModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess?: () => void; +} + +const ModelSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { + const [form] = Form.useForm(); + const { mutateAsync, isPending } = useStoreModelInDB(); + const { data: proxyConfigData, isLoading: isLoadingConfig, refetch } = useProxyConfig(ConfigType.GENERAL_SETTINGS); + + // Refetch config when modal opens to ensure we have the latest values + useEffect(() => { + if (isVisible) { + refetch(); + } + }, [isVisible, refetch]); + + // Compute initial values from fetched config data + const initialValues = useMemo(() => { + if (!proxyConfigData) { + return { + store_model_in_db: false, + }; + } + + const storeModelField = proxyConfigData.find(field => field.field_name === 'store_model_in_db'); + + return { + store_model_in_db: storeModelField?.field_value ?? false, + }; + }, [proxyConfigData]); + + const handleFormSubmit = async (formValues: StoreModelInDBParams) => { + try { + await mutateAsync(formValues, { + onSuccess: () => { + NotificationsManager.success("Model storage settings updated successfully"); + refetch(); + onSuccess?.(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save model storage settings: " + parseErrorMessage(error)); + }, + }); + } catch (error) { + NotificationsManager.fromBackend("Failed to save model storage settings: " + parseErrorMessage(error)); + } + }; + + const handleCancel = () => { + form.resetFields(); + onCancel(); + }; + + return ( + Model Settings} + open={isVisible} + footer={ + + + + + } + onCancel={handleCancel} + > +
+ f.field_name === 'store_model_in_db')?.field_description || + "If enabled, models and config are stored in and loaded from the database." + } + valuePropName="checked" + > + {isLoadingConfig ? : } + +
+ + ); +}; + +export default ModelSettingsModal;