feat(embedding): make the local model configurable

.env.example

@@ -88,9 +88,14 @@
 # BM25-only. Local embeddings are an explicit opt-in, not the keyless default.
 # `EMBEDDING_PROVIDER=local` uses Xenova/all-MiniLM-L6-v2 (384-dim) on-device;
 # the first embedding request downloads the model and needs network access.
+# That default model is English-only: most CJK text tokenizes to [UNK], so
+# non-English stores should point it at a multilingual model instead.
 
 # EMBEDDING_PROVIDER=local                       # local | openai | voyage | cohere | gemini | openrouter
 
+# AGENTMEMORY_LOCAL_EMBEDDING_MODEL=Xenova/bge-small-zh-v1.5   # any Transformers.js feature-extraction model
+# AGENTMEMORY_LOCAL_EMBEDDING_DIMENSIONS=512     # override when the model is not in the known-models table
+
 # VOYAGE_API_KEY=pa-...                          # Optimised for code embeddings
 
 # COHERE_API_KEY=...                             # General-purpose embeddings

scripts/repro-cjk-embedding.mjs

@@ -0,0 +1,60 @@
+#!/usr/bin/env node
+/**
+ * Repro: the default local embedding model is English-only, so CJK text
+ * tokenizes to [UNK] and unrelated records collapse to near-identical vectors.
+ *
+ *   node scripts/repro-cjk-embedding.mjs                 # default model
+ *   node scripts/repro-cjk-embedding.mjs Xenova/bge-small-zh-v1.5
+ *
+ * Requires @huggingface/transformers (already a dependency for local embeddings).
+ */
+import { AutoTokenizer, pipeline } from "@huggingface/transformers";
+
+const model = process.argv[2] || "Xenova/all-MiniLM-L6-v2";
+
+const sentences = [
+  "量化交易策略",
+  "记忆系统索引",
+  "回测框架",
+  "今天天气很好",
+];
+
+const tokenizer = await AutoTokenizer.from_pretrained(model);
+
+console.log(`model: ${model}\n`);
+for (const text of sentences) {
+  const ids = (await tokenizer(text)).input_ids.tolist()[0];
+  const unkCount = ids.filter(
+    (id) => tokenizer.decode([id], { skip_special_tokens: false }) === "[UNK]",
+  ).length;
+  const unkRatio = unkCount / ids.length;
+  console.log(`${text.padEnd(10)} tokens=${ids.length} [UNK]=${(unkRatio * 100).toFixed(0)}%`);
+}
+
+const extractor = await pipeline("feature-extraction", model, { dtype: "q8" });
+const vectors = [];
+for (const text of sentences) {
+  const out = await extractor([text], { pooling: "mean", normalize: true });
+  vectors.push(out.tolist()[0]);
+}
+
+const cosine = (a, b) => {
+  let dot = 0;
+  let na = 0;
+  let nb = 0;
+  for (let i = 0; i < a.length; i++) {
+    dot += a[i] * b[i];
+    na += a[i] * a[i];
+    nb += b[i] * b[i];
+  }
+  return dot / Math.sqrt(na * nb);
+};
+
+console.log("\npairwise cosine (unrelated topics should be low):");
+for (let i = 0; i < vectors.length; i++) {
+  for (let j = i + 1; j < vectors.length; j++) {
+    console.log(
+      `  ${sentences[i]} <-> ${sentences[j]}: ${cosine(vectors[i], vectors[j]).toFixed(3)}`,
+    );
+  }
+}

src/providers/embedding/local.ts

@@ -1,15 +1,66 @@
 import type { EmbeddingProvider } from "../../types.js";
+import { getEnvVar } from "../../config.js";
 
 type FeatureExtractor = (
   texts: string[],
   options: { pooling: string; normalize: boolean },
 ) => Promise<{ tolist: () => number[][] }>;
 
