@aeriondyseti/vector-memory-mcp 2.3.0-rc.2 → 2.3.0-rc.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aeriondyseti/vector-memory-mcp",
-  "version": "2.3.0-rc.2",
+  "version": "2.3.0-rc.4",
   "description": "A zero-configuration RAG memory server for MCP clients",
   "type": "module",
   "main": "server/index.ts",
@@ -47,18 +47,18 @@
   ],
   "license": "MIT",
   "dependencies": {
-    "@huggingface/transformers": "^3.8.0",
+    "@huggingface/tokenizers": "^0.1.3",
     "@lancedb/lancedb": "^0.26.2",
     "@modelcontextprotocol/sdk": "^1.0.0",
     "arg": "^5.0.2",
-    "hono": "^4.11.3"
+    "hono": "^4.11.3",
+    "onnxruntime-node": "^1.21.0"
   },
   "devDependencies": {
     "@types/bun": "latest",
     "typescript": "^5.0.0"
   },
   "trustedDependencies": [
-    "protobufjs",
-    "sharp"
+    "protobufjs"
   ]
 }

package/scripts/warmup.ts

@@ -23,20 +23,11 @@ async function warmup(): Promise<void> {
       process.exit(1);
     }
 
-    try {
-      await import("sharp");
-      console.log("  ✓ sharp loaded");
-    } catch (e) {
-      console.error("  ✗ sharp failed:", (e as Error).message);
-      process.exit(1);
-    }
-
     console.log();
 
     // Initialize embeddings service to download model
     console.log("📥 Downloading ML model (this may take a minute)...");
     console.log(`   Model: ${config.embeddingModel}`);
