Files
OllaHub/components/chat/model-selector.tsx
2026-02-22 14:06:01 -03:00

318 lines
16 KiB
TypeScript

'use client';
import * as React from 'react';
import { Check, ChevronsUpDown, Search, DownloadCloud, Cpu, Loader2 } from 'lucide-react';
import { cn } from '@/lib/utils';
import { Button } from '@/components/ui/button';
import {
Popover,
PopoverContent,
PopoverTrigger,
} from '@/components/ui/popover';
import { Tabs, TabsContent, TabsList, TabsTrigger } from '@/components/ui/tabs';
import { useModelSelection } from '@/hooks/use-model-selection';
import { useDownloadContext } from '@/contexts/download-context';
import { Badge } from '@/components/ui/badge';
import { useToast } from '@/hooks/use-toast';
export function ModelSelector() {
const [open, setOpen] = React.useState(false);
const [searchQuery, setSearchQuery] = React.useState('');
const { myModels, availableModels, selectedModel, setSelectedModel, installedVariantsByFamily, refreshModels } = useModelSelection();
const { startDownload, state: downloadState } = useDownloadContext();
const { toast } = useToast();
// Track selected variant for each model in the catalog
const [selectedVariants, setSelectedVariants] = React.useState<Record<string, string>>({});
const handleDownload = async (e: React.MouseEvent, modelName: string) => {
e.stopPropagation();
try {
await startDownload(modelName);
toast({
title: "Download Iniciado",
description: `Baixando ${modelName}...`,
});
} catch (error) {
toast({
title: "Erro",
description: "Falha ao iniciar download.",
variant: "destructive",
});
}
};
const handleSmartDownload = async () => {
const term = searchQuery.trim();
if (!term) return;
try {
await startDownload(term);
setOpen(false);
setSearchQuery("");
toast({
title: "Download Iniciado",
description: `Baixando ${term} da biblioteca Ollama...`,
});
} catch (error) {
toast({
title: "Erro",
description: "Falha ao iniciar download.",
variant: "destructive",
});
}
};
React.useEffect(() => {
if (downloadState?.isSuccess) {
refreshModels();
}
}, [downloadState?.isSuccess, refreshModels]);
const filteredMyModels = React.useMemo(() => {
if (!searchQuery) return myModels;
return myModels.filter(m =>
m.model.toLowerCase().includes(searchQuery.toLowerCase())
);
}, [myModels, searchQuery]);
const filteredAvailableModels = React.useMemo(() => {
if (!searchQuery) return availableModels;
return availableModels.filter(m =>
m.model.toLowerCase().includes(searchQuery.toLowerCase())
);
}, [availableModels, searchQuery]);
// Helper to check if a model is currently selected (partial match for variants)
const isSelected = (modelName: string) => {
if (!selectedModel) return false;
return selectedModel.startsWith(modelName) || selectedModel === modelName;
};
// Helper to extract tag/parameters from model name
const getVariantDisplay = (name: string) => {
if (name.includes(':')) {
return name.split(':')[1];
}
return 'latest';
};
return (
<Popover open={open} onOpenChange={setOpen}>
<PopoverTrigger asChild>
<Button
variant="outline"
role="combobox"
aria-expanded={open}
className="w-[260px] justify-between"
>
{selectedModel ? (
<span className="flex items-center gap-2 truncate">
<Cpu className="w-4 h-4 text-muted-foreground shrink-0" />
<span className="truncate">{selectedModel}</span>
</span>
) : (
"Selecione um modelo..."
)}
<ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
</Button>
</PopoverTrigger>
<PopoverContent className="w-[400px] p-0" align="start">
<div className="flex flex-col">
{/* Search Input */}
<div className="flex items-center border-b border-border/50 px-3 bg-muted/30">
<Search className="mr-2 h-4 w-4 shrink-0 opacity-50" />
<input
className="flex h-11 w-full rounded-md bg-transparent py-3 text-sm outline-none placeholder:text-muted-foreground disabled:cursor-not-allowed disabled:opacity-50"
placeholder="Buscar modelos..."
value={searchQuery}
onChange={(e) => setSearchQuery(e.target.value)}
onKeyDown={(e) => {
if (e.key === 'Enter' && searchQuery && filteredAvailableModels.length === 0) {
handleSmartDownload();
}
}}
/>
</div>
<Tabs defaultValue="installed" className="w-full">
<TabsList className="w-full rounded-none border-b bg-transparent p-0">
<TabsTrigger
value="installed"
className="flex-1 rounded-none border-b-2 border-transparent data-[state=active]:border-primary data-[state=active]:bg-transparent h-10"
>
Instalados
</TabsTrigger>
<TabsTrigger
value="available"
className="flex-1 rounded-none border-b-2 border-transparent data-[state=active]:border-primary data-[state=active]:bg-transparent h-10"
>
Catálogo
</TabsTrigger>
</TabsList>
<TabsContent value="installed" className="m-0 max-h-[400px] overflow-y-auto p-2">
{filteredMyModels.length === 0 ? (
<div className="py-8 text-center text-sm text-muted-foreground">
{searchQuery ? (
<span>Nenhum modelo instalado encontrado.</span>
) : (
<span>Nenhum modelo instalado.</span>
)}
</div>
) : (
<div className="flex flex-col gap-2">
{filteredMyModels.map((modelDef) => {
const family = modelDef.model;
const installedTags = installedVariantsByFamily[family] || [];
const isFamilySelected = selectedModel ? (selectedModel === family || selectedModel.startsWith(family + ":")) : false;
return (
<div
key={family}
className="relative flex flex-col rounded-md border border-transparent hover:border-border/50 bg-card/50 hover:bg-accent/30 p-3 text-sm transition-all"
>
<div className="flex items-start justify-between w-full mb-2">
<div className="flex flex-col gap-0.5 overflow-hidden pr-2">
<span className="font-bold text-base truncate flex items-center gap-2">
{modelDef.model}
{isFamilySelected && (
<Check className="h-3 w-3 text-primary animate-in fade-in zoom-in" />
)}
</span>
<span className="text-xs text-muted-foreground line-clamp-2" title={modelDef.description}>
{modelDef.description}
</span>
</div>
</div>
<div className="flex flex-wrap gap-1.5 mt-1">
{installedTags.map((fullTag) => {
const tag = getVariantDisplay(fullTag);
const isExactSelected = selectedModel === fullTag;
// Find variant def for size info if available
const variantDef = modelDef.variants.find(v => v.name === fullTag);
return (
<Badge
key={fullTag}
variant="outline"
className={cn(
"cursor-pointer text-[10px] px-2 h-5 transition-all border-transparent",
isExactSelected
? "bg-foreground text-background hover:bg-foreground/90"
: "bg-muted text-muted-foreground hover:bg-muted/80 hover:text-foreground"
)}
title={variantDef ? `Tamanho: ${variantDef.size}` : undefined}
onClick={(e) => {
e.stopPropagation();
setSelectedModel(fullTag);
setOpen(false);
}}
>
{tag}
</Badge>
);
})}
</div>
</div>
);
})}
</div>
)}
</TabsContent>
<TabsContent value="available" className="m-0 max-h-[400px] overflow-y-auto p-2">
{filteredAvailableModels.length === 0 ? (
<div className="py-6 text-center text-sm">
{searchQuery ? (
<div className="flex flex-col items-center gap-2">
<p className="text-muted-foreground">Modelo não encontrado no catálogo.</p>
<Button
variant="outline"
className="h-auto py-2 px-4 border-dashed border-zinc-500/50 hover:border-primary hover:bg-primary/10"
onClick={handleSmartDownload}
>
<DownloadCloud className="mr-2 h-4 w-4" />
<span>Baixar <span className="font-bold ml-1">{searchQuery}</span></span>
</Button>
</div>
) : (
<span className="text-muted-foreground">Nenhum modelo encontrado.</span>
)}
</div>
) : (
<div className="flex flex-col gap-2">
{filteredAvailableModels.map((modelDef) => {
// Determine active variant for this item
const activeVariantName = selectedVariants[modelDef.model] || modelDef.variants[0]?.name;
const activeVariant = modelDef.variants.find(v => v.name === activeVariantName) || modelDef.variants[0];
return (
<div
key={modelDef.model}
className="relative flex flex-col rounded-md border border-transparent hover:border-border/50 bg-card/50 hover:bg-accent/30 p-3 text-sm transition-all"
>
<div className="flex items-start justify-between w-full mb-2">
<div className="flex flex-col gap-0.5 overflow-hidden pr-2">
<span className="font-bold text-base truncate">{modelDef.model}</span>
<span className="text-xs text-muted-foreground line-clamp-2" title={modelDef.description}>
{modelDef.description}
</span>
</div>
<Button
size="icon"
variant="ghost"
className="h-8 w-8 shrink-0 text-muted-foreground hover:text-foreground"
onClick={(e) => handleDownload(e, activeVariant.name)}
disabled={downloadState.isDownloading && downloadState.currentModel === activeVariant.name}
title={`Baixar ${activeVariant.name}`}
>
{downloadState.isDownloading && downloadState.currentModel === activeVariant.name ? (
<Loader2 className="h-4 w-4 animate-spin" />
) : (
<DownloadCloud className="h-4 w-4" />
)}
</Button>
</div>
{/* Variants Pills */}
<div className="flex flex-wrap gap-1.5 mt-1">
{modelDef.variants.map((variant) => {
const isVariantSelected = activeVariant.name === variant.name;
return (
<Badge
key={variant.name}
variant="outline"
className={cn(
"cursor-pointer text-[10px] px-2 h-5 transition-all border-transparent",
isVariantSelected
? "bg-foreground text-background hover:bg-foreground/90"
: "bg-muted text-muted-foreground hover:bg-muted/80 hover:text-foreground"
)}
title={`Tamanho em disco: ${variant.size}`}
onClick={(e) => {
e.stopPropagation();
setSelectedVariants(prev => ({
...prev,
[modelDef.model]: variant.name
}));
}}
>
{getVariantDisplay(variant.name)}
</Badge>
);
})}
</div>
</div>
)})}
</div>
)}
</TabsContent>
</Tabs>
</div>
</PopoverContent>
</Popover>
);
}