Douwe Osinga
ddd35f6d47
Co-authored-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: jh-block <jhugo@block.xyz> Co-authored-by: Spence <spencermartin@squareup.com> Co-authored-by: Michael Neale <michael.neale@gmail.com>
398 lines
15 KiB
TypeScript
398 lines
15 KiB
TypeScript
import { useState, useEffect, useCallback, useRef } from 'react';
|
|
import { useConfig } from './ConfigContext';
|
|
import {
|
|
listLocalModels,
|
|
downloadHfModel,
|
|
getLocalModelDownloadProgress,
|
|
cancelLocalModelDownload,
|
|
type DownloadProgress,
|
|
type LocalModelResponse,
|
|
} from '../api';
|
|
import { toastService } from '../toasts';
|
|
import { trackOnboardingSetupFailed } from '../utils/analytics';
|
|
import { Goose } from './icons';
|
|
|
|
interface LocalModelSetupProps {
|
|
onSuccess: () => void;
|
|
onCancel: () => void;
|
|
}
|
|
|
|
const formatBytes = (bytes: number): string => {
|
|
if (bytes < 1024) return `${bytes}B`;
|
|
if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(0)}KB`;
|
|
if (bytes < 1024 * 1024 * 1024) return `${(bytes / (1024 * 1024)).toFixed(0)}MB`;
|
|
return `${(bytes / (1024 * 1024 * 1024)).toFixed(1)}GB`;
|
|
};
|
|
|
|
const formatSize = (bytes: number): string => {
|
|
const mb = bytes / (1024 * 1024);
|
|
return mb >= 1024 ? `${(mb / 1024).toFixed(1)}GB` : `${mb.toFixed(0)}MB`;
|
|
};
|
|
|
|
type SetupPhase = 'loading' | 'select' | 'downloading' | 'error';
|
|
|
|
export function LocalModelSetup({ onSuccess, onCancel }: LocalModelSetupProps) {
|
|
const { upsert } = useConfig();
|
|
const [phase, setPhase] = useState<SetupPhase>('loading');
|
|
const [models, setModels] = useState<LocalModelResponse[]>([]);
|
|
const [selectedModelId, setSelectedModelId] = useState<string | null>(null);
|
|
const [downloadProgress, setDownloadProgress] = useState<DownloadProgress | null>(null);
|
|
const [errorMessage, setErrorMessage] = useState<string | null>(null);
|
|
const [showAllModels, setShowAllModels] = useState(false);
|
|
const pollRef = useRef<ReturnType<typeof setInterval> | null>(null);
|
|
|
|
const cleanup = useCallback(() => {
|
|
if (pollRef.current) {
|
|
clearInterval(pollRef.current);
|
|
pollRef.current = null;
|
|
}
|
|
}, []);
|
|
|
|
useEffect(() => cleanup, [cleanup]);
|
|
|
|
useEffect(() => {
|
|
const load = async () => {
|
|
try {
|
|
const response = await listLocalModels();
|
|
if (response.data) {
|
|
setModels(response.data);
|
|
|
|
const alreadyDownloaded = response.data.find((m) => m.status.state === 'Downloaded');
|
|
if (alreadyDownloaded) {
|
|
setSelectedModelId(alreadyDownloaded.id);
|
|
} else {
|
|
const recommended = response.data.find((m: LocalModelResponse) => m.recommended);
|
|
if (recommended) setSelectedModelId(recommended.id);
|
|
}
|
|
}
|
|
} catch (error) {
|
|
console.error('Failed to load local models:', error);
|
|
setErrorMessage('Failed to load available models. Please try again.');
|
|
setPhase('error');
|
|
return;
|
|
}
|
|
setPhase('select');
|
|
};
|
|
load();
|
|
}, []);
|
|
|
|
const finishSetup = async (modelId: string) => {
|
|
await upsert('GOOSE_PROVIDER', 'local', false);
|
|
await upsert('GOOSE_MODEL', modelId, false);
|
|
toastService.success({
|
|
title: 'Local Model Ready',
|
|
msg: `Running entirely on your machine with ${modelId}.`,
|
|
});
|
|
onSuccess();
|
|
};
|
|
|
|
const startDownload = async (modelId: string) => {
|
|
setPhase('downloading');
|
|
setDownloadProgress(null);
|
|
setErrorMessage(null);
|
|
|
|
const model = models.find((m) => m.id === modelId);
|
|
if (!model) {
|
|
setErrorMessage('Model not found');
|
|
setPhase('error');
|
|
return;
|
|
}
|
|
|
|
try {
|
|
await downloadHfModel({ body: { spec: model.id } });
|
|
} catch (error) {
|
|
console.error('Failed to start download:', error);
|
|
setErrorMessage('Failed to start download. Please try again.');
|
|
trackOnboardingSetupFailed('local', 'download_start_failed');
|
|
setPhase('error');
|
|
return;
|
|
}
|
|
|
|
pollRef.current = setInterval(async () => {
|
|
try {
|
|
const response = await getLocalModelDownloadProgress({ path: { model_id: modelId } });
|
|
if (response.data) {
|
|
setDownloadProgress(response.data);
|
|
if (response.data.status === 'completed') {
|
|
cleanup();
|
|
await finishSetup(modelId);
|
|
} else if (response.data.status === 'failed') {
|
|
cleanup();
|
|
setErrorMessage(response.data.error || 'Download failed.');
|
|
trackOnboardingSetupFailed('local', response.data.error || 'download_failed');
|
|
setPhase('error');
|
|
} else if (response.data.status === 'cancelled') {
|
|
cleanup();
|
|
setPhase('select');
|
|
}
|
|
}
|
|
} catch {
|
|
cleanup();
|
|
setErrorMessage('Lost connection to download. Please try again.');
|
|
trackOnboardingSetupFailed('local', 'progress_poll_failed');
|
|
setPhase('error');
|
|
}
|
|
}, 500);
|
|
};
|
|
|
|
const handleCancel = async () => {
|
|
if (phase === 'downloading' && selectedModelId) {
|
|
cleanup();
|
|
try {
|
|
await cancelLocalModelDownload({ path: { model_id: selectedModelId } });
|
|
} catch {
|
|
// best-effort
|
|
}
|
|
setDownloadProgress(null);
|
|
setPhase('select');
|
|
} else {
|
|
onCancel();
|
|
}
|
|
};
|
|
|
|
const handlePrimaryAction = async () => {
|
|
if (!selectedModelId) return;
|
|
const model = models.find((m) => m.id === selectedModelId);
|
|
if (!model) return;
|
|
if (model.status.state === 'Downloaded') {
|
|
await finishSetup(model.id);
|
|
} else {
|
|
await startDownload(model.id);
|
|
}
|
|
};
|
|
|
|
const recommended = models.find((m) => m.recommended);
|
|
const otherModels = models.filter((m) => m.id !== recommended?.id);
|
|
const selectedModel = models.find((m) => m.id === selectedModelId);
|
|
|
|
if (phase === 'loading') {
|
|
return (
|
|
<div className="flex flex-col items-center justify-center py-16">
|
|
<div className="animate-spin rounded-full h-8 w-8 border-t-2 border-b-2 border-text-muted mb-4"></div>
|
|
<p className="text-text-muted text-sm">Checking available models...</p>
|
|
</div>
|
|
);
|
|
}
|
|
|
|
return (
|
|
<div className="space-y-6">
|
|
{/* Header */}
|
|
<div className="text-left space-y-3">
|
|
<div className="origin-bottom-left goose-icon-animation">
|
|
<Goose className="size-6 sm:size-8" />
|
|
</div>
|
|
<h1 className="text-2xl sm:text-4xl font-light">Run Locally</h1>
|
|
<p className="text-text-muted text-base sm:text-lg">
|
|
Download a model to run Goose entirely on your machine — no API keys, no accounts, completely free and private.
|
|
</p>
|
|
</div>
|
|
|
|
{/* Error state */}
|
|
{phase === 'error' && (
|
|
<div className="space-y-4">
|
|
<div className="border border-red-500/30 rounded-xl p-4 bg-red-500/5">
|
|
<p className="text-sm text-red-400">{errorMessage}</p>
|
|
</div>
|
|
<button
|
|
onClick={() => {
|
|
setErrorMessage(null);
|
|
setPhase('select');
|
|
}}
|
|
className="w-full px-6 py-3 bg-background-muted text-text-default rounded-lg transition-colors font-medium hover:bg-background-muted/80"
|
|
>
|
|
Try Again
|
|
</button>
|
|
<button
|
|
onClick={onCancel}
|
|
className="w-full px-6 py-3 bg-transparent text-text-muted rounded-lg hover:bg-background-muted transition-colors"
|
|
>
|
|
Back
|
|
</button>
|
|
</div>
|
|
)}
|
|
|
|
{/* Model selection */}
|
|
{phase === 'select' && (
|
|
<div className="space-y-5">
|
|
{/* Recommended model card */}
|
|
{recommended && (
|
|
<div
|
|
onClick={() => setSelectedModelId(recommended.id)}
|
|
className={`relative w-full p-4 sm:p-6 border rounded-xl cursor-pointer transition-all duration-200 group ${
|
|
selectedModelId === recommended.id
|
|
? 'border-blue-500 bg-blue-500/5'
|
|
: 'border-border-subtle hover:border-border-default'
|
|
}`}
|
|
>
|
|
<div className="absolute -top-2 -right-2 sm:-top-3 sm:-right-3 z-10">
|
|
<span className="inline-block px-2 py-1 text-xs font-medium bg-blue-600 text-white rounded-full">
|
|
Best for your machine
|
|
</span>
|
|
</div>
|
|
<div className="flex items-start gap-3">
|
|
<input
|
|
type="radio"
|
|
checked={selectedModelId === recommended.id}
|
|
onChange={() => setSelectedModelId(recommended.id)}
|
|
className="cursor-pointer flex-shrink-0 mt-1"
|
|
/>
|
|
<div className="flex-1 min-w-0">
|
|
<div className="flex items-center gap-2 flex-wrap">
|
|
<span className="font-medium text-text-default text-sm sm:text-base">
|
|
{recommended.display_name}
|
|
</span>
|
|
{recommended.status.state === 'Downloaded' && (
|
|
<span className="text-xs bg-green-600 text-white px-2 py-0.5 rounded-full">
|
|
Ready
|
|
</span>
|
|
)}
|
|
</div>
|
|
<p className="text-text-muted text-xs mt-1">
|
|
{formatSize(recommended.size_bytes)}
|
|
</p>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
)}
|
|
|
|
{/* Expandable other models */}
|
|
{otherModels.length > 0 && (
|
|
<div>
|
|
<button
|
|
onClick={() => setShowAllModels(!showAllModels)}
|
|
className="text-sm text-blue-500 hover:text-blue-400 transition-colors flex items-center gap-1"
|
|
>
|
|
{showAllModels ? 'Hide other sizes' : `Show ${otherModels.length} other sizes`}
|
|
<svg
|
|
className={`w-3.5 h-3.5 transition-transform ${showAllModels ? 'rotate-180' : ''}`}
|
|
fill="none"
|
|
stroke="currentColor"
|
|
viewBox="0 0 24 24"
|
|
>
|
|
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M19 9l-7 7-7-7" />
|
|
</svg>
|
|
</button>
|
|
|
|
{showAllModels && (
|
|
<div className="mt-3 space-y-2">
|
|
{otherModels.map((model) => (
|
|
<div
|
|
key={model.id}
|
|
onClick={() => setSelectedModelId(model.id)}
|
|
className={`w-full p-4 border rounded-xl cursor-pointer transition-all duration-200 ${
|
|
selectedModelId === model.id
|
|
? 'border-blue-500 bg-blue-500/5'
|
|
: 'border-border-subtle hover:border-border-default'
|
|
}`}
|
|
>
|
|
<div className="flex items-start gap-3">
|
|
<input
|
|
type="radio"
|
|
checked={selectedModelId === model.id}
|
|
onChange={() => setSelectedModelId(model.id)}
|
|
className="cursor-pointer flex-shrink-0 mt-0.5"
|
|
/>
|
|
<div className="flex-1 min-w-0">
|
|
<div className="flex items-center gap-2 flex-wrap">
|
|
<span className="font-medium text-text-default text-sm">{model.display_name}</span>
|
|
<span className="text-xs text-text-muted">{formatSize(model.size_bytes)}</span>
|
|
{model.status.state === 'Downloaded' && (
|
|
<span className="text-xs bg-green-600 text-white px-2 py-0.5 rounded-full">
|
|
Ready
|
|
</span>
|
|
)}
|
|
</div>
|
|
</div>
|
|
</div>
|
|
</div>
|
|
))}
|
|
</div>
|
|
)}
|
|
</div>
|
|
)}
|
|
|
|
{/* Primary action */}
|
|
<button
|
|
onClick={handlePrimaryAction}
|
|
disabled={!selectedModelId}
|
|
className="w-full px-6 py-3 bg-background-muted text-text-default rounded-lg transition-colors font-medium disabled:opacity-40 disabled:cursor-not-allowed hover:bg-background-muted/80"
|
|
>
|
|
{selectedModel?.status.state === 'Downloaded'
|
|
? `Use ${selectedModel.display_name}`
|
|
: selectedModel
|
|
? `Download ${selectedModel.display_name} (${formatSize(selectedModel.size_bytes)})`
|
|
: 'Select a model'}
|
|
</button>
|
|
|
|
<button
|
|
onClick={onCancel}
|
|
className="w-full px-6 py-3 bg-transparent text-text-muted rounded-lg hover:bg-background-muted transition-colors"
|
|
>
|
|
Back
|
|
</button>
|
|
</div>
|
|
)}
|
|
|
|
{/* Downloading state */}
|
|
{phase === 'downloading' && selectedModel && (
|
|
<div className="space-y-6">
|
|
<div className="border border-border-subtle rounded-xl p-5 sm:p-6 bg-background-default">
|
|
<p className="font-medium text-text-default text-sm sm:text-base mb-4">
|
|
Downloading {selectedModel.display_name}
|
|
</p>
|
|
|
|
{downloadProgress ? (
|
|
<div className="space-y-3">
|
|
{/* Progress bar */}
|
|
<div className="w-full bg-background-subtle rounded-full h-2 overflow-hidden">
|
|
<div
|
|
className="bg-blue-500 h-2 rounded-full transition-all duration-500 ease-out"
|
|
style={{ width: `${downloadProgress.progress_percent}%` }}
|
|
/>
|
|
</div>
|
|
|
|
{/* Stats row */}
|
|
<div className="flex justify-between text-xs text-text-muted">
|
|
<span>
|
|
{formatBytes(downloadProgress.bytes_downloaded)} of{' '}
|
|
{formatBytes(downloadProgress.total_bytes)}
|
|
</span>
|
|
<span>{downloadProgress.progress_percent.toFixed(0)}%</span>
|
|
</div>
|
|
|
|
<div className="flex justify-between text-xs text-text-muted">
|
|
{downloadProgress.speed_bps ? (
|
|
<span>{formatBytes(downloadProgress.speed_bps)}/s</span>
|
|
) : (
|
|
<span />
|
|
)}
|
|
{downloadProgress.eta_seconds != null && downloadProgress.eta_seconds > 0 && (
|
|
<span>
|
|
~{downloadProgress.eta_seconds < 60
|
|
? `${Math.round(downloadProgress.eta_seconds)}s`
|
|
: `${Math.round(downloadProgress.eta_seconds / 60)}m`}{' '}
|
|
remaining
|
|
</span>
|
|
)}
|
|
</div>
|
|
</div>
|
|
) : (
|
|
<div className="flex items-center gap-3">
|
|
<div className="animate-spin rounded-full h-4 w-4 border-t-2 border-b-2 border-text-muted"></div>
|
|
<span className="text-sm text-text-muted">Starting download...</span>
|
|
</div>
|
|
)}
|
|
</div>
|
|
|
|
<button
|
|
onClick={handleCancel}
|
|
className="w-full px-6 py-3 bg-transparent text-text-muted rounded-lg hover:bg-background-muted transition-colors border border-border-subtle"
|
|
>
|
|
Cancel Download
|
|
</button>
|
|
</div>
|
|
)}
|
|
</div>
|
|
);
|
|
}
|