diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.test.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.test.tsx new file mode 100644 index 00000000000..194f717a92f --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.test.tsx @@ -0,0 +1,82 @@ +import { render, screen } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { describe, expect, it, vi } from "vitest"; +import { BudgetFallbacksEditor } from "./BudgetFallbacksEditor"; + +const MODELS = ["gpt-4", "gpt-3.5-turbo", "claude-3", "claude-haiku"]; + +describe("BudgetFallbacksEditor", () => { + it("renders empty state with add button", () => { + const onChange = vi.fn(); + render(); + expect(screen.getByText("Add Budget Fallback")).toBeTruthy(); + expect(screen.getByText(/reroute to fallback models/)).toBeTruthy(); + }); + + it("renders existing entries from value prop", () => { + const onChange = vi.fn(); + render( + , + ); + expect(screen.getByText("IF BUDGET EXCEEDED, TRY")).toBeTruthy(); + expect(screen.getByText("Primary Model")).toBeTruthy(); + expect(screen.getByText("Fallback Models")).toBeTruthy(); + }); + + it("adds a new empty entry when clicking add button", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + render(); + + await user.click(screen.getByText("Add Budget Fallback")); + expect(screen.getByText("Primary Model")).toBeTruthy(); + expect(onChange).toHaveBeenCalledWith({}); + }); + + it("removes an entry and emits updated dict", async () => { + const user = userEvent.setup(); + const onChange = vi.fn(); + const { container } = render( + , + ); + + const removeButtons = container.querySelectorAll(".relative > button[type='button']"); + expect(removeButtons.length).toBe(2); + + await user.click(removeButtons[0]); + expect(onChange).toHaveBeenLastCalledWith({ "claude-3": ["claude-haiku"] }); + }); + + it("renders multiple entries for multiple fallback groups", () => { + const onChange = vi.fn(); + render( + , + ); + const labels = screen.getAllByText("Primary Model"); + expect(labels.length).toBe(2); + }); + + it("shows ordering hint when multiple fallback models are configured", () => { + const onChange = vi.fn(); + render( + , + ); + expect(screen.getByText(/first model still within its own budget/)).toBeTruthy(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx new file mode 100644 index 00000000000..e58d289df8a --- /dev/null +++ b/ui/litellm-dashboard/src/components/key_team_helpers/BudgetFallbacksEditor.tsx @@ -0,0 +1,153 @@ +import { Button, Select, Tooltip } from "antd"; +import { ArrowDown, Plus, X } from "lucide-react"; +import React, { useState } from "react"; + +interface FallbackEntry { + id: string; + primaryModel: string | null; + fallbackModels: string[]; +} + +interface BudgetFallbacksEditorProps { + value: Record; + onChange: (v: Record) => void; + availableModels: string[]; +} + +const entriesToDict = (entries: readonly FallbackEntry[]): Record => + Object.fromEntries( + entries + .filter( + (e): e is FallbackEntry & { primaryModel: string } => e.primaryModel !== null && e.fallbackModels.length > 0, + ) + .map((e) => [e.primaryModel, e.fallbackModels]), + ); + +const dictToEntries = (dict: Record): FallbackEntry[] => { + const keys = Object.keys(dict); + if (keys.length === 0) return []; + return keys.map((model, i) => ({ + id: String(i + 1), + primaryModel: model, + fallbackModels: dict[model], + })); +}; + +export function BudgetFallbacksEditor({ value, onChange, availableModels }: BudgetFallbacksEditorProps) { + const [entries, setEntries] = useState(() => dictToEntries(value)); + + const emitChange = (updated: FallbackEntry[]) => { + setEntries(updated); + onChange(entriesToDict(updated)); + }; + + const addEntry = () => { + emitChange([...entries, { id: Date.now().toString(), primaryModel: null, fallbackModels: [] }]); + }; + + const removeEntry = (id: string) => { + emitChange(entries.filter((e) => e.id !== id)); + }; + + const updateEntry = (id: string, patch: Partial) => { + emitChange(entries.map((e) => (e.id === id ? { ...e, ...patch } : e))); + }; + + const usedPrimaryModels = new Set(entries.map((e) => e.primaryModel).filter(Boolean)); + + if (entries.length === 0) { + return ( +
+
+ When a model exceeds its per-model budget, requests automatically reroute to fallback models +
+ +
+ ); + } + + return ( +
+
+ When a model exceeds its per-model budget, requests automatically reroute to fallback models +
+ {entries.map((entry) => { + const availablePrimaryOptions = availableModels.filter( + (m) => m === entry.primaryModel || !usedPrimaryModels.has(m), + ); + const availableFallbackOptions = availableModels.filter((m) => m !== entry.primaryModel); + + return ( +
+ + +
+ + updateEntry(entry.id, { fallbackModels: values })} + disabled={!entry.primaryModel} + showSearch + filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())} + options={availableFallbackOptions.map((m) => ({ label: m, value: m }))} + getPopupContainer={(trigger) => trigger.parentElement || document.body} + maxTagCount="responsive" + maxTagPlaceholder={(omittedValues) => ( + v).join(", ")} + > + +{omittedValues.length} more + + )} + /> + {entry.fallbackModels.length > 1 && ( +
+ Tried in order; first model still within its own budget is used +
+ )} +
+
+ ); + })} + +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 60568da48ed..99fe23792d6 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -97,6 +97,7 @@ export interface KeyResponse { agent_access_groups?: string[]; }; access_group_ids?: string[]; + budget_fallbacks?: Record; budget_limits?: Array<{ budget_duration: string; max_budget: number; reset_at?: string }>; auto_rotate?: boolean; rotation_interval?: string; diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 2a448d63300..443cbd34e32 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -28,6 +28,7 @@ import TeamDropdown from "../common_components/team_dropdown"; import OrganizationDropdown from "../common_components/OrganizationDropdown"; import ProjectDropdown from "../common_components/ProjectDropdown"; import { CreateUserButton } from "../CreateUserButton"; +import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; import { Team } from "../key_team_helpers/key_list"; @@ -202,6 +203,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp const [rotationInterval, setRotationInterval] = useState("30d"); const [routerSettings, setRouterSettings] = useState(null); const [budgetLimits, setBudgetLimits] = useState([]); + const [budgetFallbacks, setBudgetFallbacks] = useState>({}); const [routerSettingsKey, setRouterSettingsKey] = useState(0); const [agentsList, setAgentsList] = useState<{ agent_id: string; agent_name: string }[]>([]); const [selectedAgentId, setSelectedAgentId] = useState(null); @@ -220,6 +222,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setSelectedOrganizationId(null); setSelectedProjectId(null); setBudgetLimits([]); + setBudgetFallbacks({}); }; const handleCancel = () => { @@ -239,6 +242,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp setSelectedOrganizationId(null); setSelectedProjectId(null); setBudgetLimits([]); + setBudgetFallbacks({}); }; useEffect(() => { @@ -536,6 +540,10 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp formValues.budget_limits = validWindows; } + if (Object.keys(budgetFallbacks).length > 0) { + formValues.budget_fallbacks = budgetFallbacks; + } + let response; if (keyOwner === "service_account") { response = await keyCreateServiceAccountCall(accessToken, formValues); @@ -558,6 +566,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp NotificationsManager.success("Virtual Key Created"); form.resetFields(); setBudgetLimits([]); + setBudgetFallbacks({}); localStorage.removeItem("userData" + userID); } catch (error) { console.log("error in create key:", error); @@ -1076,6 +1085,23 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp > + + Budget Fallbacks{" "} + + + + + } + > + + ( Array.isArray(keyData.budget_limits) ? keyData.budget_limits : [], ); + const [budgetFallbacks, setBudgetFallbacks] = useState>( + keyData.budget_fallbacks && typeof keyData.budget_fallbacks === "object" ? keyData.budget_fallbacks : {}, + ); const { data: organizations, isLoading: isOrganizationsLoading } = useOrganizations(); const { data: projects } = useProjects(); const { data: uiSettingsData } = useUISettings(); @@ -304,6 +308,8 @@ export function KeyEditView({ values.budget_limits = []; } + values.budget_fallbacks = budgetFallbacks; + await onSubmit(values); } finally { setIsKeySaving(false); @@ -472,6 +478,23 @@ export function KeyEditView({ + + Budget Fallbacks{" "} + + + + + } + > + + + diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 76d6f7ece22..197f87b6dfe 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -748,6 +748,21 @@ export default function KeyInfoView({ + {currentKeyData.budget_fallbacks && Object.keys(currentKeyData.budget_fallbacks).length > 0 && ( +
+ Budget Fallbacks +
+ {Object.entries(currentKeyData.budget_fallbacks).map(([model, fallbacks]) => ( +
+ {model} + -> + {fallbacks.join(", ")} +
+ ))} +
+
+ )} +
Tags