+/**
+ * Default on-device model. English-only: its WordPiece vocabulary contains
+ * roughly 500 CJK characters, so most Chinese/Japanese/Korean text tokenizes
+ * to `[UNK]`. Set AGENTMEMORY_LOCAL_EMBEDDING_MODEL to a multilingual model
+ * (e.g. Xenova/bge-small-zh-v1.5, Xenova/multilingual-e5-small) for non-English
+ * stores.
+ */
+export const DEFAULT_LOCAL_MODEL = "Xenova/all-MiniLM-L6-v2";
+
+/**
+ * Output widths of the models we ship defaults for. The dimension guard in
+ * index.ts throws on a mismatch, so an unknown model must supply
+ * AGENTMEMORY_LOCAL_EMBEDDING_DIMENSIONS explicitly.
+ */
+const MODEL_DIMENSIONS: Record<string, number> = {
+  "Xenova/all-MiniLM-L6-v2": 384,
+  "Xenova/all-MiniLM-L12-v2": 384,
+  "Xenova/bge-small-zh-v1.5": 512,
+  "Xenova/bge-base-zh-v1.5": 768,
+  "Xenova/multilingual-e5-small": 384,
+  "Xenova/multilingual-e5-base": 768,
+  "Xenova/multilingual-e5-large": 1024,
+};
+
+const DEFAULT_DIMENSIONS = MODEL_DIMENSIONS[DEFAULT_LOCAL_MODEL];
+
+function parseDimensions(raw: string | undefined): number | undefined {
+  if (raw === undefined || raw.trim().length === 0) return undefined;
+  const parsed = parseInt(raw, 10);
+  if (!Number.isFinite(parsed) || parsed <= 0) {
+    throw new Error(
+      `AGENTMEMORY_LOCAL_EMBEDDING_DIMENSIONS must be a positive integer, got: ${raw}`,
+    );
+  }
+  return parsed;
+}
+
 export class LocalEmbeddingProvider implements EmbeddingProvider {
   readonly name = "local";
-  readonly dimensions = 384;
+  readonly dimensions: number;
+  readonly model: string;
   private extractor: FeatureExtractor | null = null;
 
+  constructor(options: { model?: string; dimensions?: number } = {}) {
+    this.model =
+      options.model?.trim() ||
+      getEnvVar("AGENTMEMORY_LOCAL_EMBEDDING_MODEL")?.trim() ||
+      DEFAULT_LOCAL_MODEL;
+    this.dimensions =
+      options.dimensions ??
+      parseDimensions(getEnvVar("AGENTMEMORY_LOCAL_EMBEDDING_DIMENSIONS")) ??
+      MODEL_DIMENSIONS[this.model] ??
+      DEFAULT_DIMENSIONS;
+  }
+
   async embed(text: string): Promise<Float32Array> {
     const [result] = await this.embedBatch([text]);
     return result;
@@ -39,7 +90,7 @@ export class LocalEmbeddingProvider implements EmbeddingProvider {
     }
     this.extractor = (await transformers.pipeline(
       "feature-extraction",
-      "Xenova/all-MiniLM-L6-v2",
+      this.model,
       { dtype: "q8" },
     )) as FeatureExtractor;
     return this.extractor;

test/local-embedding-provider.test.ts

@@ -2,16 +2,20 @@ import { describe, it, expect, vi, afterEach } from "vitest";
 
 afterEach(() => {
   vi.doUnmock("@huggingface/transformers");
+  vi.unstubAllEnvs();
   vi.resetModules();
 });
 
+async function loadFresh() {
+  vi.resetModules();
+  const mod = await import("../src/providers/embedding/local.js");
+  return mod.LocalEmbeddingProvider;
+}
+
 describe("LocalEmbeddingProvider (package unavailable)", () => {
   it("throws clean install hint when @huggingface/transformers is missing", async () => {
     vi.doMock("@huggingface/transformers");
-    vi.resetModules();
-    const { LocalEmbeddingProvider: Fresh } = await import(
-      "../src/providers/embedding/local.js"
-    );
+    const Fresh = await loadFresh();
     await expect(new Fresh().embed("hello")).rejects.toThrow(
       "Install @huggingface/transformers for local embeddings",
     );
@@ -25,15 +29,12 @@ describe("LocalEmbeddingProvider (with loaded pipeline)", () => {
     }));
     const pipeline = vi.fn(() => Promise.resolve(extractor));
     vi.doMock("@huggingface/transformers", () => ({ pipeline }));
-    vi.resetModules();
     return { pipeline, extractor };
   }
 
   it("calls pipeline with dtype: q8, passes extractor opts, returns mapped Float32Array", async () => {
     const { pipeline, extractor } = mockSuccessModule();
-    const { LocalEmbeddingProvider: Fresh } = await import(
-      "../src/providers/embedding/local.js"
-    );
+    const Fresh = await loadFresh();
     const vec = await new Fresh().embed("hello");
 
     expect(pipeline).toHaveBeenCalledWith(
@@ -51,12 +52,71 @@ describe("LocalEmbeddingProvider (with loaded pipeline)", () => {
 
   it("embedBatch returns one Float32Array per input text", async () => {
     mockSuccessModule();
-    const { LocalEmbeddingProvider: Fresh } = await import(
-      "../src/providers/embedding/local.js"
-    );
+    const Fresh = await loadFresh();
     const vecs = await new Fresh().embedBatch(["a", "b", "c"]);
 
     expect(vecs).toHaveLength(3);
     for (const v of vecs) expect(v).toBeInstanceOf(Float32Array);
   });
 });
+
+describe("LocalEmbeddingProvider (model configuration)", () => {
+  it("defaults to the English MiniLM model at 384 dims", async () => {
+    const Fresh = await loadFresh();
+    const provider = new Fresh();
+    expect(provider.model).toBe("Xenova/all-MiniLM-L6-v2");
+    expect(provider.dimensions).toBe(384);
+  });
+
+  it("reads model and infers dimensions from AGENTMEMORY_LOCAL_EMBEDDING_MODEL", async () => {
+    vi.stubEnv("AGENTMEMORY_LOCAL_EMBEDDING_MODEL", "Xenova/bge-small-zh-v1.5");
+    const Fresh = await loadFresh();
+    const provider = new Fresh();
+    expect(provider.model).toBe("Xenova/bge-small-zh-v1.5");
+    expect(provider.dimensions).toBe(512);
+  });
+
+  it("honours an explicit AGENTMEMORY_LOCAL_EMBEDDING_DIMENSIONS override", async () => {
+    vi.stubEnv("AGENTMEMORY_LOCAL_EMBEDDING_MODEL", "Xenova/some-unknown-model");
+    vi.stubEnv("AGENTMEMORY_LOCAL_EMBEDDING_DIMENSIONS", "768");
+    const Fresh = await loadFresh();
+    const provider = new Fresh();
+    expect(provider.dimensions).toBe(768);
+  });
+
+  it("rejects a non-positive dimensions override", async () => {
+    vi.stubEnv("AGENTMEMORY_LOCAL_EMBEDDING_DIMENSIONS", "0");
+    const Fresh = await loadFresh();
+    expect(() => new Fresh()).toThrow(
+      "AGENTMEMORY_LOCAL_EMBEDDING_DIMENSIONS must be a positive integer",
+    );
+  });
+
+  it("passes the configured model id to the pipeline", async () => {
+    vi.stubEnv("AGENTMEMORY_LOCAL_EMBEDDING_MODEL", "Xenova/bge-small-zh-v1.5");
+    const { pipeline } = (() => {
+      const extractor = vi.fn(async (texts: string[]) => ({
+        tolist: () => texts.map(() => new Array(512).fill(0.01)),
+      }));
+      const pipeline = vi.fn(() => Promise.resolve(extractor));
+      vi.doMock("@huggingface/transformers", () => ({ pipeline }));
+      return { pipeline };
+    })();
+    const Fresh = await loadFresh();
+    await new Fresh().embed("记忆系统");
+
+    expect(pipeline).toHaveBeenCalledWith(
+      "feature-extraction",
+      "Xenova/bge-small-zh-v1.5",
+      { dtype: "q8" },
+    );
+  });
+
+  it("allows constructor options to override the environment", async () => {
+    vi.stubEnv("AGENTMEMORY_LOCAL_EMBEDDING_MODEL", "Xenova/bge-small-zh-v1.5");
+    const Fresh = await loadFresh();
+    const provider = new Fresh({ model: "Xenova/multilingual-e5-base", dimensions: 768 });
+    expect(provider.model).toBe("Xenova/multilingual-e5-base");
+    expect(provider.dimensions).toBe(768);
+  });
+});