/** * Unified Embedding Manager * * THE single source of truth for all embedding operations in Brainy. * Uses Candle WASM inference for universal compatibility. * * Features: * - Singleton pattern ensures ONE model instance * - Candle WASM (no transformers.js or ONNX Runtime dependency) * - Bundled model (no runtime downloads) * - Works everywhere: Node.js, Bun, Bun --compile, browsers * - Memory monitoring */ import { Vector, EmbeddingFunction } from '../coreTypes.js' import { WASMEmbeddingEngine } from './wasm/index.js' // 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 * * Now powered by Candle WASM for universal compatibility. */ export class EmbeddingManager { private engine: WASMEmbeddingEngine private precision: ModelPrecision = 'q8' private modelName = 'all-MiniLM-L6-v2' private initialized = false private initTime: number | null = null private embedCount = 0 private locked = false private constructor() { this.engine = WASMEmbeddingEngine.getInstance() // Log deferred to init() โ€” at construction time we don't know if a plugin // (like Cortex) will replace the WASM embedder with a native one. } /** * 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 const isTestMode = process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__ if (isTestMode) { // Production safeguard if (process.env.NODE_ENV === 'production') { throw new Error( 'CRITICAL: Mock embeddings detected in production environment! ' + 'BRAINY_UNIT_TEST or __BRAINY_UNIT_TEST__ is set while NODE_ENV=production. ' + 'This is a security risk. Remove test flags before deploying to production.' ) } 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.engine.isInitialized()) { 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() try { await this.engine.initialize() // Lock precision after successful initialization this.locked = true this.initialized = true this.initTime = Date.now() - startTime // Log success const memoryMB = this.getMemoryUsage() console.log(`๐Ÿ“Š Precision: Q8 | Memory: ${memoryMB}MB`) console.log('๐Ÿ”’ Configuration locked') } catch (error) { this.initialized = false throw new Error( `Failed to initialize embedding model: ${error instanceof Error ? error.message : String(error)}` ) } } /** * Generate embeddings */ async embed(text: string | string[] | Record): Promise { // Check for unit test environment const isTestMode = process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__ if (isTestMode) { if (process.env.NODE_ENV === 'production') { throw new Error('CRITICAL: Mock embeddings in production!') } return this.getMockEmbedding(text) } // Ensure initialized await this.init() // Normalize input to string let input: string if (Array.isArray(text)) { input = text.map((t) => (typeof t === 'string' ? t : String(t))).join(' ') } else if (typeof text === 'string') { input = text } else if (typeof text === 'object') { input = JSON.stringify(text) } else { console.warn('EmbeddingManager.embed received unexpected input type:', typeof text) input = String(text) } // Generate embedding using WASM engine const embedding = await this.engine.embed(input) // Validate dimensions if (embedding.length !== 384) { console.warn(`Unexpected embedding dimension: ${embedding.length}`) 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[] | Record): Vector { 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 } this.embedCount++ return vector } /** * Batch embed multiple texts using native WASM batch API * * Uses the engine's embedBatch() for a single WASM forward pass * instead of N individual embed() calls. * * Large batches (>MICRO_BATCH_SIZE) are split into micro-batches * with event loop yielding between each, preventing the synchronous * WASM call from blocking the server for hundreds of milliseconds. * * @param texts Array of strings to embed * @returns Array of embedding vectors (384 dimensions each) */ async embedBatch(texts: string[], options?: { signal?: AbortSignal }): Promise { if (texts.length === 0) return [] const isTestMode = process.env.BRAINY_UNIT_TEST === 'true' || (globalThis as any).__BRAINY_UNIT_TEST__ if (isTestMode) { if (process.env.NODE_ENV === 'production') { throw new Error('CRITICAL: Mock embeddings in production!') } return texts.map(t => this.getMockEmbedding(t)) } await this.init() // Small batches: single WASM call (no overhead) const MICRO_BATCH_SIZE = 20 if (texts.length <= MICRO_BATCH_SIZE) { const results = await this.engine.embedBatch(texts) this.embedCount += texts.length return results } // Large batches: micro-batch with event loop yielding // Each micro-batch of ~20 texts blocks ~10-30ms, then yields // so other requests (HTTP, timers, I/O) can proceed const allResults: number[][] = [] for (let i = 0; i < texts.length; i += MICRO_BATCH_SIZE) { if (options?.signal?.aborted) { return allResults } const batch = texts.slice(i, i + MICRO_BATCH_SIZE) const batchResults = await this.engine.embedBatch(batch) allResults.push(...batchResults) this.embedCount += batch.length // Yield to event loop between micro-batches if (i + MICRO_BATCH_SIZE < texts.length) { await new Promise(resolve => setTimeout(resolve, 0)) } } return allResults } /** * Get embedding function for compatibility */ getEmbeddingFunction(): EmbeddingFunction { return async (data: string | string[] | Record): Promise => { return await this.embed(data) } } /** * 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 { const engineStats = this.engine.getStats() return { initialized: this.initialized, precision: this.precision, modelName: this.modelName, embedCount: this.embedCount + engineStats.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[] | Record ): 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() }