agent-memory-store 0.0.5 → 0.0.7

package/src/embeddings.js

@@ -0,0 +1,124 @@
+/**
+ * Local embedding generation via @huggingface/transformers.
+ *
+ * Uses the all-MiniLM-L6-v2 model (384 dimensions) running locally via ONNX Runtime.
+ * Model is auto-downloaded (~23MB) on first use and cached in ~/.cache/huggingface/.
+ *
+ * Graceful degradation: if the model fails to load, all functions return null
+ * and the system falls back to BM25-only search.
+ */
+
+let pipelineInstance = null;
+let loadFailed = false;
+let loadingPromise = null;
+
+/**
+ * Lazily initializes the feature-extraction pipeline.
+ * Returns null if the model cannot be loaded.
+ * Ensures only one load attempt runs at a time.
+ */
+async function getPipeline() {
+  if (pipelineInstance) return pipelineInstance;
+  if (loadFailed) return null;
+
+  // Deduplicate concurrent load attempts
+  if (loadingPromise) return loadingPromise;
+
+  loadingPromise = (async () => {
+    try {
+      process.stderr.write(
+        "[agent-memory-store] Loading embedding model (first run downloads ~23MB)...\n",
+      );
+      const { pipeline } = await import("@huggingface/transformers");
+      pipelineInstance = await pipeline(
+        "feature-extraction",
+        "Xenova/all-MiniLM-L6-v2",
+        { dtype: "fp32" },
+      );
+      process.stderr.write(
+        "[agent-memory-store] Embedding model loaded successfully.\n",
+      );
+      return pipelineInstance;
+    } catch (err) {
+      loadFailed = true;
+      process.stderr.write(
+        `[agent-memory-store] Embedding model failed to load: ${err.message}\n` +
+          `[agent-memory-store] Falling back to BM25-only search.\n`,
+      );
+      return null;
+    } finally {
+      loadingPromise = null;
+    }
+  })();
+
+  return loadingPromise;
+}
+
+/**
+ * Generates an embedding for a single text string.
+ *
+ * @param {string} text - Text to embed (topic + tags + content)
+ * @returns {Promise<Float32Array|null>} 384-dim embedding or null if unavailable
+ */
+export async function embed(text) {
+  const extractor = await getPipeline();
+  if (!extractor) return null;
+
+  try {
+    const output = await extractor(text, {
+      pooling: "mean",
+      normalize: true,
+    });
+    return new Float32Array(output.data);
+  } catch (err) {
+    process.stderr.write(
+      `[agent-memory-store] Embedding error: ${err.message}\n`,
+    );
+    return null;
+  }
+}
+
+/**
+ * Generates embeddings for multiple texts.
+ *
+ * @param {string[]} texts
+ * @returns {Promise<Array<Float32Array|null>>}
+ */
+export async function embedBatch(texts) {
+  const results = [];
+  for (const text of texts) {
+    results.push(await embed(text));
+  }
+  return results;
+}
+
+/**
+ * Prepares searchable text from chunk fields for embedding.
+ *
+ * @param {object} chunk
+ * @param {string} chunk.topic
+ * @param {string[]|string} chunk.tags
+ * @param {string} chunk.content
+ * @returns {string}
+ */
+export function prepareText({ topic, tags, content }) {
+  const tagStr = Array.isArray(tags) ? tags.join(" ") : tags || "";
+  // Truncate content to ~800 chars to stay within model token limit
+  const truncated = content.length > 800 ? content.slice(0, 800) : content;
+  return `${topic} ${tagStr} ${truncated}`.trim();
+}
+
+/**
+ * Returns whether the embedding model is available.
+ */
+export function isEmbeddingAvailable() {
+  return pipelineInstance !== null && !loadFailed;
+}
+
+/**
+ * Pre-warms the embedding model (call during startup).
+ * Non-blocking — failures are silently handled.
+ */
+export async function warmup() {
+  await getPipeline();
+}

package/src/index.js

