agent-memory-store 0.0.4 → 0.0.6

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;
+}

package/src/store.js

@@ -1,119 +1,93 @@
 /**
- * context-store: file-based persistent memory for multi-agent systems.
+ * context-store: SQLite-backed persistent memory for multi-agent systems.
  *
- * Storage layout:
- *   <CONTEXT_STORE_PATH>/
- *     chunks/   → one .md file per chunk, YAML frontmatter + markdown body
- *     state/    → one .json file per key (session state / pipeline variables)
- *
- * Chunk file format:
- *   ---
- *   id: <sha1-10>
- *   topic: "Descriptive title of the chunk"
- *   agent: agent-id
- *   tags: [tag1, tag2]
- *   importance: low | medium | high | critical
- *   updated: ISO-8601
- *   expires: ISO-8601   # optional — omit for permanent chunks
- *   ---
- *   Markdown content here.
+ * Storage: single SQLite database at <STORE_PATH>/store.db
+ * Search: hybrid (BM25 via FTS5 + vector cosine similarity + RRF)
+ * Embeddings: local via @huggingface/transformers (graceful fallback to BM25-only)
  */
 
-import fs from "fs/promises";
-import path from "path";
-import matter from "gray-matter";
 import { createHash } from "crypto";
-import { BM25 } from "./bm25.js";
-
-const STORE_PATH = process.env.AGENT_STORE_PATH
-  ? path.resolve(process.env.AGENT_STORE_PATH)
-  : path.join(process.cwd(), ".agent-memory-store");
-
-const CHUNKS_DIR = path.join(STORE_PATH, "chunks");
-const STATE_DIR = path.join(STORE_PATH, "state");
-
-/** Ensures storage directories exist. */
-async function ensureDirs() {
-  await fs.mkdir(CHUNKS_DIR, { recursive: true });
-  await fs.mkdir(STATE_DIR, { recursive: true });
-}
+import {
+  getDb,
+  insertChunk,
+  getChunk,
+  deleteChunkById,
+  listChunksDb,
+  getStateDb,
+  setStateDb,
+  updateEmbedding,
+  getChunksWithoutEmbedding,
+} from "./db.js";
+import { hybridSearch } from "./search.js";
+import { embed, prepareText, warmup } from "./embeddings.js";
+import { migrateIfNeeded } from "./migrate.js";
 
 /**
- * Generates a stable short ID from agent + topic + current timestamp.
- * @param {string} agentId
- * @param {string} topic
- * @returns {string} 10-char hex string
+ * Initializes the store: runs migration if needed, warms up DB and embeddings.
+ * Called once at startup.
  */
