/** * 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 { Vector, EmbeddingFunction } from '../coreTypes.js' import { pipeline, env } from '@huggingface/transformers' import { existsSync } from 'fs' import { join } from 'path' // Types export type ModelPrecision = 'q8' | 'fp32' interface EmbeddingStats { initialized: boolean precision: ModelPrecision modelName: string embedCount: number initTime: number | null memoryMB: number | null } // Global state for true singleton across entire process let globalInstance: EmbeddingManager | null = null let globalInitPromise: Promise | null = null /** * Unified Embedding Manager - Clean, simple, reliable */ export class EmbeddingManager { private model: any = null private precision: ModelPrecision private modelName = 'Xenova/all-MiniLM-L6-v2' private initialized = false private initTime: number | null = null private embedCount = 0 private locked = false private constructor() { // Determine precision - Q8 by default this.precision = this.determinePrecision() console.log(`๐ŸŽฏ EmbeddingManager: Using ${this.precision.toUpperCase()} precision`) } /** * Get the singleton instance */ static getInstance(): EmbeddingManager { if (!globalInstance) { globalInstance = new EmbeddingManager() } return globalInstance } /** * Initialize the model (happens once) */ async init(): Promise { // In unit test mode, skip real model initialization if (process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__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 */ private async performInit(): Promise { 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: any = { 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: string | string[]): Promise { // Check for unit test environment - use mocks to prevent ONNX conflicts if (process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__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) as number[] // 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 */ private getMockEmbedding(text: string | string[]): Vector { // 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(): EmbeddingFunction { return async (data: string | string[]): Promise => { return await this.embed(data) } } /** * Determine model precision */ private determinePrecision(): ModelPrecision { // 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 */ private getModelsPath(): string { // 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 */ private getMemoryUsage(): number | null { if (typeof process !== 'undefined' && process.memoryUsage) { const usage = process.memoryUsage() return Math.round(usage.heapUsed / 1024 / 1024) } return null } /** * Get current statistics */ getStats(): EmbeddingStats { return { initialized: this.initialized, precision: this.precision, modelName: this.modelName, embedCount: this.embedCount, initTime: this.initTime, memoryMB: this.getMemoryUsage() } } /** * Check if initialized */ isInitialized(): boolean { return this.initialized } /** * Get current precision */ getPrecision(): ModelPrecision { return this.precision } /** * Validate precision matches expected */ validatePrecision(expected: ModelPrecision): void { 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: string | string[]): Promise { return await embeddingManager.embed(text) } /** * Get embedding function for compatibility */ export function getEmbeddingFunction(): EmbeddingFunction { return embeddingManager.getEmbeddingFunction() } /** * Get statistics */ export function getEmbeddingStats(): EmbeddingStats { return embeddingManager.getStats() }