'use client'; import React, { useState, useEffect, useCallback } from 'react'; import { Button } from '@/components/ui/button'; import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle, } from '@/components/ui/dialog'; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue, } from '@/components/ui/select'; import { AdminService, AdminModel } from '@/lib/api/services/admin'; import { toast } from 'sonner'; import { Download, AlertCircle, CheckCircle, Loader2 } from 'lucide-react'; import { Alert, AlertDescription } from '@/components/ui/alert'; import { Checkbox } from '@/components/ui/checkbox'; import { ScrollArea } from '@/components/ui/scroll-area'; import { Badge } from '@/components/ui/badge'; interface CollectModelsDialogProps { isOpen: boolean; onClose: () => void; onSuccess?: () => void; } export function CollectModelsDialog({ isOpen, onClose, onSuccess, }: CollectModelsDialogProps) { const [selectedProvider, setSelectedProvider] = useState(null); const [providers, setProviders] = useState< Array<{ id: number; provider_type: string; base_url: string }> >([]); const [remoteModels, setRemoteModels] = useState([]); const [selectedModels, setSelectedModels] = useState>(new Set()); const [isLoadingProviders, setIsLoadingProviders] = useState(false); const [isLoadingModels, setIsLoadingModels] = useState(false); const [isSubmitting, setIsSubmitting] = useState(false); const [result, setResult] = useState<{ added: number; skipped: number; errors: string[]; } | null>(null); const loadRemoteModels = useCallback(async () => { if (!selectedProvider) return; setIsLoadingModels(true); setRemoteModels([]); setSelectedModels(new Set()); try { const data = await AdminService.getProviderModels(selectedProvider); const dbModelIds = new Set(data.db_models.map((m) => m.id)); const availableRemoteModels = data.remote_models.filter( (m: AdminModel) => !dbModelIds.has(m.id) ); setRemoteModels(availableRemoteModels); } catch { toast.error('Failed to fetch models from provider'); } finally { setIsLoadingModels(false); } }, [selectedProvider]); useEffect(() => { if (isOpen) { loadProviders(); } }, [isOpen]); useEffect(() => { if (selectedProvider) { loadRemoteModels(); } }, [selectedProvider, loadRemoteModels]); const loadProviders = async () => { setIsLoadingProviders(true); try { const data = await AdminService.getUpstreamProviders(); setProviders(data); } catch { toast.error('Failed to load providers'); } finally { setIsLoadingProviders(false); } }; const toggleModel = (modelId: string) => { const newSelected = new Set(selectedModels); if (newSelected.has(modelId)) { newSelected.delete(modelId); } else { newSelected.add(modelId); } setSelectedModels(newSelected); }; const selectAll = () => { setSelectedModels(new Set(remoteModels.map((m) => m.id))); }; const deselectAll = () => { setSelectedModels(new Set()); }; const handleCollect = async () => { if (selectedModels.size === 0) { toast.error('Please select at least one model'); return; } setIsSubmitting(true); let added = 0; let skipped = 0; const errors: string[] = []; try { for (const modelId of Array.from(selectedModels)) { const model = remoteModels.find((m) => m.id === modelId); if (!model) continue; try { await AdminService.createModel({ ...model, upstream_provider_id: selectedProvider, enabled: true, created: Math.floor(Date.now() / 1000), }); added++; } catch (err) { const errorMessage = err instanceof Error ? err.message : 'Unknown error'; if (errorMessage.includes('already exists')) { skipped++; } else { errors.push(`${modelId}: ${errorMessage}`); } } } setResult({ added, skipped, errors }); if (added > 0) { toast.success(`Successfully added ${added} models`); onSuccess?.(); } if (errors.length > 0) { toast.warning(`Completed with ${errors.length} errors`); } } catch { toast.error('Failed to collect models'); } finally { setIsSubmitting(false); } }; const handleClose = () => { if (!isSubmitting) { setSelectedProvider(null); setRemoteModels([]); setSelectedModels(new Set()); setResult(null); onClose(); } }; return ( Collect Models from Provider Fetch models from an upstream provider and add them to your database
{isLoadingModels && (
Loading models from provider...
)} {!isLoadingModels && selectedProvider && remoteModels.length > 0 && ( <>
{remoteModels.length} models available
{remoteModels.map((model) => (
toggleModel(model.id)} disabled={isSubmitting} />
{model.id}
{model.description || model.name}
{model.context_length && ( {model.context_length.toLocaleString()} tokens )}
))}
)} {!isLoadingModels && selectedProvider && remoteModels.length === 0 && ( No new models available from this provider. All models may already be in your database. )} {result && (
Collection Results:
• {result.added} models added
• {result.skipped} models skipped {result.errors.length > 0 && ( <>
• {result.errors.length} errors occurred )}
{result.errors.length > 0 && ( Errors:
    {result.errors.slice(0, 3).map((error, index) => (
  • {error}
  • ))} {result.errors.length > 3 && (
  • ... and {result.errors.length - 3} more errors
  • )}
)}
)}
); }