lorehub/dist/core/embeddings.js

Capture and surface the collective wisdom of your codebase

import { pipeline } from '@xenova/transformers';
import { ConfigManager } from './config.js';
export const EMBEDDING_MODELS = {
    'all-MiniLM-L6-v2': {
        model: 'Xenova/all-MiniLM-L6-v2',
        dimensions: 384
    },
    'all-mpnet-base-v2': {
        model: 'Xenova/all-mpnet-base-v2',
        dimensions: 768
    },
    'gte-small': {
        model: 'Xenova/gte-small',
        dimensions: 384
    }
};
export class EmbeddingService {
    static instance;
    embedder = null;
    initPromise = null;
    currentModel;
    constructor() {
        const config = ConfigManager.getInstance();
        // Use config first, then env var, then default
        this.currentModel = config.get('embeddingModel') ||
            process.env.LOREHUB_EMBEDDING_MODEL ||
            'all-mpnet-base-v2';
    }
    async switchModel(modelName) {
        if (!EMBEDDING_MODELS[modelName]) {
            throw new Error(`Unknown embedding model: ${modelName}`);
        }
        // Reset the embedder to force re-initialization
        this.embedder = null;
        this.initPromise = null;
        this.currentModel = modelName;
        // Save to config
        const config = ConfigManager.getInstance();
        config.update({
            embeddingModel: modelName,
            embeddingDimensions: EMBEDDING_MODELS[modelName].dimensions
        });
        // Re-initialize with new model
        await this.initialize();
    }
    static getInstance() {
        if (!EmbeddingService.instance) {
            EmbeddingService.instance = new EmbeddingService();
        }
        return EmbeddingService.instance;
    }
    getModelConfig() {
        const config = EMBEDDING_MODELS[this.currentModel];
        if (!config) {
            throw new Error(`Unknown embedding model: ${this.currentModel}`);
        }
        return config;
    }
    getCurrentModel() {
        return this.currentModel;
    }
    async initialize() {
        if (this.embedder)
            return;
        // Only initialize once
        if (this.initPromise) {
            await this.initPromise;
            return;
        }
        this.initPromise = (async () => {
            const config = this.getModelConfig();
            // Use the configured model for semantic search
            this.embedder = await pipeline('feature-extraction', config.model);
        })();
        await this.initPromise;
    }
    async generateEmbedding(text) {
        await this.initialize();
        if (!this.embedder) {
            throw new Error('Embedder not initialized');
        }
        // Generate embeddings with mean pooling and normalization
        const output = await this.embedder(text, {
            pooling: 'mean',
            normalize: true
        });
        // Convert to Float32Array
        return new Float32Array(output.data);
    }
    async generateBatchEmbeddings(texts) {
        await this.initialize();
        if (!this.embedder) {
            throw new Error('Embedder not initialized');
        }
        // Process in batches for efficiency
        const batchSize = 32;
        const embeddings = [];
        for (let i = 0; i < texts.length; i += batchSize) {
            const batch = texts.slice(i, i + batchSize);
            const batchEmbeddings = await Promise.all(batch.map(text => this.generateEmbedding(text)));
            embeddings.push(...batchEmbeddings);
        }
        return embeddings;
    }
    // Combine lore content, why, and sigils for better semantic representation
    formatLoreForEmbedding(lore) {
        const parts = [lore.content];
        if (lore.why) {
            parts.push(`Context: ${lore.why}`);
        }
        if (lore.type) {
            parts.push(`Type: ${lore.type}`);
        }
        if (lore.sigils && lore.sigils.length > 0) {
            parts.push(`Sigils: ${lore.sigils.join(', ')}`);
        }
        return parts.join(' ');
    }
    // Get embedding dimension based on current model
    get dimension() {
        return this.getModelConfig().dimensions;
    }
}
//# sourceMappingURL=embeddings.js.map