diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.test.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.test.tsx new file mode 100644 index 00000000000..225c308af91 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.test.tsx @@ -0,0 +1,87 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import EditFallbacks, { Fallbacks } from "./EditFallbacks"; +import * as fetchModelsModule from "@/components/llm_calls/fetch_models"; + +vi.mock("@/components/llm_calls/fetch_models", () => ({ + fetchAvailableModels: vi.fn(), +})); + +const renderWithQueryClient = (ui: React.ReactElement) => { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + return render({ui}); +}; + +describe("EditFallbacks", () => { + const accessToken = "test-token"; + const fallbackEntry = { "gpt-4": ["gpt-3.5-turbo", "claude-3-opus"] }; + const value: Fallbacks = [{ "gpt-4": ["gpt-3.5-turbo", "claude-3-opus"] }, { "claude-3-opus": ["gpt-4"] }]; + + const setup = (overrides: Partial> = {}) => { + const onChange = overrides.onChange ?? vi.fn().mockResolvedValue(undefined); + const onClose = overrides.onClose ?? vi.fn(); + renderWithQueryClient( + , + ); + return { onChange, onClose }; + }; + + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(fetchModelsModule.fetchAvailableModels).mockResolvedValue([ + { model_group: "gpt-4", mode: "chat" }, + { model_group: "gpt-3.5-turbo", mode: "chat" }, + { model_group: "claude-3-opus", mode: "chat" }, + { model_group: "gemini-pro", mode: "chat" }, + ]); + }); + + it("prefills the existing fallback chain for the primary model", async () => { + setup(); + await waitFor(() => { + expect(screen.getByText("gpt-3.5-turbo")).toBeInTheDocument(); + expect(screen.getByText("claude-3-opus")).toBeInTheDocument(); + }); + }); + + it("removes a fallback model and saves only the edited entry", async () => { + const user = userEvent.setup(); + const onChange = vi.fn().mockResolvedValue(undefined); + const onClose = vi.fn(); + setup({ onChange, onClose }); + + await screen.findByText("gpt-3.5-turbo"); + await user.click(screen.getByTestId("remove-fallback-gpt-3.5-turbo")); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onChange).toHaveBeenCalledWith([{ "gpt-4": ["claude-3-opus"] }, { "claude-3-opus": ["gpt-4"] }]); + }); + await waitFor(() => expect(onClose).toHaveBeenCalled()); + }); + + it("blocks saving with an empty fallback chain", async () => { + const user = userEvent.setup(); + const onChange = vi.fn().mockResolvedValue(undefined); + setup({ fallbackEntry: { "gpt-4": ["gpt-3.5-turbo"] }, onChange }); + + await screen.findByText("gpt-3.5-turbo"); + await user.click(screen.getByTestId("remove-fallback-gpt-3.5-turbo")); + + const saveButton = screen.getByRole("button", { name: /save changes/i }); + expect(saveButton).toBeDisabled(); + expect(onChange).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx new file mode 100644 index 00000000000..938e1104301 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/EditFallbacks.tsx @@ -0,0 +1,106 @@ +/** + * Modal for editing an existing fallback entry + * Lets the user add/remove models from a primary model's fallback chain + * Reuses FallbackGroupConfig with the primary model locked + */ + +import { Button } from "antd"; +import { useQuery } from "@tanstack/react-query"; +import { Pencil } from "lucide-react"; +import React, { useMemo, useState } from "react"; +import { fetchAvailableModels } from "@/components/llm_calls/fetch_models"; +import NotificationManager from "../../../molecules/notifications_manager"; +import { AddFallbacksModal } from "./AddFallbacksModal"; +import { FallbackGroup, FallbackGroupConfig } from "./FallbackGroupConfig"; + +export type FallbackEntry = { [modelName: string]: string[] }; +export type Fallbacks = FallbackEntry[]; + +interface EditFallbacksProps { + accessToken: string; + fallbackEntry: FallbackEntry; + value: Fallbacks; + onChange: (fallbacks: Fallbacks) => Promise; + onClose: () => void; + maxFallbacks?: number; +} + +const toGroup = (entry: FallbackEntry): FallbackGroup => { + const primaryModel = Object.keys(entry)[0] ?? null; + return { + id: "edit", + primaryModel, + fallbackModels: primaryModel ? [...(entry[primaryModel] ?? [])] : [], + }; +}; + +export default function EditFallbacks({ + accessToken, + fallbackEntry, + value, + onChange, + onClose, + maxFallbacks = 10, +}: EditFallbacksProps) { + const [group, setGroup] = useState(() => toGroup(fallbackEntry)); + const [isSaving, setIsSaving] = useState(false); + + const { data: modelGroups = [] } = useQuery({ + queryKey: ["availableModels", "fallbacks"], + queryFn: () => fetchAvailableModels(accessToken), + enabled: Boolean(accessToken), + }); + + const availableModels = useMemo( + () => Array.from(new Set(modelGroups.map((option) => option.model_group))).sort(), + [modelGroups], + ); + + const handleSave = async () => { + const primaryModel = group.primaryModel; + if (!primaryModel) { + return; + } + + const updatedFallbacks = (value || []).map((entry) => + primaryModel in entry ? { ...entry, [primaryModel]: group.fallbackModels } : entry, + ); + + setIsSaving(true); + try { + await onChange(updatedFallbacks); + NotificationManager.success(`Fallbacks for ${primaryModel} updated successfully!`); + onClose(); + } catch (error) { + console.error("Error updating fallbacks:", error); + } finally { + setIsSaving(false); + } + }; + + return ( + + +
+ + +
+
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx index e91c938b87a..381818a53f5 100644 --- a/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx +++ b/ui/litellm-dashboard/src/components/Settings/RouterSettings/Fallbacks/FallbackGroupConfig.tsx @@ -18,9 +18,16 @@ interface FallbackGroupConfigProps { onChange: (updatedGroup: FallbackGroup) => void; availableModels: string[]; maxFallbacks: number; + disablePrimaryModel?: boolean; } -export function FallbackGroupConfig({ group, onChange, availableModels, maxFallbacks }: FallbackGroupConfigProps) { +export function FallbackGroupConfig({ + group, + onChange, + availableModels, + maxFallbacks, + disablePrimaryModel = false, +}: FallbackGroupConfigProps) { // Filter available options for fallbacks (exclude primary only, allow already selected to be shown for deselection) const availableFallbackOptions = availableModels.filter((m) => m !== group.primaryModel); @@ -70,12 +77,13 @@ export function FallbackGroupConfig({ group, onChange, availableModels, maxFallb placeholder="Select primary model" value={group.primaryModel} onChange={handlePrimaryChange} + disabled={disablePrimaryModel} showSearch getPopupContainer={(trigger) => trigger.parentElement || document.body} filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())} options={availableModels.map((m) => ({ label: m, value: m }))} /> - {!group.primaryModel && ( + {!disablePrimaryModel && !group.primaryModel && (
Select a model to begin configuring fallbacks @@ -176,6 +184,7 @@ export function FallbackGroupConfig({ group, onChange, availableModels, maxFallb