-    console.log(`   Cache: ~/.cache/huggingface/`);
     console.log();
 
     const embeddings = new EmbeddingsService(

package/server/core/embeddings.service.ts

@@ -1,9 +1,17 @@
-import { pipeline, type FeatureExtractionPipeline } from "@huggingface/transformers";
+import * as ort from "onnxruntime-node";
+import { Tokenizer } from "@huggingface/tokenizers";
+import { join, dirname } from "path";
+import { mkdir } from "fs/promises";
+import { existsSync } from "fs";
+
+const HF_CDN = "https://huggingface.co";
+const MAX_SEQ_LENGTH = 512;
 
 export class EmbeddingsService {
   private modelName: string;
-  private extractor: FeatureExtractionPipeline | null = null;
-  private initPromise: Promise<FeatureExtractionPipeline> | null = null;
+  private session: ort.InferenceSession | null = null;
+  private tokenizer: Tokenizer | null = null;
+  private initPromise: Promise<void> | null = null;
   private _dimension: number;
 
   constructor(modelName: string, dimension: number) {
@@ -15,27 +23,71 @@ export class EmbeddingsService {
     return this._dimension;
   }
 
-  private async getExtractor(): Promise<FeatureExtractionPipeline> {
-    if (this.extractor) {
-      return this.extractor;
-    }
-
+  private async initialize(): Promise<void> {
+    if (this.session) return;
     if (!this.initPromise) {
-      this.initPromise = pipeline(
-        "feature-extraction",
-        this.modelName,
-        { dtype: "fp32" } as any
-      ) as Promise<FeatureExtractionPipeline>;
+      this.initPromise = this._init();
     }
+    await this.initPromise;
+  }
+
+  private get cacheDir(): string {
+    const packageRoot = join(dirname(Bun.main), "..");
+    return join(packageRoot, ".cache", "models", this.modelName);
+  }
 
-    this.extractor = await this.initPromise;
-    return this.extractor;
+  private async downloadIfMissing(fileName: string): Promise<string> {
+    const filePath = join(this.cacheDir, fileName);
+    if (existsSync(filePath)) return filePath;
+
+    const url = `${HF_CDN}/${this.modelName}/resolve/main/${fileName}`;
+    await mkdir(dirname(filePath), { recursive: true });
+    const response = await fetch(url);
+    if (!response.ok) throw new Error(`Failed to download ${url}: ${response.status}`);
+    const buffer = await response.arrayBuffer();
+    await Bun.write(filePath, buffer);
+    return filePath;
+  }
+
+  private async _init(): Promise<void> {
+    const modelPath = await this.downloadIfMissing("onnx/model.onnx");
+    const tokenizerJsonPath = await this.downloadIfMissing("tokenizer.json");
+    const tokenizerConfigPath = await this.downloadIfMissing("tokenizer_config.json");
+
+    this.session = await ort.InferenceSession.create(modelPath, {
+      executionProviders: ["cpu"],
+    });
+
+    const tokenizerJson = await Bun.file(tokenizerJsonPath).json();
+    const tokenizerConfig = await Bun.file(tokenizerConfigPath).json();
+    this.tokenizer = new Tokenizer(tokenizerJson, tokenizerConfig);
   }
 
   async embed(text: string): Promise<number[]> {
-    const extractor = await this.getExtractor();
-    const output = await extractor(text, { pooling: "mean", normalize: true });
-    return Array.from(output.data as Float32Array);
+    await this.initialize();
+
+    const encoded = this.tokenizer!.encode(text);
+
+    // Truncate to model's max sequence length
+    const seqLen = Math.min(encoded.ids.length, MAX_SEQ_LENGTH);
+    const ids = encoded.ids.slice(0, seqLen);
+    const mask = encoded.attention_mask.slice(0, seqLen);
+
+    const inputIds = BigInt64Array.from(ids.map(BigInt));
+    const attentionMask = BigInt64Array.from(mask.map(BigInt));
+    const tokenTypeIds = new BigInt64Array(seqLen); // zeros for single-sequence input
+
+    const feeds: Record<string, ort.Tensor> = {
+      input_ids: new ort.Tensor("int64", inputIds, [1, seqLen]),
+      attention_mask: new ort.Tensor("int64", attentionMask, [1, seqLen]),
+      token_type_ids: new ort.Tensor("int64", tokenTypeIds, [1, seqLen]),
+    };
+
+    const output = await this.session!.run(feeds);
+    const lastHidden = output["last_hidden_state"];
+
+    const pooled = this.meanPool(lastHidden.data as Float32Array, mask, seqLen);
+    return this.normalize(pooled);
   }
 
   async embedBatch(texts: string[]): Promise<number[][]> {
@@ -45,4 +97,29 @@ export class EmbeddingsService {
     }
     return results;
   }
+
+  private meanPool(data: Float32Array, mask: number[], seqLen: number): number[] {
+    const dim = this._dimension;
+    const pooled = new Array(dim).fill(0);
+    let maskSum = 0;
+    for (let t = 0; t < seqLen; t++) {
+      if (mask[t]) {
+        maskSum += 1;
+        for (let d = 0; d < dim; d++) {
+          pooled[d] += data[t * dim + d];
+        }
+      }
+    }
+    for (let d = 0; d < dim; d++) {
+      pooled[d] /= maskSum;
+    }
+    return pooled;
+  }
+
+  private normalize(vec: number[]): number[] {
+    let norm = 0;
+    for (const v of vec) norm += v * v;
+    norm = Math.sqrt(norm);
+    return vec.map(v => v / norm);
+  }
 }

package/server/core/migrations.ts

@@ -127,8 +127,43 @@ export async function backfillVectors(
   db: Database,
   embeddings: EmbeddingsService,
 ): Promise<void> {
+  // Fast sentinel check: skip the LEFT JOIN queries entirely when backfill is done
+  const sentinel = db
+    .prepare("SELECT 1 FROM memories_vec LIMIT 1")
+    .get();
+  const memoriesExist = db.prepare("SELECT 1 FROM memories LIMIT 1").get();
+  const convosExist = db.prepare("SELECT 1 FROM conversation_history LIMIT 1").get();
+
+  // If vec tables have data and source tables have data, backfill is likely complete.
+  // Only run the expensive LEFT JOIN when there's reason to suspect gaps.
+  const convoSentinel = db
+    .prepare("SELECT 1 FROM conversation_history_vec LIMIT 1")
+    .get();
+  const mayNeedMemoryBackfill = memoriesExist && !sentinel;
+  const mayNeedConvoBackfill = convosExist && !convoSentinel;
+
+  // If both vec tables are populated, do a quick count check to confirm
+  if (!mayNeedMemoryBackfill && !mayNeedConvoBackfill) {
+    if (memoriesExist) {
+      const gap = db.prepare(
+        `SELECT 1 FROM memories m LEFT JOIN memories_vec v ON m.id = v.id
+         WHERE v.id IS NULL OR length(v.vector) = 0 LIMIT 1`,
+      ).get();
+      if (!gap && convosExist) {
+        const convoGap = db.prepare(
+          `SELECT 1 FROM conversation_history c LEFT JOIN conversation_history_vec v ON c.id = v.id
+           WHERE v.id IS NULL OR length(v.vector) = 0 LIMIT 1`,
+        ).get();
+        if (!convoGap) return;
+      } else if (!gap && !convosExist) {
+        return;
+      }
+    } else {
+      return; // No data at all
+    }
+  }
+
   // ── Memories ──────────────────────────────────────────────────────
-  // Catch both missing rows (v.id IS NULL) and corrupt 0-byte BLOBs
   const missingMemories = db
     .prepare(
       `SELECT m.id, m.content, json_extract(m.metadata, '$.type') AS type
@@ -151,14 +186,27 @@ export async function backfillVectors(
       new Array(embeddings.dimension).fill(0),
     );
 
-    for (const row of missingMemories) {
-      // Waypoints use a zero vector (not semantically searched)
-      const blob =
-        row.type === "waypoint"
-          ? zeroVector
-          : serializeVector(await embeddings.embed(row.content));
+    // Separate waypoints from content that needs embedding
+    const toEmbed = missingMemories.filter((r) => r.type !== "waypoint");
+    const waypoints = missingMemories.filter((r) => r.type === "waypoint");
+
+    // Batch embed all non-waypoint content
+    const vectors = toEmbed.length > 0
+      ? await embeddings.embedBatch(toEmbed.map((r) => r.content))
+      : [];
 
-      insertVec.run(row.id, blob);
+    db.exec("BEGIN");
+    try {
+      for (const row of waypoints) {
+        insertVec.run(row.id, zeroVector);
+      }
+      for (let i = 0; i < toEmbed.length; i++) {
+        insertVec.run(toEmbed[i].id, serializeVector(vectors[i]));
+      }
+      db.exec("COMMIT");
+    } catch (e) {
+      db.exec("ROLLBACK");
+      throw e;
     }
 
     console.error(
@@ -185,17 +233,27 @@ export async function backfillVectors(
       "INSERT OR REPLACE INTO conversation_history_vec (id, vector) VALUES (?, ?)",
     );
 
-    for (let i = 0; i < missingConvos.length; i++) {
-      const row = missingConvos[i];
-      const vec = serializeVector(await embeddings.embed(row.content));
-      insertConvoVec.run(row.id, vec);
+    // Batch embed in chunks of 32
+    const BATCH_SIZE = 32;
+    db.exec("BEGIN");
+    try {
+      for (let i = 0; i < missingConvos.length; i += BATCH_SIZE) {
+        const batch = missingConvos.slice(i, i + BATCH_SIZE);
+        const vecs = await embeddings.embedBatch(batch.map((r) => r.content));
+        for (let j = 0; j < batch.length; j++) {
+          insertConvoVec.run(batch[j].id, serializeVector(vecs[j]));
+        }
 
-      // Log progress every 100 chunks
-      if ((i + 1) % 100 === 0) {
-        console.error(
-          `[vector-memory-mcp]   ...${i + 1}/${missingConvos.length} conversation chunks`,
-        );
+        if ((i + BATCH_SIZE) % 100 < BATCH_SIZE) {
+          console.error(
+            `[vector-memory-mcp]   ...${Math.min(i + BATCH_SIZE, missingConvos.length)}/${missingConvos.length} conversation chunks`,
+          );
+        }
       }
+      db.exec("COMMIT");
+    } catch (e) {
+      db.exec("ROLLBACK");
+      throw e;
     }
 
     console.error(

package/server/index.ts

@@ -25,17 +25,15 @@ async function main(): Promise<void> {
   const overrides = parseCliArgs(args);
   const config = loadConfig(overrides);
 
-  // Initialize database
+  // Initialize database and backfill any missing vectors before services start
   const db = connectToDatabase(config.dbPath);
+  const embeddings = new EmbeddingsService(config.embeddingModel, config.embeddingDimension);
+  await backfillVectors(db, embeddings);
 
   // Initialize layers
   const repository = new MemoryRepository(db);
-  const embeddings = new EmbeddingsService(config.embeddingModel, config.embeddingDimension);
   const memoryService = new MemoryService(repository, embeddings);
 
-  // Backfill any missing vectors (e.g. after vec0-to-BLOB migration)
-  await backfillVectors(db, embeddings);
-
   if (config.pluginMode) {
     console.error("[vector-memory-mcp] Running in plugin mode");
   }