/** * Unified Embedding Manager * * THE single source of truth for all embedding operations in Brainy. * Combines model management, precision configuration, and embedding generation * into one clean, maintainable class. * * Features: * - Singleton pattern ensures ONE model instance * - Automatic Q8 (default) or FP32 precision * - Model downloading and caching * - Thread-safe initialization * - Memory monitoring * * This replaces: SingletonModelManager, TransformerEmbedding, ModelPrecisionManager, * hybridModelManager, universalMemoryManager, and more. */ import { pipeline, env } from '@huggingface/transformers'; import { existsSync } from 'fs'; import { join } from 'path'; // Global state for true singleton across entire process let globalInstance = null; let globalInitPromise = null; /** * Unified Embedding Manager - Clean, simple, reliable */ export class EmbeddingManager { constructor() { this.model = null; this.modelName = 'Xenova/all-MiniLM-L6-v2'; this.initialized = false; this.initTime = null; this.embedCount = 0; this.locked = false; // Determine precision - Q8 by default this.precision = this.determinePrecision(); console.log(`๐ŸŽฏ EmbeddingManager: Using ${this.precision.toUpperCase()} precision`); } /** * Get the singleton instance */ static getInstance() { if (!globalInstance) { globalInstance = new EmbeddingManager(); } return globalInstance; } /** * Initialize the model (happens once) */ async init() { // In unit test mode, skip real model initialization if (process.env.BRAINY_UNIT_TEST === 'true' || globalThis.__BRAINY_UNIT_TEST__) { if (!this.initialized) { this.initialized = true; this.initTime = 1; // Mock init time console.log('๐Ÿงช EmbeddingManager: Using mocked embeddings for unit tests'); } return; } // Already initialized if (this.initialized && this.model) { return; } // Initialization in progress if (globalInitPromise) { await globalInitPromise; return; } // Start initialization globalInitPromise = this.performInit(); try { await globalInitPromise; } finally { globalInitPromise = null; } } /** * Perform actual initialization */ async performInit() { const startTime = Date.now(); console.log(`๐Ÿš€ Initializing embedding model (${this.precision.toUpperCase()})...`); try { // Configure transformers.js environment const modelsPath = this.getModelsPath(); env.cacheDir = modelsPath; env.allowLocalModels = true; env.useFSCache = true; // Check if models exist locally const modelPath = join(modelsPath, ...this.modelName.split('/')); const hasLocalModels = existsSync(modelPath); if (hasLocalModels) { console.log('โœ… Using cached models from:', modelPath); } // Configure pipeline options for the selected precision const pipelineOptions = { cache_dir: modelsPath, local_files_only: false, // Specify precision dtype: this.precision, quantized: this.precision === 'q8', // Memory optimizations session_options: { enableCpuMemArena: false, enableMemPattern: false, interOpNumThreads: 1, intraOpNumThreads: 1, graphOptimizationLevel: 'disabled' } }; // Load the model this.model = await pipeline('feature-extraction', this.modelName, pipelineOptions); // Lock precision after successful initialization this.locked = true; this.initialized = true; this.initTime = Date.now() - startTime; // Log success const memoryMB = this.getMemoryUsage(); console.log(`โœ… Model loaded in ${this.initTime}ms`); console.log(`๐Ÿ“Š Precision: ${this.precision.toUpperCase()} | Memory: ${memoryMB}MB`); console.log(`๐Ÿ”’ Configuration locked`); } catch (error) { this.initialized = false; this.model = null; throw new Error(`Failed to initialize embedding model: ${error instanceof Error ? error.message : String(error)}`); } } /** * Generate embeddings */ async embed(text) { // Check for unit test environment - use mocks to prevent ONNX conflicts if (process.env.BRAINY_UNIT_TEST === 'true' || globalThis.__BRAINY_UNIT_TEST__) { return this.getMockEmbedding(text); } // Ensure initialized await this.init(); if (!this.model) { throw new Error('Model not initialized'); } // Handle array input const input = Array.isArray(text) ? text.join(' ') : text; // Generate embedding const output = await this.model(input, { pooling: 'mean', normalize: true }); // Extract embedding vector const embedding = Array.from(output.data); // Validate dimensions if (embedding.length !== 384) { console.warn(`Unexpected embedding dimension: ${embedding.length}`); // Pad or truncate if (embedding.length < 384) { return [...embedding, ...new Array(384 - embedding.length).fill(0)]; } else { return embedding.slice(0, 384); } } this.embedCount++; return embedding; } /** * Generate mock embeddings for unit tests */ getMockEmbedding(text) { // Use the same mock logic as setup-unit.ts for consistency const input = Array.isArray(text) ? text.join(' ') : text; const str = typeof input === 'string' ? input : JSON.stringify(input); const vector = new Array(384).fill(0); // Create semi-realistic embeddings based on text content for (let i = 0; i < Math.min(str.length, 384); i++) { vector[i] = (str.charCodeAt(i % str.length) % 256) / 256; } // Add position-based variation for (let i = 0; i < 384; i++) { vector[i] += Math.sin(i * 0.1 + str.length) * 0.1; } // Track mock embedding count this.embedCount++; return vector; } /** * Get embedding function for compatibility */ getEmbeddingFunction() { return async (data) => { return await this.embed(data); }; } /** * Determine model precision */ determinePrecision() { // Check environment variable overrides if (process.env.BRAINY_MODEL_PRECISION === 'fp32') { return 'fp32'; } if (process.env.BRAINY_MODEL_PRECISION === 'q8') { return 'q8'; } if (process.env.BRAINY_FORCE_FP32 === 'true') { return 'fp32'; } // Default to Q8 - optimal for most use cases return 'q8'; } /** * Get models directory path */ getModelsPath() { // Check various possible locations const paths = [ process.env.BRAINY_MODELS_PATH, './models', join(process.cwd(), 'models'), join(process.env.HOME || '', '.brainy', 'models') ]; for (const path of paths) { if (path && existsSync(path)) { return path; } } // Default return join(process.cwd(), 'models'); } /** * Get memory usage in MB */ getMemoryUsage() { if (typeof process !== 'undefined' && process.memoryUsage) { const usage = process.memoryUsage(); return Math.round(usage.heapUsed / 1024 / 1024); } return null; } /** * Get current statistics */ getStats() { return { initialized: this.initialized, precision: this.precision, modelName: this.modelName, embedCount: this.embedCount, initTime: this.initTime, memoryMB: this.getMemoryUsage() }; } /** * Check if initialized */ isInitialized() { return this.initialized; } /** * Get current precision */ getPrecision() { return this.precision; } /** * Validate precision matches expected */ validatePrecision(expected) { if (this.locked && expected !== this.precision) { throw new Error(`Precision mismatch! System using ${this.precision.toUpperCase()} ` + `but ${expected.toUpperCase()} was requested. Cannot mix precisions.`); } } } // Export singleton instance and convenience functions export const embeddingManager = EmbeddingManager.getInstance(); /** * Direct embed function */ export async function embed(text) { return await embeddingManager.embed(text); } /** * Get embedding function for compatibility */ export function getEmbeddingFunction() { return embeddingManager.getEmbeddingFunction(); } /** * Get statistics */ export function getEmbeddingStats() { return embeddingManager.getStats(); } //# sourceMappingURL=EmbeddingManager.js.map