-function generateId(agentId, topic) {
-  const seed = `${agentId}:${topic}:${Date.now()}:${Math.random()}`;
-  return createHash("sha1").update(seed).digest("hex").slice(0, 10);
+export async function initStore() {
+  // Ensure DB is ready (also runs schema + expiry purge)
+  getDb();
+
+  // Migrate from filesystem if needed
+  await migrateIfNeeded();
+
+  // Warm up embedding model in background (non-blocking)
+  warmup().then(() => backfillEmbeddings());
 }
 
 /**
- * Reads all non-expired chunks from disk.
- * Expired chunks are automatically deleted on read.
- *
- * @returns {Promise<Array<{ file: string, meta: object, content: string }>>}
+ * Background task: computes embeddings for chunks that don't have one yet.
  */
-async function loadAllChunks() {
-  await ensureDirs();
-  const files = await fs.readdir(CHUNKS_DIR).catch(() => []);
-  const chunks = [];
+async function backfillEmbeddings() {
+  const chunks = getChunksWithoutEmbedding();
+  if (!chunks.length) return;
 
-  await Promise.all(
-    files.map(async (file) => {
-      if (!file.endsWith(".md")) return;
-      try {
-        const raw = await fs.readFile(path.join(CHUNKS_DIR, file), "utf8");
-        const { data: meta, content } = matter(raw);
+  process.stderr.write(
+    `[agent-memory-store] Backfilling embeddings for ${chunks.length} chunks...\n`,
+  );
 
-        if (meta.expires && new Date(meta.expires) < new Date()) {
-          await fs.unlink(path.join(CHUNKS_DIR, file)).catch(() => {});
-          return;
-        }
+  for (const chunk of chunks) {
+    const text = prepareText({
+      topic: chunk.topic,
+      tags: chunk.tags,
+      content: chunk.content,
+    });
+    const embedding = await embed(text);
+    if (embedding) {
+      updateEmbedding(chunk.id, embedding);
+    }
+  }
 
-        chunks.push({ file, meta, content: content.trim() });
-      } catch {
-        // Skip unreadable files silently
-      }
-    }),
+  process.stderr.write(
+    `[agent-memory-store] Embedding backfill complete.\n`,
   );
-
-  return chunks;
 }
 
 /**
- * Builds a BM25 index from a list of chunks.
- * Searchable text = topic + tags + agent + body content.
- *
- * @param {Array} chunks
- * @returns {BM25}
+ * Generates a stable short ID from agent + topic + current timestamp.
+ * @param {string} agentId
+ * @param {string} topic
+ * @returns {string} 10-char hex string
  */
-function buildIndex(chunks) {
-  const engine = new BM25();
-  for (const c of chunks) {
-    const searchText = [
-      c.meta.topic || "",
-      (c.meta.tags || []).join(" "),
-      c.meta.agent || "",
-      c.content,
-    ].join(" ");
-    engine.addDocument(c.file, searchText, c.meta);
-  }
-  return engine;
+function generateId(agentId, topic) {
+  const seed = `${agentId}:${topic}:${Date.now()}:${Math.random()}`;
+  return createHash("sha1").update(seed).digest("hex").slice(0, 10);
 }
 
 // ─── Public API ───────────────────────────────────────────────────────────────
 
 /**
- * Searches chunks by relevance using BM25, with optional tag and agent filters.
+ * Searches chunks using hybrid search (BM25 + vector + RRF).
  *
  * @param {object} opts
- * @param {string}   opts.query     - Search query text
- * @param {string[]} [opts.tags]    - Filter: only chunks matching any of these tags
- * @param {string}   [opts.agent]   - Filter: only chunks written by this agent
- * @param {number}   [opts.topK]    - Max results (default: 6)
- * @param {number}   [opts.minScore] - Minimum BM25 score threshold (default: 0.1)
+ * @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 searchChunks({
@@ -122,53 +96,23 @@ export async function searchChunks({
   agent,
   topK = 6,
   minScore = 0.1,
+  mode = "hybrid",
 }) {
-  const chunks = await loadAllChunks();
-  if (chunks.length === 0) return [];
-
-  const engine = buildIndex(chunks);
-  const hasFilter = tags.length > 0 || !!agent;
-
-  const filter = hasFilter
-    ? (meta) => {
-        if (agent && meta.agent !== agent) return false;
-        if (tags.length > 0 && !tags.some((t) => (meta.tags || []).includes(t)))
-          return false;
-        return true;
-      }
-    : null;
-
-  const hits = engine.search(query, topK, filter);
-  const byFile = Object.fromEntries(chunks.map((c) => [c.file, c]));
-
-  return hits
-    .filter((h) => h.score >= minScore)
-    .map((h) => {
-      const c = byFile[h.id];
-      return {
-        id: c.meta.id,
-        topic: c.meta.topic,
-        agent: c.meta.agent,
-        tags: c.meta.tags || [],
-        importance: c.meta.importance || "medium",
-        score: Math.round(h.score * 100) / 100,
-        content: c.content,
-        updated: c.meta.updated,
-      };
-    });
+  return hybridSearch({ query, tags, agent, topK, minScore, mode });
 }
 
 /**
- * Writes a new chunk to disk.
+ * Writes a new chunk to the database.
+ * Embedding is computed asynchronously in the background.
  *
  * @param {object} opts
- * @param {string}   opts.topic      - Short descriptive title
- * @param {string}   opts.content    - Markdown body content
- * @param {string}   [opts.agent]    - Agent identifier
- * @param {string[]} [opts.tags]     - Search tags
- * @param {string}   [opts.importance] - low | medium | high | critical
- * @param {number}   [opts.ttlDays]  - Days until auto-expiry (omit = permanent)
- * @returns {Promise<{ id, file, topic, tags, importance }>}
+ * @param {string}   opts.topic
+ * @param {string}   opts.content
+ * @param {string}   [opts.agent]
+ * @param {string[]} [opts.tags]
+ * @param {string}   [opts.importance]
+ * @param {number}   [opts.ttlDays]
+ * @returns {Promise<{ id, topic, tags, importance }>}
  */
 export async function writeChunk({
   topic,
@@ -178,29 +122,32 @@ export async function writeChunk({
   importance = "medium",
   ttlDays = null,
 }) {
-  await ensureDirs();
-
   const id = generateId(agent, topic);
   const now = new Date().toISOString();
-  const expires = ttlDays
+  const expiresAt = ttlDays
     ? new Date(Date.now() + ttlDays * 86_400_000).toISOString()
     : null;
 
-  const meta = {
+  insertChunk({
     id,
     topic,
     agent,
     tags,
     importance,
-    updated: now,
-    ...(expires ? { expires } : {}),
-  };
-
-  const fileContent = matter.stringify(`\n${content}\n`, meta);
-  const filename = `${id}.md`;
-
-  await fs.writeFile(path.join(CHUNKS_DIR, filename), fileContent, "utf8");
-  return { id, file: filename, topic, tags, importance };
+    content,
+    embedding: null, // Computed in background
+    createdAt: now,
+    updatedAt: now,
+    expiresAt,
+  });
+
+  // Compute embedding in background (non-blocking)
+  const text = prepareText({ topic, tags, content });
+  embed(text).then((embedding) => {
+    if (embedding) updateEmbedding(id, embedding);
+  });
+
+  return { id, topic, tags, importance };
 }
 
 /**
@@ -210,27 +157,36 @@ export async function writeChunk({
  * @returns {Promise<{ meta: object, content: string } | null>}
  */
 export async function readChunk(id) {
-  const chunks = await loadAllChunks();
-  return chunks.find((c) => c.meta.id === id) ?? null;
+  const chunk = getChunk(id);
+  if (!chunk) return null;
+
+  return {
+    meta: {
+      id: chunk.id,
+      topic: chunk.topic,
+      agent: chunk.agent,
+      tags: chunk.tags,
+      importance: chunk.importance,
+      updated: chunk.updatedAt,
+      ...(chunk.expiresAt ? { expires: chunk.expiresAt } : {}),
+    },
+    content: chunk.content,
+  };
 }
 
 /**
  * Deletes a chunk by ID.
  *
  * @param {string} id
- * @returns {Promise<boolean>} true if deleted, false if not found
+ * @returns {Promise<boolean>}
  */
 export async function deleteChunk(id) {
-  const chunks = await loadAllChunks();
-  const target = chunks.find((c) => c.meta.id === id);
-  if (!target) return false;
-  await fs.unlink(path.join(CHUNKS_DIR, target.file));
-  return true;
+  return deleteChunkById(id);
 }
 
 /**
  * Lists chunk metadata without loading full content.
- * Results are sorted by most recently updated.
+ * Sorted by most recently updated.
  *
  * @param {object}   opts
  * @param {string}   [opts.agent]
@@ -238,23 +194,7 @@ export async function deleteChunk(id) {
  * @returns {Promise<Array>}
  */
 export async function listChunks({ agent, tags = [] } = {}) {
-  const chunks = await loadAllChunks();
-  return chunks
-    .filter((c) => {
-      if (agent && c.meta.agent !== agent) return false;
-      if (tags.length > 0 && !tags.some((t) => (c.meta.tags || []).includes(t)))
-        return false;
-      return true;
-    })
-    .map((c) => ({
-      id: c.meta.id,
-      topic: c.meta.topic,
-      agent: c.meta.agent,
-      tags: c.meta.tags || [],
-      importance: c.meta.importance || "medium",
-      updated: c.meta.updated,
-    }))
-    .sort((a, b) => new Date(b.updated) - new Date(a.updated));
+  return listChunksDb({ agent, tags });
 }
 
 /**
@@ -264,29 +204,16 @@ export async function listChunks({ agent, tags = [] } = {}) {
  * @returns {Promise<any | null>}
  */
 export async function getState(key) {
-  await ensureDirs();
-  try {
-    const raw = await fs.readFile(path.join(STATE_DIR, `${key}.json`), "utf8");
-    return JSON.parse(raw).value;
-  } catch {
-    return null;
-  }
+  return getStateDb(key);
 }
 
 /**
- * Writes a session state variable (any JSON-serializable value).
+ * Writes a session state variable.
  *
  * @param {string} key
  * @param {any}    value
  * @returns {Promise<{ key: string, updated: string }>}
  */
 export async function setState(key, value) {
-  await ensureDirs();
-  const updated = new Date().toISOString();
-  await fs.writeFile(
-    path.join(STATE_DIR, `${key}.json`),
-    JSON.stringify({ key, value, updated }, null, 2),
-    "utf8",
-  );
-  return { key, updated };
+  return setStateDb(key, value);
 }