/** * Universal Memory Manager for Embeddings * * Works in ALL environments: Node.js, browsers, serverless, workers * Solves transformers.js memory leak with environment-specific strategies */ import { Vector, EmbeddingFunction } from '../coreTypes.js' // Environment detection const isNode = typeof process !== 'undefined' && process.versions?.node const isBrowser = typeof window !== 'undefined' && typeof document !== 'undefined' const isServerless = typeof process !== 'undefined' && ( process.env.VERCEL || process.env.NETLIFY || process.env.AWS_LAMBDA_FUNCTION_NAME || process.env.FUNCTIONS_WORKER_RUNTIME ) interface MemoryStats { embeddings: number memoryUsage: string restarts: number strategy: string } export class UniversalMemoryManager { private embeddingFunction: any = null private embedCount = 0 private restartCount = 0 private lastRestart = 0 private strategy: string private maxEmbeddings: number constructor() { // CRITICAL FIX: Never use worker threads with ONNX Runtime // Worker threads cause HandleScope V8 API errors due to isolate issues // Always use direct embedding on main thread for ONNX compatibility if (isServerless) { this.strategy = 'serverless-restart' this.maxEmbeddings = 50 // Restart frequently in serverless } else if (isNode && !isBrowser) { // CHANGED: Use direct strategy instead of node-worker to avoid V8 isolate issues this.strategy = 'node-direct' this.maxEmbeddings = 200 // Main thread can handle more with single model instance } else if (isBrowser) { this.strategy = 'browser-dispose' this.maxEmbeddings = 25 // Browser memory is limited } else { this.strategy = 'fallback-dispose' this.maxEmbeddings = 75 } console.log(`๐Ÿง  Universal Memory Manager: Using ${this.strategy} strategy`) console.log('โœ… UNIVERSAL: Memory-safe embedding system initialized') } async getEmbeddingFunction(): Promise { return async (data: string | string[]): Promise => { return this.embed(data) } } async embed(data: string | string[]): Promise { // Check if we need to restart/cleanup await this.checkMemoryLimits() // Ensure embedding function is available await this.ensureEmbeddingFunction() // Perform embedding const result = await this.embeddingFunction.embed(data) this.embedCount++ return result } private async checkMemoryLimits(): Promise { if (this.embedCount >= this.maxEmbeddings) { console.log(`๐Ÿ”„ Memory cleanup: ${this.embedCount} embeddings processed`) await this.cleanup() } } private async ensureEmbeddingFunction(): Promise { if (this.embeddingFunction) { return } switch (this.strategy) { case 'node-direct': await this.initNodeDirect() break case 'serverless-restart': await this.initServerless() break case 'browser-dispose': await this.initBrowser() break default: await this.initFallback() } } private async initNodeDirect(): Promise { if (isNode) { // CRITICAL: Use direct embedding to avoid worker thread V8 isolate issues // This prevents HandleScope errors and ensures single model instance console.log('โœ… Using Node.js direct embedding (main thread - ONNX compatible)') await this.initDirect() } } private async initServerless(): Promise { // In serverless, use direct embedding but restart more aggressively await this.initDirect() console.log('โœ… Using serverless strategy with aggressive cleanup') } private async initBrowser(): Promise { // In browser, use direct embedding with disposal await this.initDirect() console.log('โœ… Using browser strategy with disposal') } private async initFallback(): Promise { await this.initDirect() console.log('โœ… Using fallback direct embedding strategy') } private async initDirect(): Promise { try { // Dynamic import to handle different environments const { TransformerEmbedding } = await import('../utils/embedding.js') this.embeddingFunction = new TransformerEmbedding({ verbose: false, precision: 'fp32', localFilesOnly: process.env.BRAINY_ALLOW_REMOTE_MODELS !== 'true' }) await this.embeddingFunction.init() console.log('โœ… Direct embedding function initialized') } catch (error) { throw new Error(`Failed to initialize embedding function: ${error instanceof Error ? error.message : String(error)}`) } } private async cleanup(): Promise { const startTime = Date.now() try { // Strategy-specific cleanup switch (this.strategy) { case 'node-worker': if (this.embeddingFunction?.forceRestart) { await this.embeddingFunction.forceRestart() } break case 'serverless-restart': // In serverless, create new instance if (this.embeddingFunction?.dispose) { this.embeddingFunction.dispose() } this.embeddingFunction = null break case 'browser-dispose': // In browser, try disposal if (this.embeddingFunction?.dispose) { this.embeddingFunction.dispose() } // Force garbage collection if available if (typeof window !== 'undefined' && (window as any).gc) { (window as any).gc() } break default: // Fallback: dispose and recreate if (this.embeddingFunction?.dispose) { this.embeddingFunction.dispose() } this.embeddingFunction = null } this.embedCount = 0 this.restartCount++ this.lastRestart = Date.now() const cleanupTime = Date.now() - startTime console.log(`๐Ÿงน Memory cleanup completed in ${cleanupTime}ms (strategy: ${this.strategy})`) } catch (error) { console.warn('โš ๏ธ Cleanup failed:', error instanceof Error ? error.message : String(error)) // Force null assignment as last resort this.embeddingFunction = null } } getMemoryStats(): MemoryStats { let memoryUsage = 'unknown' // Get memory stats based on environment if (isNode && typeof process !== 'undefined') { const mem = process.memoryUsage() memoryUsage = `${(mem.heapUsed / 1024 / 1024).toFixed(2)} MB` } else if (isBrowser && (performance as any).memory) { const mem = (performance as any).memory memoryUsage = `${(mem.usedJSHeapSize / 1024 / 1024).toFixed(2)} MB` } return { embeddings: this.embedCount, memoryUsage, restarts: this.restartCount, strategy: this.strategy } } async dispose(): Promise { if (this.embeddingFunction) { if (this.embeddingFunction.dispose) { await this.embeddingFunction.dispose() } this.embeddingFunction = null } } } // Export singleton instance export const universalMemoryManager = new UniversalMemoryManager() // Export convenience function export async function getUniversalEmbeddingFunction(): Promise { return universalMemoryManager.getEmbeddingFunction() } // Export memory stats function export function getEmbeddingMemoryStats(): MemoryStats { return universalMemoryManager.getMemoryStats() }