diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 34443ec87f..1184183088 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -5,6 +5,7 @@ import { } from "@/components/assistant-ui/model-selector"; import { Thread } from "@/components/assistant-ui/thread"; import { Button } from "@/components/ui/button"; +import { Spinner } from "@/components/ui/spinner"; import { SidebarProvider, SidebarTrigger, useSidebar } from "@/components/ui/sidebar"; import { Sheet, @@ -283,7 +284,8 @@ export function ChatPage(): ReactElement { const modelsFromStore = useChatRuntimeStore((state) => state.models); const lorasFromStore = useChatRuntimeStore((state) => state.loras); const modelsError = useChatRuntimeStore((state) => state.modelsError); - const { refresh, selectModel, ejectModel } = useChatModelRuntime(); + const { refresh, selectModel, ejectModel, loadingModel } = + useChatModelRuntime(); const refreshRef = useRef(refresh); const selectModelRef = useRef(selectModel); @@ -518,6 +520,17 @@ export function ChatPage(): ReactElement { contentDataTour="chat-model-selector-popover" className="max-w-[62vw] sm:max-w-none" /> + {loadingModel ? ( +
+ + + Downloading model… + +
+ ) : null} {modelsError && (
diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index 03aba01736..175ccfba65 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -1,4 +1,4 @@ -import { useCallback } from "react"; +import { useCallback, useState } from "react"; import { toast } from "sonner"; import { getInferenceStatus, @@ -116,6 +116,11 @@ export function useChatModelRuntime() { const setCheckpoint = useChatRuntimeStore((state) => state.setCheckpoint); const clearCheckpoint = useChatRuntimeStore((state) => state.clearCheckpoint); + const [loadingModel, setLoadingModel] = useState<{ + id: string; + displayName: string; + } | null>(null); + const refresh = useCallback(async () => { setModelsError(null); try { @@ -157,6 +162,7 @@ export function useChatModelRuntime() { const displayName = model?.name || lora?.name || modelId; setModelsError(null); + setLoadingModel({ id: modelId, displayName }); try { async function performLoad(): Promise { if (params.checkpoint) { @@ -176,19 +182,20 @@ export function useChatModelRuntime() { await refresh(); } - let description = "Base model selected."; - if (isLora) { - description = "Fine-tuned (LoRA) selected."; - } + const loadPromise = performLoad().finally(() => { + setLoadingModel(null); + }); - await toast.promise(performLoad(), { - loading: `Loading ${displayName}`, + await toast.promise(loadPromise, { + loading: "Downloading model…", success: `${displayName} loaded`, error: (err) => err instanceof Error ? err.message : "Failed to load model", - description, + description: + "This may include downloading. Large models can take a while.", }); } catch (error) { + setLoadingModel(null); const message = error instanceof Error ? error.message : "Failed to load model"; setModelsError(message); @@ -227,5 +234,6 @@ export function useChatModelRuntime() { refresh, selectModel, ejectModel, + loadingModel, }; }