103 lines
3.2 KiB
TypeScript
103 lines
3.2 KiB
TypeScript
import { useState, useMemo } from 'react';
|
|
import { useLocalModels, LocalModel } from '@/hooks/use-local-models';
|
|
import { MODEL_REGISTRY, ModelDefinition } from '@/lib/model-registry';
|
|
|
|
const STORAGE_KEY = 'last_used_model';
|
|
|
|
export interface UseModelSelectionResult {
|
|
myModels: ModelDefinition[];
|
|
availableModels: ModelDefinition[];
|
|
selectedModel: string | null;
|
|
setSelectedModel: (model: string) => void;
|
|
isLoading: boolean;
|
|
installedVariantsByFamily: Record<string, string[]>;
|
|
refreshModels: () => Promise<void>;
|
|
}
|
|
|
|
export function useModelSelection(): UseModelSelectionResult {
|
|
const { models: installedModels, loading: isLoadingModels, refresh: refreshLocalModels } = useLocalModels();
|
|
const [selectedModel, setSelectedModelState] = useState<string | null>(null);
|
|
|
|
const { myModels, availableModels } = useMemo(() => {
|
|
const installedNames = new Set(installedModels.map(m => m.name));
|
|
const my: ModelDefinition[] = [];
|
|
const available: ModelDefinition[] = [];
|
|
const registryByModel = new Map<string, ModelDefinition>();
|
|
for (const def of MODEL_REGISTRY) {
|
|
registryByModel.set(def.model, def);
|
|
}
|
|
|
|
for (const def of MODEL_REGISTRY) {
|
|
const isInstalled = def.variants.some(v => installedNames.has(v.name));
|
|
if (isInstalled) {
|
|
my.push(def);
|
|
} else {
|
|
available.push(def);
|
|
}
|
|
}
|
|
|
|
const installedByBase = new Map<string, LocalModel[]>();
|
|
for (const m of installedModels) {
|
|
const base = m.name.split(':')[0];
|
|
const list = installedByBase.get(base) || [];
|
|
list.push(m);
|
|
installedByBase.set(base, list);
|
|
}
|
|
|
|
for (const [base, models] of installedByBase.entries()) {
|
|
if (!registryByModel.has(base)) {
|
|
const variants = models.map((m) => ({
|
|
name: m.name,
|
|
size: m.size,
|
|
context: 'Unknown',
|
|
inputType: 'Text',
|
|
}));
|
|
my.push({
|
|
model: base,
|
|
url: '',
|
|
description: 'Modelo instalado localmente',
|
|
variantsCount: variants.length,
|
|
variants,
|
|
});
|
|
}
|
|
}
|
|
|
|
return { myModels: my, availableModels: available };
|
|
}, [installedModels]);
|
|
|
|
const installedVariantsByFamily = useMemo(() => {
|
|
const map: Record<string, string[]> = {};
|
|
for (const m of installedModels) {
|
|
const base = m.name.split(':')[0];
|
|
if (!map[base]) map[base] = [];
|
|
map[base].push(m.name);
|
|
}
|
|
return map;
|
|
}, [installedModels]);
|
|
|
|
const defaultSelectedModel = useMemo(() => {
|
|
if (isLoadingModels) return null;
|
|
const lastUsed = typeof window !== 'undefined' ? localStorage.getItem(STORAGE_KEY) : null;
|
|
const installedNames = installedModels.map(m => m.name);
|
|
if (lastUsed && installedNames.includes(lastUsed)) return lastUsed;
|
|
if (installedModels.length > 0) return installedModels[0].name;
|
|
return null;
|
|
}, [installedModels, isLoadingModels]);
|
|
|
|
// Persistence wrapper
|
|
const setSelectedModel = (model: string) => {
|
|
setSelectedModelState(model);
|
|
localStorage.setItem(STORAGE_KEY, model);
|
|
};
|
|
|
|
return {
|
|
myModels,
|
|
availableModels,
|
|
selectedModel: selectedModel ?? defaultSelectedModel,
|
|
setSelectedModel,
|
|
isLoading: isLoadingModels,
|
|
installedVariantsByFamily,
|
|
refreshModels: refreshLocalModels,
|
|
};
|
|
}
|