@amemhq/core 1.0.0 → 1.1.0

package/dist/index.cjs

@@ -33,6 +33,7 @@ __export(index_exports, {
   DEFAULT_CRUD_UPDATE_MIN_SIM: () => DEFAULT_CRUD_UPDATE_MIN_SIM,
   DEFAULT_EMBEDDING_MODEL: () => DEFAULT_EMBEDDING_MODEL,
   EmbeddingDimensionMismatchError: () => EmbeddingDimensionMismatchError,
+  EmbeddingModelMismatchError: () => EmbeddingModelMismatchError,
   addEpisodic: () => addEpisodic,
   addMemory: () => addMemory,
   canRead: () => canRead,
@@ -47,8 +48,11 @@ __export(index_exports, {
   encode: () => encode,
   ensureCollection: () => ensureCollection,
   generateReviewBatch: () => generateReviewBatch,
+  getEmbeddingDevice: () => getEmbeddingDevice,
   getEmbeddingDim: () => getEmbeddingDim,
+  getEmbeddingDtype: () => getEmbeddingDtype,
   getEmbeddingModel: () => getEmbeddingModel,
+  getEmbeddingPooling: () => getEmbeddingPooling,
   getNote: () => getNote,
   invalidateNote: () => invalidateNote,
   isModelLoaded: () => isModelLoaded,
@@ -82,47 +86,85 @@ function getDataDir() {
 // src/embedding.ts
 var pipeline = null;
 var extractor = null;
-var loadedModelName = null;
+var loadedKey = null;
 var cachedDim = null;
 var DEFAULT_EMBEDDING_MODEL = "Xenova/paraphrase-multilingual-MiniLM-L12-v2";
 function getEmbeddingModel() {
   return process.env.AMEM_EMBED_MODEL?.trim() || DEFAULT_EMBEDDING_MODEL;
 }
+var CLS_POOLED_MODELS = /* @__PURE__ */ new Set([
+  "bge-m3",
+  "bge-base-zh-v1.5",
+  "bge-small-zh-v1.5",
+  "bge-base-en-v1.5",
+  "bge-small-en-v1.5",
+  "bge-large-en-v1.5",
+  "gte-multilingual-base",
+  "gte-modernbert-base",
+  "gte-large-en-v1.5",
+  "snowflake-arctic-embed-m",
+  "snowflake-arctic-embed-l"
+]);
+function getEmbeddingPooling() {
+  const explicit = process.env.AMEM_EMBED_POOLING?.trim().toLowerCase();
+  if (explicit === "mean" || explicit === "cls") return explicit;
+  const basename2 = getEmbeddingModel().split("/").pop()?.toLowerCase() ?? "";
+  return CLS_POOLED_MODELS.has(basename2) ? "cls" : "mean";
+}
+function getEmbeddingDevice() {
+  return process.env.AMEM_EMBED_DEVICE?.trim() || void 0;
+}
+function getEmbeddingDtype() {
+  return process.env.AMEM_EMBED_DTYPE?.trim() || void 0;
+}
+function extractorKey() {
+  return `${getEmbeddingModel()}|${getEmbeddingDevice() ?? ""}|${getEmbeddingDtype() ?? ""}`;
+}
 async function getExtractor() {
-  const wanted = getEmbeddingModel();
-  if (extractor && loadedModelName === wanted) return extractor;
+  const wanted = extractorKey();
+  if (extractor && loadedKey === wanted) return extractor;
   if (!pipeline) {
     const mod = await import("@huggingface/transformers");
     pipeline = mod.pipeline;
   }
-  extractor = await pipeline("feature-extraction", wanted, {
-    revision: "main"
+  const device = getEmbeddingDevice();
+  const dtype = getEmbeddingDtype();
+  extractor = await pipeline("feature-extraction", getEmbeddingModel(), {
+    revision: "main",
+    // Omitted entirely when unset, so an unconfigured install gets exactly the
+    // library defaults it got before these existed.
+    ...device ? { device } : {},
+    ...dtype ? { dtype } : {}
   });
-  loadedModelName = wanted;
+  loadedKey = wanted;
   cachedDim = null;
   return extractor;
 }
 async function getEmbeddingDim() {
-  if (cachedDim !== null && loadedModelName === getEmbeddingModel()) return cachedDim;
+  if (cachedDim !== null && loadedKey === extractorKey()) return cachedDim;
   const probe = await encode("dimension probe");
   cachedDim = probe.length;
   return cachedDim;
 }
-function meanPoolingNormalize(output, attentionMask) {
+function poolNormalize(output, attentionMask, mode) {
   const seqLen = output.length;
   const dim = output[0].length;
   const pooled = new Array(dim).fill(0);
-  let maskSum = 0;
-  for (let i = 0; i < seqLen; i++) {
-    const m = attentionMask[i];
-    maskSum += m;
+  if (mode === "cls") {
+    for (let j = 0; j < dim; j++) pooled[j] = output[0][j];
+  } else {
+    let maskSum = 0;
+    for (let i = 0; i < seqLen; i++) {
+      const m = attentionMask[i];
+      maskSum += m;
+      for (let j = 0; j < dim; j++) {
+        pooled[j] += output[i][j] * m;
+      }
+    }
     for (let j = 0; j < dim; j++) {
-      pooled[j] += output[i][j] * m;
+      pooled[j] /= Math.max(maskSum, 1e-9);
     }
   }
-  for (let j = 0; j < dim; j++) {
-    pooled[j] /= Math.max(maskSum, 1e-9);
-  }
   let norm = 0;
   for (const v of pooled) norm += v * v;
   norm = Math.sqrt(norm);
@@ -130,7 +172,8 @@ function meanPoolingNormalize(output, attentionMask) {
 }
 async function encode(text) {
   const ext = await getExtractor();
-  const result = await ext(text, { pooling: "mean", normalize: true });
+  const pooling = getEmbeddingPooling();
+  const result = await ext(text, { pooling, normalize: true });
   if (result && result.data) {
     return Array.from(result.data);
   }
@@ -146,7 +189,7 @@ async function encode(text) {
       }
       raw.push(row);
     }
-    return meanPoolingNormalize(raw, new Array(seqLen).fill(1));
+    return poolNormalize(raw, new Array(seqLen).fill(1), pooling);
   }
   throw new Error("Unexpected embedding output shape");
 }
@@ -183,7 +226,7 @@ var EmbeddingDimensionMismatchError = class extends Error {
   constructor(collection, collectionDim, modelDim, model) {
     super(
       `Collection "${collection}" stores ${collectionDim}-dimension vectors, but the embedding model "${model}" produces ${modelDim}. Qdrant fixes a collection's vector size at creation and cannot change it, so writes and searches would both fail.
-Either set AMEM_EMBED_MODEL back to the model this collection was built with, or migrate: build a new collection with the new model, backfill it, then point AMEM_COLLECTION at it. See docs/reference/embedding-models.md.`
+Either set AMEM_EMBED_MODEL back to the model this collection was built with, or ${migrationHint(collection, model)}`
     );
     this.collection = collection;
     this.collectionDim = collectionDim;
@@ -196,6 +239,35 @@ Either set AMEM_EMBED_MODEL back to the model this collection was built with, or
   modelDim;
   model;
 };
+function migrationHint(collection, targetModel) {
+  return `migrate to a new collection:
+
+  AMEM_EMBED_MODEL=${targetModel} \\
+    npx --package=@amemhq/core amem-migrate --to ${collection}_v2
+
+That is a dry run; add --apply to write. "${collection}" is only read, so nothing is lost either way. When it looks right, point whatever names this collection at the new one \u2014 AMEM_COLLECTION, or the plugin's "collection" setting if this agent has its own. See https://amem.owo.lc/reference/embedding-models.`;
+}
+var EmbeddingModelMismatchError = class extends Error {
+  constructor(collection, collectionModel, configuredModel) {
+    super(
+      `Collection "${collection}" was built with the embedding model "${collectionModel}", but this process is configured for "${configuredModel}". Both produce vectors of the same width, so nothing would fail \u2014 searches would just quietly compare vectors from two different models.
+Either set AMEM_EMBED_MODEL back to "${collectionModel}", or ` + migrationHint(collection, configuredModel)
+    );
+    this.collection = collection;
+    this.collectionModel = collectionModel;
+    this.configuredModel = configuredModel;
+    this.name = "EmbeddingModelMismatchError";
+  }
+  collection;
+  collectionModel;
+  configuredModel;
+};
+async function recordCollectionModel(collection, model) {
+  try {
+    await qdrant("PATCH", `/collections/${collection}`, { metadata: { embedding_model: model } });
+  } catch {
+  }
+}
 async function qdrant(method, path5, body) {
   const res = await fetch(`${QDRANT_URL}${path5}`, {
     method,
@@ -238,6 +310,14 @@ async function ensureCollection(collectionName) {
         throw new EmbeddingDimensionMismatchError(col, collectionDim, modelDim, getEmbeddingModel());
       }
     }
+    const recorded = existing.config?.metadata?.embedding_model;
+    const current = getEmbeddingModel();
+    if (typeof recorded === "string" && recorded !== current) {
+      throw new EmbeddingModelMismatchError(col, recorded, current);
+    }
+    if (recorded === void 0) {
+      await recordCollectionModel(col, current);
+    }
     markReady();
     return;
   }
@@ -249,6 +329,7 @@ async function ensureCollection(collectionName) {
   } catch (err) {
     if (!(err instanceof Error) || !err.message.includes("already exists")) throw err;
   }
+  await recordCollectionModel(col, getEmbeddingModel());
   await qdrant("PUT", `/collections/${col}/index`, {
     field_name: "agent_id",
     field_schema: "keyword"
@@ -294,6 +375,7 @@ async function collectionDimRaw(collection) {
 }
 async function createCollectionRaw(collection, size) {
   await qdrant("PUT", `/collections/${collection}`, { vectors: { size, distance: "Cosine" } });
+  await recordCollectionModel(collection, getEmbeddingModel());
   for (const field_name of ["agent_id", "hash", "topics", "subjects"]) {
     await qdrant("PUT", `/collections/${collection}/index`, { field_name, field_schema: "keyword" });
   }
@@ -2333,6 +2415,7 @@ function isPlausibleUpdateTarget(newEmbedding, targetEmbedding, minSimilarity) {
   DEFAULT_CRUD_UPDATE_MIN_SIM,
   DEFAULT_EMBEDDING_MODEL,
   EmbeddingDimensionMismatchError,
+  EmbeddingModelMismatchError,
   addEpisodic,
   addMemory,
   canRead,
@@ -2347,8 +2430,11 @@ function isPlausibleUpdateTarget(newEmbedding, targetEmbedding, minSimilarity) {
   encode,
   ensureCollection,
   generateReviewBatch,
+  getEmbeddingDevice,
   getEmbeddingDim,
+  getEmbeddingDtype,
   getEmbeddingModel,
+  getEmbeddingPooling,
   getNote,
   invalidateNote,
   isModelLoaded,