@sconedev/ai_toolkit/dist/models/modelRegistry.js

Simplify AI integration in web apps with local and offline model support

import { isBrowser } from '../utils/browser-check';
import { downloadModelForOfflineUse } from '../core/onnxRuntime';
import * as fs from 'fs';
import * as path from 'path';
// Local model registry
let localModelRegistry = {};
// Models directory for Node.js
const MODEL_CACHE_DIR = isBrowser() ? '' : path.join(process.cwd(), '.model-cache');
/**
 * Register a local model
 */
export function registerLocalModel(config) {
    localModelRegistry[config.name] = {
        ...config,
        isLocal: true
    };
}
/**
 * Get a local model configuration
 */
export function getLocalModel(name) {
    return localModelRegistry[name] || null;
}
/**
 * List all registered local models
 */
export function listLocalModels() {
    return Object.keys(localModelRegistry);
}
/**
 * Download a remote model and register it for local use
 */
export async function downloadAndRegisterModel(modelUrl, config, options) {
    try {
        // Download the model
        const result = await downloadModelForOfflineUse(modelUrl, {
            modelName: config.name,
            progressCallback: options?.progressCallback,
            saveToIndexedDB: options?.saveToIndexedDB,
            saveToCache: options?.saveToCache
        });
        // Create the model config
        const modelConfig = {
            ...config,
            modelUrl,
            isLocal: true,
            localPath: typeof result === 'string' ? result : undefined
        };
        // Register the model
        registerLocalModel(modelConfig);
        return modelConfig;
    }
    catch (error) {
        throw new Error(`Failed to download and register model: ${error instanceof Error ? error.message : String(error)}`);
    }
}
/**
 * Initialize the model registry on startup
 */
export async function initModelRegistry() {
    if (isBrowser()) {
        try {
            // In browser, try to load model registry from localStorage
            const savedRegistry = localStorage.getItem('onnx_model_registry');
            if (savedRegistry) {
                localModelRegistry = JSON.parse(savedRegistry);
            }
        }
        catch (err) {
            console.warn('Failed to load model registry from localStorage:', err);
        }
    }
    else {
        // In Node.js, scan the model cache directory
        try {
            if (!fs.existsSync(MODEL_CACHE_DIR)) {
                fs.mkdirSync(MODEL_CACHE_DIR, { recursive: true });
            }
            // Look for a registry file
            const registryPath = path.join(MODEL_CACHE_DIR, 'registry.json');
            if (fs.existsSync(registryPath)) {
                try {
                    const registryContent = fs.readFileSync(registryPath, 'utf8');
                    localModelRegistry = JSON.parse(registryContent);
                }
                catch (err) {
                    console.warn('Failed to parse model registry file:', err);
                }
            }
            // Scan directory for .onnx files not in registry
            const files = fs.readdirSync(MODEL_CACHE_DIR);
            for (const file of files) {
                if (file.endsWith('.onnx')) {
                    const modelName = path.basename(file, '.onnx');
                    if (!localModelRegistry[modelName]) {
                        // Add to registry with minimal info
                        localModelRegistry[modelName] = {
                            name: modelName,
                            modelUrl: file,
                            inputNames: [], // These will need to be filled in manually
                            outputNames: [], // These will need to be filled in manually
                            isLocal: true,
                            localPath: path.join(MODEL_CACHE_DIR, file)
                        };
                    }
                }
            }
            // Save the updated registry
            saveModelRegistry();
        }
        catch (err) {
            console.warn('Failed to initialize model registry:', err);
        }
    }
}
/**
 * Save the model registry
 */
export function saveModelRegistry() {
    if (isBrowser()) {
        try {
            localStorage.setItem('onnx_model_registry', JSON.stringify(localModelRegistry));
        }
        catch (err) {
            console.warn('Failed to save model registry to localStorage:', err);
        }
    }
    else {
        try {
            const registryPath = path.join(MODEL_CACHE_DIR, 'registry.json');
            fs.writeFileSync(registryPath, JSON.stringify(localModelRegistry, null, 2));
        }
        catch (err) {
            console.warn('Failed to save model registry to file:', err);
        }
    }
}
// Initialize the registry when the module loads
initModelRegistry().catch(console.error);