@reftrixmcp/ml 0.1.7

package/dist/embeddings/service.d.ts

@@ -0,0 +1,261 @@
+import type { EmbeddingTextType, EmbeddingServiceConfig, CacheStats } from "./types.js";
+/**
+ * Default maximum cache size for LRU cache
+ * Prevents unbounded memory growth
+ */
+export declare const DEFAULT_MAX_CACHE_SIZE = 5000;
+/**
+ * Default number of inferences before recycling the ONNX pipeline.
+ *
+ * onnxruntime-node uses a native C++ arena allocator that grows monotonically
+ * during inference (~65-100MB per inference) but never returns memory to the OS.
+ * Pipeline recycling (dispose + recreate) tears down the InferenceSession via
+ * session.release(), which frees the arena's native memory back to the OS.
+ *
+ * With threshold=10, the arena is reset every 10 inferences, keeping peak
+ * additional native memory to ~650MB-1GB before each recycle. Model reload
+ * takes ~1.5-2s per recycle.
+ *
+ * v0.1.0: Reduced from 20 to 10 for universal embedding chunking.
+ * v0.1.0: Increased from 10 to 30 — chunking already limits items per chunk
+ * to 30, so threshold=30 means at most 1 recycle per chunk. With threshold=10,
+ * 3 recycles occurred per chunk, and model reload (~1.5-2s each) caused
+ * re-initialization overhead of ~45%. threshold=30 aligns with chunk size.
+ *
+ * Note: This project uses onnxruntime-node (native C++ bindings), NOT WASM.
+ * dispose() correctly frees native memory unlike WASM backing store pages.
+ */
+export declare const DEFAULT_PIPELINE_RECYCLE_THRESHOLD = 30;
+/**
+ * Calculate cosine similarity between two vectors
+ */
+export declare function cosineSimilarity(a: number[], b: number[]): number;
+/**
+ * EmbeddingService - Generates embeddings using multilingual-e5-base
+ *
+ * v0.1.0: ONNX inference is offloaded to a Worker Thread by default.
+ * The main thread retains the LRU cache and E5 prefix logic. Only the
+ * actual ONNX pipeline.call() runs in the worker, keeping the main
+ * event loop responsive for BullMQ heartbeats and IPC messages.
+ */
+export declare class EmbeddingService {
+    private config;
+    private cache;
+    private cacheHits;
+    private cacheMisses;
+    private cacheEvictions;
+    private useWorkerThread;
+    private worker;
+    private workerInitPromise;
+    private workerReady;
+    private workerRestartCount;
+    private pendingRequests;
+    /** Timestamp of last worker crash, used to enforce cooldown between restarts. */
+    private lastCrashTime;
+    private currentProvider;
+    private idleTimer;
+    private idleTimeoutMs;
+    private pipeline;
+    private initPromise;
+    /**
+     * Number of inference calls since the last pipeline recycle.
+     * In worker-thread mode, the worker manages its own counter.
+     * In fallback mode, this counter is used directly.
+     */
+    private inferencesSinceRecycle;
+    /**
+     * Total number of pipeline recycles performed during this service's lifetime.
+     * In worker-thread mode, the count is approximate (tracked by dispose calls).
+     */
+    private totalRecycles;
+    constructor(config?: EmbeddingServiceConfig);
+    /**
+     * Spawn the worker thread and send init message.
+     * Enforces a cooldown after crashes to prevent rapid spawn-crash-spawn cycles
+     * that consume memory (~10GB per CUDA init attempt).
+     */
+    private ensureWorkerReady;
+    private spawnAndInitWorker;
+    /**
+     * Send a message to the worker and return a Promise for the response.
+     */
+    private sendWorkerMessage;
+    /**
+     * Route a response from the worker to the pending request's promise.
+     */
+    private handleWorkerResponse;
+    /**
+     * Handle worker thread crash: reject all pending requests, attempt restart.
+     */
+    private handleWorkerCrash;
+    /**
+     * Check if the worker can be restarted (under max restart limit).
+     */
+    private canRestartWorker;
+    /**
+     * Initialize the in-process embedding pipeline (lazy loading).
+     * Used when worker thread is disabled.
+     *
+     * If the configured device is 'cuda' but initialization fails (e.g. because
+     * LD_LIBRARY_PATH was not set at the OS level), automatically falls back to CPU.
+     */
+    private initializeInProcess;
+    /**
+     * Recycle the in-process pipeline if threshold reached.
+     */
+    private recyclePipelineIfNeeded;
+    /**
+     * Dispose the in-process pipeline.
+     */
+    private disposeInProcess;
+    /**
+     * Generate a single embedding in-process (fallback).
+     */
+    private generateInProcess;
+    /**
+     * Generate batch embeddings in-process (fallback).
+     */
+    private generateBatchInProcess;
+    /**
+     * Set the idle timeout for automatic VRAM release.
+     *
+     * When set to a positive value (ms), the service will automatically call
+     * dispose() after no embedding generation activity for the specified duration.
+     * This frees GPU VRAM (e.g. 1,406 MiB for multilingual-e5-base on CUDA)
+     * so that other GPU consumers (e.g. Ollama Vision at 11.7 GB) can use it.
+     *
+     * dispose() preserves the CUDA provider setting, so the next inference call
+     * will re-initialize with CUDA automatically (unlike releaseGpu() which
+     * forces CPU).
+     *
+     * @param ms - Timeout in milliseconds. 0 disables the idle timer.
+     */
+    setIdleTimeout(ms: number): void;
+    /**
+     * Reset the idle timer after an embedding generation completes.
+     * Called in finally blocks of generateEmbedding() and generateBatchEmbeddings().
+     */
+    private resetIdleTimer;
+    /**
+     * Check if the model is initialized (worker thread is ready or pipeline loaded).
+     */
+    isInitialized(): boolean;
+    /**
+     * Get the total number of pipeline recycles performed.
+     */
+    getRecycleCount(): number;
+    /**
+     * Get the number of inferences since the last pipeline recycle.
+     * In worker-thread mode this is approximate (tracked via dispose calls).
+     */
+    getInferencesSinceRecycle(): number;
+    /**
+     * Generate cache key for embedding
+     */
+    private getCacheKey;
+    /**
+     * Set cache entry with LRU eviction
+     */
+    private setCacheEntry;
+    /**
+     * Get cache entry with LRU update (move to end)
+     */
+    private getCacheEntry;
+    /**
+     * Generate embedding for a single text
+     *
+     * @param text - Input text to embed
+     * @param type - Type of text ('query' or 'passage')
+     * @returns 768-dimensional normalized embedding vector
+     */
+    generateEmbedding(text: string, type: EmbeddingTextType): Promise<number[]>;
+    /**
+     * Generate embedding via Worker Thread, with crash recovery.
+     */
+    private generateViaWorker;
+    /**
+     * Generate batch embeddings via Worker Thread, with crash recovery.
+     */
+    private generateBatchViaWorker;
+    /**
+     * Generate embeddings for multiple texts
+     *
+     * @param texts - Array of input texts
+     * @param type - Type of texts ('query' or 'passage')
+     * @returns Array of 768-dimensional normalized embedding vectors
+     */
+    generateBatchEmbeddings(texts: string[], type: EmbeddingTextType): Promise<number[][]>;
+    /**
+     * Dispose the ONNX pipeline to free native memory.
+     *
+     * In worker-thread mode, sends a dispose message to the worker.
+     * In fallback mode, disposes the in-process pipeline.
+     *
+     * After disposal, the next inference call will re-initialize automatically.
+     */
+    dispose(): Promise<void>;
+    /**
+     * Terminate the worker thread completely and clean up.
+     * After this, the service cannot be used until a new instance is created.
+     */
+    terminate(): Promise<void>;
+    /**
+     * Get cache statistics
+     */
+    getCacheStats(): CacheStats;
+    /**
+     * Clear the embedding cache
+     */
+    clearCache(): void;
+    /**
+     * Whether this service is using a Worker Thread for inference.
+     */
+    isUsingWorkerThread(): boolean;
+    /**
+     * Get the current ONNX execution provider ('cpu' or 'cuda').
+     */
+    getCurrentProvider(): "cpu" | "cuda";
+    /**
+     * Switch ONNX execution provider at runtime.
+     *
+     * In worker-thread mode: sends a switch-provider message to the worker,
+     * which disposes the current pipeline and updates the device config.
+     * In in-process mode: disposes the pipeline and updates the config
+     * so the next inference call re-initializes with the new provider.
+     *
+     * @param provider - The target execution provider ('cpu' or 'cuda').
+     * @returns true if the switch succeeded and the requested provider is now
+     *          active; false if the switch was not possible (e.g. CUDA not
+     *          available) — in that case the service remains on CPU.
+     */
+    switchProvider(provider: "cpu" | "cuda"): Promise<boolean>;
+    /**
+     * Switch provider via Worker Thread message.
+     */
+    private switchProviderViaWorker;
+    /**
+     * Switch provider in-process (fallback mode).
+     */
+    private switchProviderInProcess;
+    /**
+     * Release GPU resources by disposing the pipeline and switching to CPU.
+     *
+     * Safe to call even if already on CPU — acts as a dispose + ensures
+     * the provider is set to CPU.
+     */
+    releaseGpu(): Promise<void>;
+    /**
+     * Release GPU via Worker Thread message.
+     */
+    private releaseGpuViaWorker;
+    /**
+     * Release GPU in-process (fallback mode).
+     */
+    private releaseGpuInProcess;
+    /**
+     * Get the number of worker thread restarts that have occurred.
+     */
+    getWorkerRestartCount(): number;
+}
+export declare const embeddingService: EmbeddingService;
+//# sourceMappingURL=service.d.ts.map