@@ -1,9 +1,9 @@
 #!/usr/bin/env node
 /**
- * agent-store MCP server entry point.
+ * agent-memory-store MCP server entry point.
  *
  * Exposes 7 tools to any MCP-compatible client (Claude Code, opencode, etc.):
- *   search_context  — BM25 full-text search over stored chunks
+ *   search_context  — Hybrid search (BM25 + semantic) over stored chunks
  *   write_context   — persist a new memory chunk
  *   read_context    — retrieve a chunk by ID
  *   list_context    — list chunk metadata (no body)
@@ -12,8 +12,8 @@
  *   set_state       — write a session state variable
  *
  * Usage:
- *   npx @agentops/context-store
- *   CONTEXT_STORE_PATH=/your/project/.context npx @agentops/context-store
+ *   npx agent-memory-store
+ *   AGENT_STORE_PATH=/your/project/.agent-memory-store npx agent-memory-store
  */
 
 import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
@@ -27,6 +27,7 @@ import {
   listChunks,
   getState,
   setState,
+  initStore,
 } from "./store.js";
 
 const { version } = JSON.parse(
@@ -35,6 +36,9 @@ const { version } = JSON.parse(
   ),
 );
 
+// Initialize database, run migration, warm up embeddings
+await initStore();
+
 const server = new McpServer({
   name: "context-store",
   version,
@@ -45,9 +49,10 @@ const server = new McpServer({
 server.tool(
   "search_context",
   [
-    "Search stored memory chunks by relevance using BM25 full-text ranking.",
+    "Search stored memory chunks using hybrid ranking (BM25 + semantic similarity).",
     "Call this at the start of any task to retrieve relevant prior knowledge,",
     "decisions, and outputs before generating a response.",
+    "Supports three modes: 'hybrid' (default, best quality), 'bm25' (keyword-only), 'semantic' (meaning-only).",
   ].join(" "),
   {
     query: z
@@ -75,16 +80,23 @@ server.tool(
       .min(0)
       .optional()
       .describe(
-        "Minimum BM25 relevance score. Lower = more permissive (default: 0.1).",
+        "Minimum relevance score. Lower = more permissive (default: 0.1).",
+      ),
+    search_mode: z
+      .enum(["hybrid", "bm25", "semantic"])
+      .optional()
+      .describe(
+        "Search strategy: 'hybrid' (BM25 + semantic, default), 'bm25' (keyword only), 'semantic' (embedding similarity only).",
       ),
   },
-  async ({ query, tags, agent, top_k, min_score }) => {
+  async ({ query, tags, agent, top_k, min_score, search_mode }) => {
     const results = await searchChunks({
       query,
       tags: tags ?? [],
       agent,
       topK: top_k ?? 6,
       minScore: min_score ?? 0.1,
+      mode: search_mode ?? "hybrid",
     });
 
     if (results.length === 0) {
@@ -111,9 +123,10 @@ server.tool(
 server.tool(
   "write_context",
   [
-    "Persist a memory chunk to local storage.",
+    "Persist a memory chunk to the database.",
     "Call this after completing a subtask, making a key decision,",
     "or producing output that downstream agents will need.",
+    "Embeddings are computed automatically in the background for semantic search.",
   ].join(" "),
   {
     topic: z

package/src/migrate.js

@@ -0,0 +1,124 @@
+/**
+ * Migration: filesystem-based storage → SQLite database.
+ *
+ * Runs automatically on first startup if the legacy chunks/ directory exists
+ * but store.db does not. Migrates all chunks and state, then renames the
+ * legacy directories to *_backup/.
+ */
+
+import fs from "fs/promises";
+import path from "path";
+import matter from "gray-matter";
+import { insertChunk, setStateDb, STORE_PATH } from "./db.js";
+
+const CHUNKS_DIR = path.join(STORE_PATH, "chunks");
+const STATE_DIR = path.join(STORE_PATH, "state");
+const DB_PATH = path.join(STORE_PATH, "store.db");
+
+/**
+ * Checks if migration is needed and runs it.
+ * @returns {Promise<boolean>} true if migration was performed
+ */
+export async function migrateIfNeeded() {
+  // Check if legacy chunks dir exists
+  const chunksExist = await fs
+    .stat(CHUNKS_DIR)
+    .then((s) => s.isDirectory())
+    .catch(() => false);
+
+  if (!chunksExist) return false;
+
+  // Check if DB already exists (already migrated)
+  const dbExists = await fs
+    .stat(DB_PATH)
+    .then((s) => s.isFile())
+    .catch(() => false);
+
+  if (dbExists) return false;
+
+  process.stderr.write(
+    "[agent-memory-store] Migrating filesystem storage to SQLite...\n",
+  );
+
+  let chunkCount = 0;
+  let stateCount = 0;
+
+  // Migrate chunks
+  try {
+    const files = await fs.readdir(CHUNKS_DIR);
+    for (const file of files) {
+      if (!file.endsWith(".md")) continue;
+      try {
+        const raw = await fs.readFile(path.join(CHUNKS_DIR, file), "utf8");
+        const { data: meta, content } = matter(raw);
+
+        // Skip expired chunks
+        if (meta.expires && new Date(meta.expires) < new Date()) continue;
+
+        const now = new Date().toISOString();
+        await insertChunk({
+          id: meta.id || file.replace(".md", ""),
+          topic: meta.topic || "Untitled",
+          agent: meta.agent || "global",
+          tags: meta.tags || [],
+          importance: meta.importance || "medium",
+          content: content.trim(),
+          embedding: null, // Will be computed in background
+          createdAt: meta.updated || now,
+          updatedAt: meta.updated || now,
+          expiresAt: meta.expires || null,
+        });
+        chunkCount++;
+      } catch {
+        // Skip unreadable files
+      }
+    }
+  } catch {
+    // chunks dir not readable
+  }
+
+  // Migrate state
+  try {
+    const files = await fs.readdir(STATE_DIR);
+    for (const file of files) {
+      if (!file.endsWith(".json")) continue;
+      try {
+        const raw = await fs.readFile(path.join(STATE_DIR, file), "utf8");
+        const { key, value } = JSON.parse(raw);
+        if (key) {
+          await setStateDb(key, value);
+          stateCount++;
+        }
+      } catch {
+        // Skip unreadable files
+      }
+    }
+  } catch {
+    // state dir not readable
+  }
+
+  // Rename legacy directories to backups
+  try {
+    await fs.rename(CHUNKS_DIR, CHUNKS_DIR + "_backup");
+  } catch {
+    // Rename failed — not critical
+  }
+
+  try {
+    const stateExists = await fs
+      .stat(STATE_DIR)
+      .then((s) => s.isDirectory())
+      .catch(() => false);
+    if (stateExists) {
+      await fs.rename(STATE_DIR, STATE_DIR + "_backup");
+    }
+  } catch {
+    // Rename failed — not critical
+  }
+
+  process.stderr.write(
+    `[agent-memory-store] Migration complete: ${chunkCount} chunks, ${stateCount} state entries.\n`,
+  );
+
+  return true;
+}

package/src/search.js

@@ -0,0 +1,151 @@
+/**
+ * Hybrid search engine combining FTS5 BM25 (native SQLite) and vector similarity.
+ *
+ * Search modes:
+ *   - "hybrid"   — FTS5 BM25 + vector cosine similarity merged via Reciprocal Rank Fusion
+ *   - "bm25"     — FTS5 only (no embeddings needed)
+ *   - "semantic"  — Vector similarity only
+ *
+ * Falls back to BM25-only if embeddings are not available.
+ */
+
+import { searchFTS, getAllEmbeddings, getChunk } from "./db.js";
+import { embed, isEmbeddingAvailable } from "./embeddings.js";
+
+// ─── Vector Search ──────────────────────────────────────────────────────────
+
+/**
+ * Computes cosine similarity between two Float32Arrays.
+ * Assumes both vectors are already L2-normalized (dot product = cosine sim).
+ */
+function cosineSimilarity(a, b) {
+  let dot = 0;
+  for (let i = 0; i < a.length; i++) dot += a[i] * b[i];
+  return dot;
+}
+
+/**
+ * Brute-force vector search over all chunk embeddings.
+ */
+function vectorSearch(queryEmbedding, { agent, tags = [], topK = 18 }) {
+  const embeddings = getAllEmbeddings({ agent, tags });
+  if (!embeddings.length) return [];
+
+  return embeddings
+    .map(({ id, embedding }) => ({
+      id,
+      score: cosineSimilarity(queryEmbedding, embedding),
+    }))
+    .filter((r) => r.score > 0)
+    .sort((a, b) => b.score - a.score)
+    .slice(0, topK);
+}
+
+// ─── Fusion ─────────────────────────────────────────────────────────────────
+
+/**
+ * Reciprocal Rank Fusion — merges two ranked lists into one.
+ */
+function reciprocalRankFusion(bm25Hits, vecHits, wBM25 = 0.4, wVec = 0.6) {
+  const K = 60;
+  const scores = new Map();
+
+  bm25Hits.forEach(({ id }, rank) => {
+    scores.set(id, (scores.get(id) || 0) + wBM25 / (K + rank + 1));
+  });
+
+  vecHits.forEach(({ id }, rank) => {
+    scores.set(id, (scores.get(id) || 0) + wVec / (K + rank + 1));
+  });
+
+  return [...scores.entries()]
+    .map(([id, score]) => ({ id, score }))
+    .sort((a, b) => b.score - a.score);
+}
+
+// ─── Main Search ────────────────────────────────────────────────────────────
+
+/**
+ * Main search function — performs hybrid, BM25, or semantic search.
+ *
+ * @param {object} opts
+ * @param {string}   opts.query
+ * @param {string[]} [opts.tags]
+ * @param {string}   [opts.agent]
+ * @param {number}   [opts.topK]
+ * @param {number}   [opts.minScore]
+ * @param {string}   [opts.mode]  - "hybrid" | "bm25" | "semantic"
+ * @returns {Promise<Array>}
+ */
+export async function hybridSearch({
+  query,
+  tags = [],
+  agent,
+  topK = 6,
+  minScore = 0.1,
+  mode = "hybrid",
+}) {
+  const candidateK = topK * 3;
+  const embeddingsReady = isEmbeddingAvailable();
+
+  // Determine effective mode
+  let effectiveMode = mode;
+  if ((mode === "hybrid" || mode === "semantic") && !embeddingsReady) {
+    effectiveMode = "bm25";
+  }
+
+  let fusedResults;
+
+  if (effectiveMode === "bm25") {
+    fusedResults = searchFTS({ query, agent, tags, topK: candidateK });
+  } else if (effectiveMode === "semantic") {
+    const queryEmbedding = await embed(query);
+    if (!queryEmbedding) {
+      fusedResults = searchFTS({ query, agent, tags, topK: candidateK });
+    } else {
+      fusedResults = vectorSearch(queryEmbedding, {
+        agent,
+        tags,
+        topK: candidateK,
+      });
+    }
+  } else {
+    // Hybrid: run FTS5 (sync) and embed query (async) in parallel
+    const queryEmbeddingPromise = embed(query);
+    const bm25Hits = searchFTS({ query, agent, tags, topK: candidateK });
+    const queryEmbedding = await queryEmbeddingPromise;
+
+    if (!queryEmbedding) {
+      fusedResults = bm25Hits;
+    } else {
+      const vecHits = vectorSearch(queryEmbedding, {
+        agent,
+        tags,
+        topK: candidateK,
+      });
+      fusedResults = reciprocalRankFusion(bm25Hits, vecHits);
+    }
+  }
+
+  // Take topK and enrich with full chunk data
+  const topResults = fusedResults.slice(0, topK);
+  const enriched = [];
+
+  for (const { id, score } of topResults) {
+    const chunk = getChunk(id);
+    if (!chunk) continue;
+
+    enriched.push({
+      id: chunk.id,
+      topic: chunk.topic,
+      agent: chunk.agent,
+      tags: chunk.tags,
+      importance: chunk.importance,
+      score: Math.round(score * 100) / 100,
+      content: chunk.content,
+      updated: chunk.updatedAt,
+    });
+  }
+
+  return enriched;
+}