/** * Model registry backed by list.modelgrid.com. */ import * as fs from 'node:fs/promises'; import type { IModelCatalog, IModelCatalogEntry } from '../interfaces/catalog.ts'; import { MODEL_REGISTRY, TIMING } from '../constants.ts'; import { logger } from '../logger.ts'; export class ModelRegistry { private catalogUrl: string; private cachedCatalog: IModelCatalog | null = null; private cacheTime: number = 0; constructor(catalogUrl: string = MODEL_REGISTRY.DEFAULT_CATALOG_URL) { this.catalogUrl = catalogUrl; } public setCatalogUrl(url: string): void { this.catalogUrl = url; this.cachedCatalog = null; this.cacheTime = 0; } public async fetchCatalog(forceRefresh: boolean = false): Promise { if ( !forceRefresh && this.cachedCatalog && Date.now() - this.cacheTime < TIMING.GREENLIST_CACHE_DURATION_MS ) { return this.cachedCatalog; } try { logger.dim(`Fetching model catalog from: ${this.catalogUrl}`); const catalog = await this.readCatalogSource(this.catalogUrl); if (!Array.isArray(catalog.models)) { throw new Error('Invalid catalog format: missing models array'); } this.cachedCatalog = catalog; this.cacheTime = Date.now(); logger.dim(`Loaded ${catalog.models.length} catalog models`); return catalog; } catch (error) { logger.warn( `Failed to fetch model catalog: ${error instanceof Error ? error.message : String(error)}`, ); if (!this.cachedCatalog) { logger.dim('Using fallback catalog'); return this.getFallbackCatalog(); } return this.cachedCatalog; } } public async isModelListed(modelName: string): Promise { return (await this.getModel(modelName)) !== null; } public async getModel(modelName: string): Promise { const catalog = await this.fetchCatalog(); const normalized = this.normalizeModelName(modelName); return catalog.models.find((model) => { const candidates = [model.id, ...(model.aliases || [])]; return candidates.some((candidate) => this.normalizeModelName(candidate) === normalized); }) || null; } public async getAllModels(): Promise { const catalog = await this.fetchCatalog(); return catalog.models; } public async getModelsByEngine(engine: 'vllm'): Promise { const catalog = await this.fetchCatalog(); return catalog.models.filter((model) => model.engine === engine); } public async getModelsWithinVram(maxVramGb: number): Promise { const catalog = await this.fetchCatalog(); return catalog.models.filter((model) => model.requirements.minVramGb <= maxVramGb); } public async getRecommendedEngine(modelName: string): Promise<'vllm' | null> { const model = await this.getModel(modelName); return model ? model.engine : null; } public async getMinVram(modelName: string): Promise { const model = await this.getModel(modelName); return model ? model.requirements.minVramGb : null; } public async modelFitsInVram(modelName: string, availableVramGb: number): Promise { const minVram = await this.getMinVram(modelName); if (minVram === null) { return false; } return availableVramGb >= minVram; } public async searchModels(pattern: string): Promise { const catalog = await this.fetchCatalog(); const normalizedPattern = pattern.toLowerCase(); return catalog.models.filter((model) => model.id.toLowerCase().includes(normalizedPattern) || model.aliases?.some((alias) => alias.toLowerCase().includes(normalizedPattern)) || model.metadata?.summary?.toLowerCase().includes(normalizedPattern) || model.metadata?.tags?.some((tag) => tag.toLowerCase().includes(normalizedPattern)) ); } public async getModelsByTags(tags: string[]): Promise { const catalog = await this.fetchCatalog(); const normalizedTags = tags.map((tag) => tag.toLowerCase()); return catalog.models.filter((model) => model.metadata?.tags?.some((tag) => normalizedTags.includes(tag.toLowerCase())) ); } public clearCache(): void { this.cachedCatalog = null; this.cacheTime = 0; } public async printSummary(): Promise { const catalog = await this.fetchCatalog(); logger.logBoxTitle('Model Catalog', 70, 'info'); logger.logBoxLine(`Version: ${catalog.version}`); logger.logBoxLine(`Generated: ${catalog.generatedAt}`); logger.logBoxLine(`Total Models: ${catalog.models.length}`); logger.logBoxLine(''); for (const model of catalog.models.slice(0, 10)) { logger.logBoxLine( `- ${model.id} (${model.requirements.minVramGb}GB, ${model.engine})`, ); } if (catalog.models.length > 10) { logger.logBoxLine(`... and ${catalog.models.length - 10} more`); } logger.logBoxEnd(); } private async readCatalogSource(source: string): Promise { if (source.startsWith('file://')) { const filePath = new URL(source); const content = await fs.readFile(filePath, 'utf-8'); return JSON.parse(content) as IModelCatalog; } if (source.startsWith('/')) { const content = await fs.readFile(source, 'utf-8'); return JSON.parse(content) as IModelCatalog; } const controller = new AbortController(); const timeout = setTimeout(() => controller.abort(), 30000); try { const response = await fetch(source, { signal: controller.signal, headers: { Accept: 'application/json', 'User-Agent': 'ModelGrid/1.0', }, }); if (!response.ok) { throw new Error(`HTTP ${response.status}: ${response.statusText}`); } return await response.json() as IModelCatalog; } finally { clearTimeout(timeout); } } private getFallbackCatalog(): IModelCatalog { return { version: '1.0', generatedAt: new Date().toISOString(), models: MODEL_REGISTRY.FALLBACK_CATALOG as unknown as IModelCatalogEntry[], }; } private normalizeModelName(name: string): string { return name .toLowerCase() .replace(/[^a-z0-9:/._-]/g, '') .trim(); } }