@bratcliffe909/mcp-server-segmind/dist/models/registry.js

Model Context Protocol server for Segmind API - Generate images and videos using AI models

import { z } from 'zod';
import { logger } from '../utils/logger.js';
import { ModelCategory, OutputType } from './types.js';
import { WORKING_MODELS } from './working-models.js';
export { ModelCategory, OutputType };
export class ModelRegistry {
    models = new Map();
    modelsByCategory = new Map();
    constructor() {
        this.initializeModels();
    }
    initializeModels() {
        Object.values(ModelCategory).forEach(category => {
            this.modelsByCategory.set(category, []);
        });
        WORKING_MODELS.forEach(model => {
            this.registerModel(model);
        });
        logger.info('Model registry initialized', {
            totalModels: this.models.size,
            categories: Object.values(ModelCategory).map(cat => ({
                category: cat,
                count: this.modelsByCategory.get(cat)?.length || 0,
            })),
        });
    }
    registerModel(config) {
        this.models.set(config.id, config);
        const categoryModels = this.modelsByCategory.get(config.category) || [];
        categoryModels.push(config.id);
        this.modelsByCategory.set(config.category, categoryModels);
    }
    getModel(modelId) {
        return this.models.get(modelId);
    }
    getModelsByCategory(category) {
        const modelIds = this.modelsByCategory.get(category) || [];
        return modelIds.map(id => this.models.get(id)).filter(Boolean);
    }
    getAllModels() {
        return Array.from(this.models.values());
    }
    findModelByEndpoint(endpoint) {
        return Array.from(this.models.values()).find(model => model.endpoint === endpoint);
    }
    searchModels(query) {
        const lowerQuery = query.toLowerCase();
        return Array.from(this.models.values()).filter(model => model.name.toLowerCase().includes(lowerQuery) ||
            model.description.toLowerCase().includes(lowerQuery) ||
            model.id.toLowerCase().includes(lowerQuery));
    }
    getModelCategories() {
        return Object.values(ModelCategory);
    }
    getAllCategories() {
        return Array.from(this.modelsByCategory.keys());
    }
    validateModelParameters(modelId, params) {
        const model = this.getModel(modelId);
        if (!model) {
            return { success: false, error: 'Model not found' };
        }
        try {
            const validated = model.parameters.parse(params);
            return { success: true, data: validated };
        }
        catch (error) {
            if (error instanceof z.ZodError) {
                return {
                    success: false,
                    error: error.errors.map((e) => `${e.path.join('.')}: ${e.message}`).join(', ')
                };
            }
            return { success: false, error: 'Invalid parameters' };
        }
    }
}
export const modelRegistry = new ModelRegistry();
//# sourceMappingURL=registry.js.map