@vellumai/assistant 0.3.28 → 0.4.0

package/src/memory/embedding-local.ts

@@ -1,27 +1,48 @@
-import { dirname, join } from 'node:path';
+import { existsSync, unlinkSync, writeFileSync } from 'node:fs';
+import { join } from 'node:path';
 
 import { getLogger } from '../util/logger.js';
+import { getEmbeddingModelsDir, getRootDir } from '../util/platform.js';
 import { PromiseGuard } from '../util/promise-guard.js';
 import type { EmbeddingBackend, EmbeddingRequestOptions } from './embedding-backend.js';
+import { EmbeddingRuntimeManager } from './embedding-runtime-manager.js';
 
 const log = getLogger('memory-embedding-local');
 
-type FeatureExtractionPipeline = (
-  texts: string | string[],
-  options?: { pooling?: string; normalize?: boolean },
-) => Promise<{ tolist: () => number[][] }>;
+interface WorkerResponse {
+  id?: number;
+  type?: string;
+  vectors?: number[][];
+  error?: string;
+}
 
 /**
  * Local embedding backend using @huggingface/transformers (ONNX Runtime).
  * Runs BAAI/bge-small-en-v1.5 locally — no API calls, no network required.
  *
- * The model is downloaded on first use and cached by the transformers library.
+ * Embeddings run in a **separate bun process** because compiled Bun binaries
+ * cannot resolve bare specifier imports in dynamically loaded files. The embed
+ * worker communicates via JSON-lines over stdin/stdout.
+ *
+ * The embedding runtime (onnxruntime-node + transformers + bun) is downloaded
+ * post-hatch by EmbeddingRuntimeManager.
+ *
  * Produces 384-dimensional embeddings.
  */
 export class LocalEmbeddingBackend implements EmbeddingBackend {
   readonly provider = 'local' as const;
   readonly model: string;
-  private extractor: FeatureExtractionPipeline | null = null;
+
+  // Subprocess — typed loosely to avoid coupling to Bun's Subprocess generics
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  private workerProc: any = null;
+  private stdoutBuffer = '';
+  private requestCounter = 0;
+  private pendingRequests = new Map<number, {
+    resolve: (response: WorkerResponse) => void;
+  }>();
+  private stdoutReaderActive = false;
+
   private readonly initGuard = new PromiseGuard<void>();
 
   constructor(model: string) {
@@ -35,55 +56,252 @@ export class LocalEmbeddingBackend implements EmbeddingBackend {
     await this.ensureInitialized();
 
     const results: number[][] = [];
-    // Process in batches of 32 to avoid OOM with large inputs
     const batchSize = 32;
     for (let i = 0; i < texts.length; i += batchSize) {
       if (options?.signal?.aborted) throw new DOMException('Aborted', 'AbortError');
       const batch = texts.slice(i, i + batchSize);
-      const output = await this.extractor!(batch, {
-        pooling: 'cls',
-        normalize: true,
-      });
-      const vectors = output.tolist();
-      results.push(...vectors);
+      const response = await this.sendRequest(batch);
+      if (response.error) {
+        throw new Error(`Embedding worker error: ${response.error}`);
+      }
+      if (!response.vectors) {
+        throw new Error('Embedding worker returned no vectors');
+      }
+      results.push(...response.vectors);
     }
-
     return results;
   }
 
+  private sendRequest(texts: string[]): Promise<WorkerResponse> {
+    const id = ++this.requestCounter;
+    return new Promise((resolve) => {
+      if (!this.workerProc) {
+        resolve({ id, error: 'Worker not initialized' });
+        return;
+      }
+      this.pendingRequests.set(id, { resolve });
+      this.workerProc.stdin.write(JSON.stringify({ id, texts }) + '\n');
+      this.workerProc.stdin.flush();
+    });
+  }
+
   private async ensureInitialized(): Promise<void> {
-    if (this.extractor) return;
+    if (this.workerProc) return;
     await this.initGuard.run(() => this.initialize());
   }
 
   private async initialize(): Promise<void> {
-    log.info({ model: this.model }, 'Loading local embedding model (first load downloads the model)');
-
-    // In compiled Bun binaries, bare specifier resolution for packages with
-    // subdirectory entry points (like onnxruntime-common's dist/esm/index.js)
-    // fails. Additionally, CJS/ESM dual-instance issues cause onnxruntime-node's
-    // backend registration to be invisible to transformers. To solve both, the
-    // build step pre-bundles all JS deps into a single file placed inside
-    // onnxruntime-node/dist/ so native .node binary relative paths resolve.
-    const execDir = dirname(process.execPath);
-    const bundlePath = join(execDir, 'node_modules', 'onnxruntime-node', 'dist', 'transformers-bundle.mjs');
-    let transformers: typeof import('@huggingface/transformers');
+    log.info({ model: this.model }, 'Initializing local embedding backend');
+
+    const runtimeManager = new EmbeddingRuntimeManager();
+
+    // Wait for download if in progress
+    if (!runtimeManager.isReady()) {
+      log.info('Embedding runtime not yet available, waiting for download...');
+      await runtimeManager.ensureInstalled();
+    }
+
+    const bunPath = runtimeManager.getBunPath();
+    const workerPath = runtimeManager.getWorkerPath();
+
+    if (!bunPath) {
+      throw new Error('Local embedding backend unavailable: no bun binary found');
+    }
+    if (!existsSync(workerPath)) {
+      throw new Error(`Local embedding backend unavailable: worker script not found at ${workerPath}`);
+    }
+
+    await this.startWorker(bunPath, workerPath);
+  }
+
+  private async startWorker(bunPath: string, workerPath: string): Promise<void> {
+    const embeddingModelsDir = getEmbeddingModelsDir();
+    const modelCacheDir = `${embeddingModelsDir}/model-cache`;
+
+    log.info({ bunPath, workerPath, model: this.model }, 'Spawning embedding worker process');
+
+    const proc = Bun.spawn({
+      cmd: [bunPath, workerPath, this.model, modelCacheDir],
+      stdin: 'pipe',
+      stdout: 'pipe',
+      stderr: 'pipe',
+      cwd: embeddingModelsDir,
+    });
+
+    // Type-compatible assignment
+    this.workerProc = proc;
+
+    // Start reading stdout for responses (needed for waitForReady)
+    this.startStdoutReader();
+
     try {
-      transformers = await import(bundlePath);
-    } catch {
-      // Fall back to bare specifier for dev mode (running via `bun run`, not compiled)
+      // Wait for the worker to signal it's ready (model loaded)
+      await this.waitForReady();
+    } catch (err) {
+      // Worker failed to start — collect stderr for diagnosis
+      this.workerProc = null;
+      const exitCode = await proc.exited.catch(() => undefined);
+      const stderr = await new Response(proc.stderr).text().catch(() => '');
+      if (stderr.trim()) {
+        log.warn({ stderr: stderr.trim(), exitCode, bunPath }, 'Embedding worker stderr');
+      }
+      throw new Error(
+        `Embedding worker exited (code ${exitCode ?? 'unknown'}): ${stderr.trim() || (err instanceof Error ? err.message : String(err))}`,
+      );
+    }
+
+    // Worker is running — drain stderr in background for ongoing logging
+    this.drainStderr(proc.stderr);
+
+    // Write PID file so `vellum ps` can see the embed worker
+    this.writePidFile(proc.pid);
+
+    log.info({ pid: proc.pid, model: this.model }, 'Embedding worker process started');
+  }
+
+  private drainStderr(stderr: ReadableStream<Uint8Array>): void {
+    const reader = stderr.getReader();
+    const decoder = new TextDecoder();
+    (async () => {
       try {
-        transformers = await import('@huggingface/transformers');
-      } catch (err) {
-        throw new Error(
-          `Local embedding backend unavailable: failed to load @huggingface/transformers (${err instanceof Error ? err.message : String(err)})`,
-        );
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          const text = decoder.decode(value, { stream: true }).trim();
+          if (text) log.debug({ workerStderr: text }, 'Embedding worker stderr');
+        }
+      } catch {
+        // Reader cancelled or stream errored — expected on shutdown
+      }
+    })();
+  }
+
+  private startStdoutReader(): void {
+    if (this.stdoutReaderActive || !this.workerProc) return;
+    this.stdoutReaderActive = true;
+
+    const reader = this.workerProc.stdout.getReader();
+    const decoder = new TextDecoder();
+
+    (async () => {
+      try {
+        while (true) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          this.stdoutBuffer += decoder.decode(value, { stream: true });
+          this.processStdoutBuffer();
+        }
+      } catch {
+        // Reader cancelled or stream errored
+      }
+
+      // Worker exited — reject all pending requests and clean up
+      for (const [, pending] of this.pendingRequests) {
+        pending.resolve({ error: 'Embedding worker process exited unexpectedly' });
+      }
+      this.pendingRequests.clear();
+      this.workerProc = null;
+      this.stdoutReaderActive = false;
+      this.removePidFile();
+      this.stdoutBuffer = '';
+      // Allow re-initialization on next embed() call
+      this.initGuard.reset();
+    })();
+  }
+
+  private readyResolve: (() => void) | null = null;
+  private readyReject: ((err: Error) => void) | null = null;
+
+  private processStdoutBuffer(): void {
+    let idx: number;
+    while ((idx = this.stdoutBuffer.indexOf('\n')) !== -1) {
+      const line = this.stdoutBuffer.slice(0, idx);
+      this.stdoutBuffer = this.stdoutBuffer.slice(idx + 1);
+      if (!line.trim()) continue;
+
+      let msg: WorkerResponse;
+      try {
+        msg = JSON.parse(line);
+      } catch {
+        continue; // Skip malformed lines
+      }
+
+      // Handle ready/error signals during initialization
+      if (msg.type === 'ready') {
+        this.readyResolve?.();
+        this.readyResolve = null;
+        this.readyReject = null;
+        continue;
+      }
+      if (msg.type === 'error' && this.readyReject) {
+        this.readyReject(new Error(msg.error ?? 'Worker initialization failed'));
+        this.readyResolve = null;
+        this.readyReject = null;
+        continue;
+      }
+
+      // Handle embed responses
+      if (msg.id !== undefined) {
+        const pending = this.pendingRequests.get(msg.id);
+        if (pending) {
+          this.pendingRequests.delete(msg.id);
+          pending.resolve(msg);
+        }
       }
     }
+  }
+
+  private waitForReady(): Promise<void> {
+    return new Promise<void>((resolve, reject) => {
+      this.readyResolve = resolve;
+      this.readyReject = reject;
+
+      // Timeout after 2 minutes (first model download can be slow)
+      const timeout = setTimeout(() => {
+        this.readyResolve = null;
+        this.readyReject = null;
+        reject(new Error('Embedding worker timed out waiting for model to load'));
+      }, 120_000);
+
+      // Clear timeout when resolved
+      const originalResolve = resolve;
+      this.readyResolve = () => {
+        clearTimeout(timeout);
+        originalResolve();
+      };
+      const originalReject = reject;
+      this.readyReject = (err: Error) => {
+        clearTimeout(timeout);
+        originalReject(err);
+      };
+
+      // Also handle early worker exit
+      this.workerProc?.exited.then(() => {
+        if (this.readyResolve) {
+          clearTimeout(timeout);
+          this.readyResolve = null;
+          this.readyReject = null;
+          reject(new Error('Embedding worker process exited before becoming ready'));
+        }
+      });
+    });
+  }
+
+  private static readonly PID_FILENAME = 'embed-worker.pid';
 
-    this.extractor = await transformers.pipeline('feature-extraction', this.model, {
-      dtype: 'fp32',
-    }) as unknown as FeatureExtractionPipeline;
-    log.info({ model: this.model }, 'Local embedding model loaded');
+  private writePidFile(pid: number): void {
+    try {
+      writeFileSync(join(getRootDir(), LocalEmbeddingBackend.PID_FILENAME), String(pid));
+    } catch {
+      // Best-effort — doesn't affect functionality
+    }
+  }
+
+  private removePidFile(): void {
+    try {
+      unlinkSync(join(getRootDir(), LocalEmbeddingBackend.PID_FILENAME));
+    } catch {
+      // Best-effort
+    }
   }
 }