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

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@aeriondyseti/vector-memory-mcp",
-  "version": "2.3.0-rc.3",
+  "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);
+  }
 }