@aeriondyseti/vector-memory-mcp 2.3.0 → 2.4.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aeriondyseti/vector-memory-mcp",
-  "version": "2.3.0",
+  "version": "2.4.0",
   "description": "A zero-configuration RAG memory server for MCP clients",
   "type": "module",
   "main": "server/index.ts",
@@ -9,7 +9,6 @@
   },
   "files": [
     "server",
-    "scripts",
     "README.md",
     "LICENSE"
   ],
@@ -31,6 +30,7 @@
     "test:quick": "bun test",
     "test:coverage": "bun test --preload ./tests/preload.ts --coverage",
     "benchmark": "bun test tests/benchmark.test.ts --preload ./tests/preload.ts",
+    "benchmark:update": "bun run scripts/update-benchmarks.ts",
     "test:preload": "bun run tests/preload.ts",
     "smoke": "bun run scripts/smoke-test.ts",
     "warmup": "bun run scripts/warmup.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/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,79 @@ export class EmbeddingsService {
     return this._dimension;
   }
 
-  private async getExtractor(): Promise<FeatureExtractionPipeline> {
-    if (this.extractor) {
-      return this.extractor;
-    }
+  get isReady(): boolean {
+    return this.session !== null;
+  }
 
+  async warmup(): Promise<void> {
+    await this.initialize();
+  }
+
+  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);
+  }
+
+  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"],
+    });
 
-    this.extractor = await this.initPromise;
-    return this.extractor;
+    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 +105,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");
   }

package/server/transports/http/server.ts

@@ -111,10 +111,22 @@ export function createHttpApp(memoryService: MemoryService, config: Config): Hon
         embeddingDimension: config.embeddingDimension,
         historyEnabled: config.conversationHistory.enabled,
         pluginMode: config.pluginMode,
+        embeddingReady: memoryService.getEmbeddings().isReady,
       },
     });
   });
 
