From 65203b29208bbbb0c079a9ea9f1a22fa6ffc830c Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 6 Aug 2026 14:56:57 -0700 Subject: [PATCH] fix(ui): load playground models for virtual keys via /v1/models Prefer the key-scoped OpenAI models list when the Virtual Key source is selected, enrich with mode from model_group/info, and debounce custom key input so models appear for the key's access set --- .../playground/components/chat_ui/ChatUI.tsx | 53 +++++++++--- .../components/llm_calls/fetch_models.test.ts | 77 +++++++++++++++++ .../src/components/llm_calls/fetch_models.tsx | 86 ++++++++++++++----- 3 files changed, 184 insertions(+), 32 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx index c89600b1efc..d1924d8c38a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx @@ -181,6 +181,9 @@ const ChatUI: React.FC = ({ return disabledPersonalKeyCreation ? "custom" : "session"; }); const [apiKey, setApiKey] = useState(() => getSecureItem("apiKey") || ""); + const [debouncedCustomApiKey, setDebouncedCustomApiKey] = useState(() => + (getSecureItem("apiKey") || "").trim(), + ); const [customProxyBaseUrl, setCustomProxyBaseUrl] = useState( () => sessionStorage.getItem("customProxyBaseUrl") || "", ); @@ -401,10 +404,21 @@ const ChatUI: React.FC = ({ ]); useEffect(() => { - const userApiKey = apiKeySource === "session" ? accessToken : apiKey.trim(); + if (apiKeySource === "session") { + return; + } + const timeoutId = window.setTimeout(() => { + setDebouncedCustomApiKey(apiKey.trim()); + }, 300); + return () => window.clearTimeout(timeoutId); + }, [apiKey, apiKeySource]); + + useEffect(() => { + const userApiKey = apiKeySource === "session" ? accessToken : debouncedCustomApiKey; if (!userApiKey) { setModelInfo([]); setModelLoadError(false); + setIsLoadingModels(false); return; } @@ -442,12 +456,12 @@ const ChatUI: React.FC = ({ if (!simplified) { void loadModels(); } - void loadMCPServers(); + void loadMCPServers(userApiKey); return () => { cancelled = true; }; - }, [accessToken, apiKeySource, apiKey, simplified]); + }, [accessToken, apiKeySource, debouncedCustomApiKey, simplified]); // Load tools when MCP direct mode has a server (or toolset) selected useEffect(() => { @@ -1199,15 +1213,30 @@ const ChatUI: React.FC = ({ {apiKeySource === "custom" && ( -
- - setApiKey(event.target.value)} - value={apiKey} - /> +
+
+ + setApiKey(event.target.value)} + onBlur={() => setDebouncedCustomApiKey(apiKey.trim())} + value={apiKey} + aria-label="Virtual Key" + /> +
+ {isLoadingModels && apiKey.trim() !== "" && ( +

Loading models for this key...

+ )} + {!isLoadingModels && apiKey.trim() !== "" && modelLoadError && ( +

Unable to load models for this Virtual Key.

+ )} + {!isLoadingModels && apiKey.trim() !== "" && !modelLoadError && modelInfo.length === 0 && ( +

No models available for this Virtual Key.

+ )}
)}
diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.ts b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.ts new file mode 100644 index 00000000000..f08143b52f2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.test.ts @@ -0,0 +1,77 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { fetchAvailableModels } from "./fetch_models"; + +vi.mock("@/components/networking", () => ({ + apiClient: { + get: vi.fn(), + }, +})); + +import { apiClient } from "@/components/networking"; + +const mockGet = vi.mocked(apiClient.get); + +describe("fetchAvailableModels", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("returns key-scoped models from /v1/models and attaches modes from model_group/info", async () => { + mockGet.mockImplementation(async (path: string) => { + if (path === "/model_group/info") { + return { + data: [ + { model_group: "anthropic-haiku-4-5", mode: "chat" }, + { model_group: "gpt-realtime", mode: "realtime" }, + ], + }; + } + if (path === "/v1/models") { + return { + data: [{ id: "anthropic-haiku-4-5" }], + }; + } + throw new Error(`unexpected path ${path}`); + }); + + const models = await fetchAvailableModels("sk-virtual-key"); + + expect(models).toEqual([{ model_group: "anthropic-haiku-4-5", mode: "chat" }]); + expect(mockGet).toHaveBeenCalledWith("/v1/models", { accessToken: "sk-virtual-key" }); + }); + + it("falls back to model_group/info when /v1/models is empty", async () => { + mockGet.mockImplementation(async (path: string) => { + if (path === "/model_group/info") { + return { + data: [ + { model_group: "gpt-4o", mode: "chat" }, + { id: "legacy-model", mode: "chat" }, + ], + }; + } + if (path === "/v1/models") { + return { data: [] }; + } + throw new Error(`unexpected path ${path}`); + }); + + const models = await fetchAvailableModels("sk-admin"); + expect(models.map((m) => m.model_group)).toEqual(["gpt-4o", "legacy-model"]); + }); + + it("still returns /v1/models when model_group/info fails", async () => { + mockGet.mockImplementation(async (path: string) => { + if (path === "/model_group/info") { + throw new Error("forbidden"); + } + if (path === "/v1/models") { + return { data: [{ id: "only-from-key" }] }; + } + throw new Error(`unexpected path ${path}`); + }); + + const models = await fetchAvailableModels("sk-virtual-key"); + expect(models).toEqual([{ model_group: "only-from-key", mode: undefined }]); + }); +}); diff --git a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx index 24f1e038f85..cc74355539f 100644 --- a/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx +++ b/ui/litellm-dashboard/src/components/llm_calls/fetch_models.tsx @@ -1,40 +1,86 @@ -// fetch_models.ts - -import { modelHubCall } from "@/components/networking"; +import { apiClient } from "@/components/networking"; export interface ModelGroup { model_group: string; mode?: string; } -interface AvailableModel { +interface ModelGroupInfoItem { model_group?: string | null; model_name?: string | null; id?: string | null; mode?: string | null; } +interface OpenAIModelItem { + id?: string | null; +} + +const modelNameFromGroupItem = (item: ModelGroupInfoItem): string => + (item.model_group || item.id || item.model_name || "").trim(); + +const dedupeAndSort = (models: ModelGroup[]): ModelGroup[] => { + const unique = Array.from(new Map(models.map((model) => [model.model_group, model])).values()); + unique.sort((a, b) => a.model_group.localeCompare(b.model_group)); + return unique; +}; + /** - * Fetches available models using modelHubCall and formats them for the selection dropdown. + * Loads models available to the given key. + * + * Prefers OpenAI-compatible `/v1/models` (scoped to the key's access) and enriches + * entries with `mode` from `/model_group/info` when that endpoint is available. + * Falls back to `/model_group/info` alone if `/v1/models` is empty or fails. */ export const fetchAvailableModels = async (accessToken: string): Promise => { + const modeByName = new Map(); + try { - const fetchedModels = await modelHubCall(accessToken); - - if (fetchedModels?.data.length > 0) { - const models: ModelGroup[] = fetchedModels.data - .map((item: AvailableModel) => ({ - model_group: item.model_group || item.id || item.model_name || "", - mode: item.mode || undefined, - })) - .filter((model: ModelGroup) => model.model_group !== ""); - - models.sort((a, b) => a.model_group.localeCompare(b.model_group)); - return Array.from(new Map(models.map((model) => [model.model_group, model])).values()); + const groupInfo = await apiClient.get<{ data?: ModelGroupInfoItem[] }>("/model_group/info", { + accessToken, + }); + for (const item of groupInfo?.data ?? []) { + const name = modelNameFromGroupItem(item); + if (name) { + modeByName.set(name, item.mode || undefined); + } } - return []; } catch (error) { - console.error("Error fetching model info:", error); - throw error; + console.error("Error fetching model group info:", error); } + + try { + const listed = await apiClient.get<{ data?: OpenAIModelItem[] }>("/v1/models", { + accessToken, + }); + const fromList = (listed?.data ?? []) + .map((item) => { + const name = (item.id || "").trim(); + if (!name) { + return null; + } + return { + model_group: name, + mode: modeByName.get(name), + } satisfies ModelGroup; + }) + .filter((model): model is ModelGroup => model != null); + + if (fromList.length > 0) { + return dedupeAndSort(fromList); + } + } catch (error) { + console.error("Error fetching /v1/models:", error); + } + + if (modeByName.size > 0) { + return dedupeAndSort( + Array.from(modeByName.entries()).map(([model_group, mode]) => ({ + model_group, + mode, + })), + ); + } + + return []; };