package/dist/embeddings/service.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"service.d.ts","sourceRoot":"","sources":["../../src/embeddings/service.ts"],"names":[],"mappings":"AAqBA,OAAO,KAAK,EAAE,iBAAiB,EAAE,sBAAsB,EAAE,UAAU,EAAE,MAAM,YAAY,CAAC;AAoBxF;;;GAGG;AACH,eAAO,MAAM,sBAAsB,OAAO,CAAC;AAE3C;;;;;;;;;;;;;;;;;;;;GAoBG;AACH,eAAO,MAAM,kCAAkC,KAAK,CAAC;AA2ErD;;GAEG;AACH,wBAAgB,gBAAgB,CAAC,CAAC,EAAE,MAAM,EAAE,EAAE,CAAC,EAAE,MAAM,EAAE,GAAG,MAAM,CAwBjE;AAgBD;;;;;;;GAOG;AACH,qBAAa,gBAAgB;IAC3B,OAAO,CAAC,MAAM,CAAmC;IACjD,OAAO,CAAC,KAAK,CAAoC;IACjD,OAAO,CAAC,SAAS,CAAK;IACtB,OAAO,CAAC,WAAW,CAAK;IACxB,OAAO,CAAC,cAAc,CAAK;IAG3B,OAAO,CAAC,eAAe,CAAU;IACjC,OAAO,CAAC,MAAM,CAAuB;IACrC,OAAO,CAAC,iBAAiB,CAA8B;IACvD,OAAO,CAAC,WAAW,CAAS;IAC5B,OAAO,CAAC,kBAAkB,CAAK;IAC/B,OAAO,CAAC,eAAe,CAA0C;IACjE,iFAAiF;IACjF,OAAO,CAAC,aAAa,CAAK;IAG1B,OAAO,CAAC,eAAe,CAAyB;IAGhD,OAAO,CAAC,SAAS,CAA8C;IAC/D,OAAO,CAAC,aAAa,CAAa;IAGlC,OAAO,CAAC,QAAQ,CAAmC;IACnD,OAAO,CAAC,WAAW,CAA8B;IAEjD;;;;OAIG;IACH,OAAO,CAAC,sBAAsB,CAAK;IAEnC;;;OAGG;IACH,OAAO,CAAC,aAAa,CAAK;gBAEd,MAAM,GAAE,sBAA2B;IA+B/C;;;;OAIG;YACW,iBAAiB;YAuBjB,kBAAkB;IA6EhC;;OAEG;IACH,OAAO,CAAC,iBAAiB;IAoBzB;;OAEG;IACH,OAAO,CAAC,oBAAoB;IAS5B;;OAEG;IACH,OAAO,CAAC,iBAAiB;IAoBzB;;OAEG;IACH,OAAO,CAAC,gBAAgB;IAQxB;;;;;;OAMG;YACW,mBAAmB;IAiEjC;;OAEG;YACW,uBAAuB;IAgCrC;;OAEG;YACW,gBAAgB;IAa9B;;OAEG;YACW,iBAAiB;IA6B/B;;OAEG;YACW,sBAAsB;IA8CpC;;;;;;;;;;;;;OAaG;IACH,cAAc,CAAC,EAAE,EAAE,MAAM,GAAG,IAAI;IAahC;;;OAGG;IACH,OAAO,CAAC,cAAc;IAmCtB;;OAEG;IACH,aAAa,IAAI,OAAO;IAOxB;;OAEG;IACH,eAAe,IAAI,MAAM;IAIzB;;;OAGG;IACH,yBAAyB,IAAI,MAAM;IAInC;;OAEG;IACH,OAAO,CAAC,WAAW;IAInB;;OAEG;IACH,OAAO,CAAC,aAAa;IAiBrB;;OAEG;IACH,OAAO,CAAC,aAAa;IASrB;;;;;;OAMG;IACG,iBAAiB,CAAC,IAAI,EAAE,MAAM,EAAE,IAAI,EAAE,iBAAiB,GAAG,OAAO,CAAC,MAAM,EAAE,CAAC;IAwDjF;;OAEG;YACW,iBAAiB;IA0B/B;;OAEG;YACW,sBAAsB;IAyBpC;;;;;;OAMG;IACG,uBAAuB,CAAC,KAAK,EAAE,MAAM,EAAE,EAAE,IAAI,EAAE,iBAAiB,GAAG,OAAO,CAAC,MAAM,EAAE,EAAE,CAAC;IA6E5F;;;;;;;OAOG;IACG,OAAO,IAAI,OAAO,CAAC,IAAI,CAAC;IAkC9B;;;OAGG;IACG,SAAS,IAAI,OAAO,CAAC,IAAI,CAAC;IAwChC;;OAEG;IACH,aAAa,IAAI,UAAU;IAS3B;;OAEG;IACH,UAAU,IAAI,IAAI;IAYlB;;OAEG;IACH,mBAAmB,IAAI,OAAO;IAI9B;;OAEG;IACH,kBAAkB,IAAI,KAAK,GAAG,MAAM;IAIpC;;;;;;;;;;;;OAYG;IACG,cAAc,CAAC,QAAQ,EAAE,KAAK,GAAG,MAAM,GAAG,OAAO,CAAC,OAAO,CAAC;IAgBhE;;OAEG;YACW,uBAAuB;IAkCrC;;OAEG;YACW,uBAAuB;IAwDrC;;;;;OAKG;IACG,UAAU,IAAI,OAAO,CAAC,IAAI,CAAC;IAajC;;OAEG;YACW,mBAAmB;IAgBjC;;OAEG;YACW,mBAAmB;IAMjC;;OAEG;IACH,qBAAqB,IAAI,MAAM;CAGhC;AAGD,eAAO,MAAM,gBAAgB,kBAAyB,CAAC"}