+  // Warmup endpoint — triggers ONNX model load if not already cached
+  app.post("/warmup", async (c) => {
+    const embeddings = memoryService.getEmbeddings();
+    if (embeddings.isReady) {
+      return c.json({ status: "already_warm" });
+    }
+    const start = Date.now();
+    await embeddings.warmup();
+    return c.json({ status: "warmed", elapsed: Date.now() - start });
+  });
+
   // Search endpoint
   app.post("/search", async (c) => {
     try {

package/scripts/lancedb-extract.ts

@@ -1,181 +0,0 @@
-#!/usr/bin/env bun
-/**
- * Standalone LanceDB data extractor — runs in a child process so that
- * @lancedb/lancedb native bindings never coexist with bun:sqlite's
- * extension loading in the same process.
- *
- * Usage: bun scripts/lancedb-extract.ts <lance-db-path>
- * Output: JSON on stdout — { memories: Row[], conversations: Row[] }
- */
-
-const source = process.argv[2];
-if (!source) {
-  console.error("Usage: bun scripts/lancedb-extract.ts <lance-db-path>");
-  process.exit(1);
-}
-
-// Arrow TimeUnit enum → divisor to convert to milliseconds.
-// 0=SECOND, 1=MILLISECOND, 2=MICROSECOND, 3=NANOSECOND
-// Negative divisor = multiply (seconds → ms needs ×1000).
-const TIME_UNIT_TO_MS_DIVISOR: Record<number, bigint> = {
-  0: -1000n,   // seconds → ms (multiply by 1000)
-  1: 1n,       // ms → no conversion
-  2: 1000n,    // μs → ms
-  3: 1000000n, // ns → ms
-};
-
-function buildTimestampDivisors(schema: any): Map<string, bigint> {
-  const map = new Map<string, bigint>();
-  for (const field of schema.fields) {
-    if (field.type.typeId === 10) {
-      map.set(field.name, TIME_UNIT_TO_MS_DIVISOR[field.type.unit] ?? 1n);
-    }
-  }
-  return map;
-}
-
-function columnValue(batch: any, colName: string, rowIdx: number): unknown {
-  const col = batch.getChild(colName);
-  if (!col) return undefined;
-  try {
-    return col.get(rowIdx);
-  } catch {
-    // Arrow's getter can throw on BigInt timestamps exceeding MAX_SAFE_INTEGER;
-    // fall back to the raw typed array.
-    let offset = rowIdx;
-    for (const data of col.data) {
-      if (offset < data.length) {
-        return (data.values instanceof BigInt64Array || data.values instanceof BigUint64Array)
-          ? data.values[offset]
-          : null;
-      }
-      offset -= data.length;
-    }
-    return null;
-  }
-}
-
-function toEpochMs(value: unknown, divisor: bigint = 1n): number {
-  if (value == null) return Date.now();
-  if (value instanceof Date) return value.getTime();
-  if (typeof value === "bigint") {
-    if (divisor < 0n) return Number(value * -divisor);  // seconds → ms
-    if (divisor === 1n) return Number(value);
-    return Number(value / divisor);
-  }
-  if (typeof value === "number") {
-    if (divisor < 0n) return value * Number(-divisor);
-    if (divisor === 1n) return value;
-    return Math.floor(value / Number(divisor));
-  }
-  return Date.now();
-}
-
-function toFloatArray(vec: unknown): number[] {
-  if (Array.isArray(vec)) return vec;
-  if (vec instanceof Float32Array) return Array.from(vec);
-  if (vec && typeof (vec as any).toArray === "function") {
-    return Array.from((vec as any).toArray());
-  }
-  if (ArrayBuffer.isView(vec)) {
-    const view = vec as DataView;
-    return Array.from(new Float32Array(view.buffer, view.byteOffset, view.byteLength / 4));
-  }
-  return [];
-}
-
-const BATCH_SIZE = 100;
-const lancedb = await import("@lancedb/lancedb");
-const db = await lancedb.connect(source);
-const tableNames = await db.tableNames();
-console.error(`Found tables: ${tableNames.join(", ")}`);
-
-const result: { memories: any[]; conversations: any[] } = {
-  memories: [],
-  conversations: [],
-};
-
-if (tableNames.includes("memories")) {
-  const table = await db.openTable("memories");
-  const total = await table.countRows();
-  console.error(`Reading ${total} memories...`);
-
-  // Paginated scan — query().toArrow() without offset/limit returns
-  // non-deterministic results that can duplicate some rows and skip others.
-  const schemaSample = await table.query().limit(1).toArrow();
-  const tsDivisors = buildTimestampDivisors(schemaSample.schema);
-  const seen = new Map<string, any>();
-
-  for (let offset = 0; offset < total; offset += BATCH_SIZE) {
-    const arrowTable = await table.query().offset(offset).limit(BATCH_SIZE).toArrow();
-    for (const batch of arrowTable.batches) {
-      for (let i = 0; i < batch.numRows; i++) {
-        const id = columnValue(batch, "id", i) as string;
-        const content = columnValue(batch, "content", i) as string;
-        const lastAccessed = columnValue(batch, "last_accessed", i);
-        const accessedMs = lastAccessed != null ? toEpochMs(lastAccessed, tsDivisors.get("last_accessed")) : null;
-        // Deduplicate by ID: prefer most recently accessed, then longest content.
-        const existing = seen.get(id);
-        if (existing) {
-          const existingAccess = existing.last_accessed ?? 0;
-          const newAccess = accessedMs ?? 0;
-          if (newAccess < existingAccess) continue;
-          if (newAccess === existingAccess && content.length <= existing.content.length) continue;
-        }
-        seen.set(id, {
-          id,
-          content,
-          metadata: columnValue(batch, "metadata", i) ?? "{}",
-          vector: toFloatArray(columnValue(batch, "vector", i)),
-          created_at: toEpochMs(columnValue(batch, "created_at", i), tsDivisors.get("created_at")),
-          updated_at: toEpochMs(columnValue(batch, "updated_at", i), tsDivisors.get("updated_at")),
-          last_accessed: accessedMs,
-          superseded_by: columnValue(batch, "superseded_by", i) ?? null,
-          usefulness: columnValue(batch, "usefulness", i) ?? 0,
-          access_count: columnValue(batch, "access_count", i) ?? 0,
-        });
-      }
-    }
-  }
-  result.memories = [...seen.values()];
-  console.error(`  ${result.memories.length} unique memories read (${total} rows scanned)`);
-}
-
-if (tableNames.includes("conversation_history")) {
-  const table = await db.openTable("conversation_history");
-  const total = await table.countRows();
-  console.error(`Reading ${total} conversation chunks...`);
-
-  const schemaSample = await table.query().limit(1).toArrow();
-  const tsDivisors = buildTimestampDivisors(schemaSample.schema);
-  const seen = new Map<string, any>();
-
-  for (let offset = 0; offset < total; offset += BATCH_SIZE) {
-    const arrowTable = await table.query().offset(offset).limit(BATCH_SIZE).toArrow();
-    for (const batch of arrowTable.batches) {
-      for (let i = 0; i < batch.numRows; i++) {
-        const id = columnValue(batch, "id", i) as string;
-        const content = columnValue(batch, "content", i) as string;
-        const existing = seen.get(id);
-        if (existing && existing.content.length >= content.length) continue;
-        seen.set(id, {
-          id,
-          content,
-          metadata: columnValue(batch, "metadata", i) ?? "{}",
-          vector: toFloatArray(columnValue(batch, "vector", i)),
-          created_at: toEpochMs(columnValue(batch, "created_at", i), tsDivisors.get("created_at")),
-          session_id: columnValue(batch, "session_id", i),
-          role: columnValue(batch, "role", i),
-          message_index_start: columnValue(batch, "message_index_start", i) ?? 0,
-          message_index_end: columnValue(batch, "message_index_end", i) ?? 0,
-          project: columnValue(batch, "project", i) ?? "",
-        });
-      }
-    }
-  }
-  result.conversations = [...seen.values()];
-  console.error(`  ${result.conversations.length} unique conversation chunks read (${total} rows scanned)`);
-}
-
-await db.close?.();
-process.stdout.write(JSON.stringify(result));