@cohortapp/agent-sdk 2.3.0

package/lib/rag/indexer.mjs

@@ -0,0 +1,425 @@
+/**
+ * Maestro — RAG indexer (sqlite-fts5)
+ *
+ * Walks the agent's corpus and writes chunked content into a
+ * sqlite-fts5 virtual table at `state/rag/index/agent.db`. Voice
+ * sessions and other Claude sessions then query it via
+ * `lib/rag/search.mjs` for at-call context.
+ *
+ * Why FTS5 (and not vector embeddings, yet):
+ *   - Free, instant, no network round-trip.
+ *   - Good enough for the "what did we decide about X" pattern that
+ *     dominates voice queries — most lookups reference specific
+ *     names / projects / decision IDs.
+ *   - Semantic embeddings can be layered on top later for queries
+ *     FTS5 can't answer ("anything about culture we wrote last quarter?").
+ *
+ * Chunking:
+ *   - Markdown: split on H2/H3 headings, sub-split anything > 1500 chars
+ *     with 200-char overlap.
+ *   - YAML queues: each item becomes a chunk.
+ *   - Default: fixed-window 1500-char chunks with 200-char overlap.
+ *
+ * Idempotency: each file's mtime is tracked in the `files` table.
+ * Unchanged files are skipped; deleted files trigger purge.
+ *
+ * Public API:
+ *   buildIndex({ agentRoot, force? }) → { added, updated, deleted, took_ms }
+ *   indexedAt(agentRoot, sourcePath) → epoch_ms | null
+ */
+
+import { readFileSync, statSync, existsSync, readdirSync, mkdirSync } from "node:fs";
+import { join, resolve, relative } from "node:path";
+
+let _Database = null;
+async function loadDatabase() {
+  if (_Database) return _Database;
+  try {
+    const mod = await import("better-sqlite3");
+    _Database = mod.default || mod;
+    return _Database;
+  } catch (err) {
+    throw new Error(
+      "RAG requires `better-sqlite3`. Install in the agent repo: " +
+        "`npm i better-sqlite3`. Underlying: " + err.message
+    );
+  }
+}
+
+const DEFAULT_SOURCES = [
+  { dir: "memory",             kind: "memory",     recurse: -1 },
+  { dir: "knowledge",          kind: "knowledge",  recurse: -1 },
+  { dir: "state/queues",       kind: "queue",      recurse: 0 },
+  { dir: "state/dashboards",   kind: "dashboard",  recurse: 0 },
+  { dir: "config",             kind: "config",     recurse: 0 },
+  { dir: "agents",             kind: "agent_def",  recurse: -1 },
+  { dir: "workflows",          kind: "workflow",   recurse: -1 },
+  { dir: "outputs/briefs",     kind: "brief",      recurse: -1 },
+  { dir: "outputs/memos",      kind: "memo",       recurse: -1 },
+  { dir: "policies",           kind: "policy",     recurse: -1 },
+];
+
+const EXTENSIONS = /\.(md|yaml|yml|json|jsonl|txt)$/i;
+const MAX_FILE_BYTES = 512 * 1024;
+const CHUNK_TARGET_CHARS = 1500;
+const CHUNK_OVERLAP_CHARS = 200;
+
+const SCHEMA = `
+CREATE VIRTUAL TABLE IF NOT EXISTS chunks USING fts5(
+  body,
+  source_path UNINDEXED,
+  kind UNINDEXED,
+  title UNINDEXED,
+  heading UNINDEXED,
+  tokenize = "porter unicode61 remove_diacritics 2"
+);
+
+CREATE TABLE IF NOT EXISTS files (
+  source_path TEXT PRIMARY KEY,
+  kind TEXT NOT NULL,
+  mtime_ms INTEGER NOT NULL,
+  size_bytes INTEGER NOT NULL,
+  chunk_count INTEGER NOT NULL,
+  indexed_at INTEGER NOT NULL
+);
+
+-- Semantic vectors for each chunk. chunk_rowid mirrors chunks.rowid.
+-- Vectors are L2-normalized at write time so cosine similarity reduces
+-- to a plain dot product at read time.
+CREATE TABLE IF NOT EXISTS embeddings (
+  chunk_rowid INTEGER PRIMARY KEY,
+  model TEXT NOT NULL,
+  dim INTEGER NOT NULL,
+  vector BLOB NOT NULL
+);
+
+CREATE INDEX IF NOT EXISTS files_kind ON files(kind);
+CREATE INDEX IF NOT EXISTS files_indexed_at ON files(indexed_at);
+CREATE INDEX IF NOT EXISTS embeddings_model ON embeddings(model);
+`;
+
+function openDir(agentRoot) {
+  const indexDir = join(agentRoot, "state/rag/index");
+  mkdirSync(indexDir, { recursive: true });
+  return indexDir;
+}
+
+function* walk(rootAgentDir, source) {
+  const startAbs = join(rootAgentDir, source.dir);
+  if (!existsSync(startAbs)) return;
+  const queue = [{ path: startAbs, depth: 0 }];
+  while (queue.length) {
+    const { path, depth } = queue.shift();
+    let entries;
+    try { entries = readdirSync(path, { withFileTypes: true }); }
+    catch { continue; }
+    for (const ent of entries) {
+      if (ent.name.startsWith(".")) continue;
+      const full = join(path, ent.name);
+      if (ent.isDirectory()) {
+        if (source.recurse === -1 || depth < source.recurse) {
+          queue.push({ path: full, depth: depth + 1 });
+        }
+        continue;
+      }
+      if (!ent.isFile()) continue;
+      if (!EXTENSIONS.test(ent.name)) continue;
+      try {
+        const st = statSync(full);
+        if (st.size > MAX_FILE_BYTES) continue;
+        yield {
+          absPath: full,
+          relPath: relative(rootAgentDir, full),
+          kind: source.kind,
+          mtime_ms: Math.floor(st.mtimeMs),
+          size: st.size,
+        };
+      } catch { /* */ }
+    }
+  }
+}
+
+function chunkText(text, sourcePath, kind) {
+  if (!text || !text.trim()) return [];
+
+  if (kind === "queue" && /\nitems:/i.test(text)) {
+    const out = [];
+    const parts = text.split(/\n\s*-\s+/);
+    const header = parts.shift() || "";
+    for (let i = 0; i < parts.length; i++) {
+      const body = "- " + parts[i].trim();
+      out.push({ heading: `item ${i + 1}`, body: header.split("items:")[0] + "\n" + body });
+    }
+    if (!out.length) out.push({ heading: "", body: text });
+    return out;
+  }
+
+  // JSONL (e.g. collective memory cards): one chunk per record so recall
+  // returns a single card, not a window of several. Falls back to the raw line.
+  if (sourcePath.endsWith(".jsonl")) {
+    const out = [];
+    for (const line of text.split("\n")) {
+      const t = line.trim();
+      if (!t) continue;
+      try {
+        const o = JSON.parse(t);
+        if (o && o.title) {
+          const tags = Array.isArray(o.tags) && o.tags.length ? ` [${o.tags.join(", ")}]` : "";
+          const body = o.body ? ` — ${o.body}` : "";
+          out.push({ heading: o.kind || "", body: `(${o.kind || "note"}) ${o.title}${body}${tags}` });
+          continue;
+        }
+      } catch { /* not a card line */ }
+      out.push({ heading: "", body: t });
+    }
+    return out.length ? out : [{ heading: "", body: text }];
+  }
+
+  if (sourcePath.endsWith(".md")) {
+    const sections = text.split(/^(#{2,3}\s+.+)$/m);
+    const out = [];
+    let current = { heading: "(intro)", body: "" };
+    for (const seg of sections) {
+      if (/^#{2,3}\s+/.test(seg)) {
+        if (current.body.trim()) out.push(current);
+        current = { heading: seg.replace(/^#+\s+/, "").trim(), body: seg + "\n" };
+      } else {
+        current.body += seg;
+      }
+    }
+    if (current.body.trim()) out.push(current);
+    const final = [];
+    for (const s of out) {
+      if (s.body.length <= CHUNK_TARGET_CHARS + 500) { final.push(s); continue; }
+      let i = 0;
+      while (i < s.body.length) {
+        const end = Math.min(s.body.length, i + CHUNK_TARGET_CHARS);
+        final.push({ heading: s.heading, body: s.body.slice(i, end) });
+        if (end >= s.body.length) break;
+        const next = end - CHUNK_OVERLAP_CHARS;
+        // Guarantee forward progress: minimum step = 1 char so we
+        // can't loop forever on pathologically small sections.
+        i = next > i ? next : i + 1;
+      }
+    }
+    return final;
+  }
+
+  const out = [];
+  let i = 0;
+  while (i < text.length) {
+    const end = Math.min(text.length, i + CHUNK_TARGET_CHARS);
+    out.push({ heading: "", body: text.slice(i, end) });
+    if (end >= text.length) break;
+    const next = end - CHUNK_OVERLAP_CHARS;
+    i = next > i ? next : i + 1;
+  }
+  return out;
+}
+
+function deriveTitle(relPath, content) {
+  if (relPath.endsWith(".md")) {
+    const m = content.match(/^#\s+(.+)$/m);
+    if (m) return m[1].trim().slice(0, 200);
+  }
+  for (const re of [/^title:\s*(.+)$/m, /^name:\s*(.+)$/m, /^decision_id:\s*(.+)$/m]) {
+    const m = content.match(re);
+    if (m) return m[1].trim().replace(/^["']|["']$/g, "").slice(0, 200);
+  }
+  return relPath.split("/").pop().replace(/\.(md|yaml|yml|json|jsonl|txt)$/i, "");
+}
+
+/**
+ * Build (or incrementally update) the index.
+ *
+ * @param {object} opts
+ * @param {string} opts.agentRoot
+ * @param {boolean} [opts.force]
+ * @param {Array}   [opts.sources]
+ * @param {function} [opts.log]
+ * @param {boolean} [opts.embeddings]    Generate semantic vectors (default:
+ *                                       true if OPENAI_API_KEY is set).
+ * @param {string}  [opts.embedModel]    "text-embedding-3-small" (default)
+ *                                       or "text-embedding-3-large".
+ * @returns {Promise<{added,updated,deleted,unchanged,chunks,embeds,took_ms,cost_usd}>}
+ */
+export async function buildIndex(opts = {}) {
+  const Database = await loadDatabase();
+  const root = resolve(
+    opts.agentRoot || process.env.AGENT_ROOT || process.env.AGENT_DIR || process.cwd()
+  );
+  const indexDir = openDir(root);
+  const dbPath = join(indexDir, "agent.db");
+  const sources = opts.sources || DEFAULT_SOURCES;
+  const log = opts.log || (() => {});
+  const wantEmbeddings = opts.embeddings ?? !!process.env.OPENAI_API_KEY;
+  const embedModel = opts.embedModel || "text-embedding-3-small";
+  const t0 = Date.now();
+
+  const db = new Database(dbPath);
+  db.pragma("journal_mode = WAL");
+  db.pragma("synchronous = NORMAL");
+  db.exec(SCHEMA);
+
+  const existingRows = db.prepare("SELECT source_path, mtime_ms FROM files").all();
+  const existing = new Map(existingRows.map((r) => [r.source_path, r.mtime_ms]));
+
+  const result = { added: 0, updated: 0, deleted: 0, unchanged: 0, chunks: 0, embeds: 0, purged: 0, cost_usd: 0 };
+
+  const insertChunk = db.prepare(
+    "INSERT INTO chunks (body, source_path, kind, title, heading) VALUES (?, ?, ?, ?, ?)"
+  );
+  const deleteChunks = db.prepare("DELETE FROM chunks WHERE source_path = ?");
+  const upsertFile = db.prepare(
+    "INSERT INTO files (source_path, kind, mtime_ms, size_bytes, chunk_count, indexed_at) " +
+      "VALUES (?, ?, ?, ?, ?, ?) " +
+      "ON CONFLICT(source_path) DO UPDATE SET " +
+      "kind=excluded.kind, mtime_ms=excluded.mtime_ms, size_bytes=excluded.size_bytes, " +
+      "chunk_count=excluded.chunk_count, indexed_at=excluded.indexed_at"
+  );
+  const deleteEmbeddingsForFile = db.prepare(
+    "DELETE FROM embeddings WHERE chunk_rowid IN (SELECT rowid FROM chunks WHERE source_path = ?)"
+  );
+
+  // Track newly-inserted chunks so we can backfill embeddings outside
+  // the synchronous transaction (network calls don't fit inside one).
+  /** @type {Array<{ rowid: number, body: string }>} */
+  const newChunks = [];
+  const seen = new Set();
+
+  const txn = db.transaction((files) => {
+    for (const f of files) {
+      seen.add(f.relPath);
+      const wasIndexed = existing.has(f.relPath);
+      const unchanged = wasIndexed && existing.get(f.relPath) === f.mtime_ms && !opts.force;
+      if (unchanged) { result.unchanged++; continue; }
+
+      let body = "";
+      try { body = readFileSync(f.absPath, "utf-8"); }
+      catch { continue; }
+      const title = deriveTitle(f.relPath, body);
+      const chunks = chunkText(body, f.relPath, f.kind);
+
+      // Remove old chunks AND their embeddings before re-inserting.
+      deleteEmbeddingsForFile.run(f.relPath);
+      deleteChunks.run(f.relPath);
+      for (const c of chunks) {
+        const info = insertChunk.run(c.body, f.relPath, f.kind, title, c.heading || "");
+        newChunks.push({ rowid: info.lastInsertRowid, body: c.body });
+        result.chunks++;
+      }
+      upsertFile.run(f.relPath, f.kind, f.mtime_ms, f.size, chunks.length, Date.now());
+
+      if (wasIndexed) result.updated++;
+      else result.added++;
+    }
+  });
+
+  const allFiles = [];
+  for (const source of sources) {
+    for (const f of walk(root, source)) allFiles.push(f);
+  }
+  log("info", `walking ${allFiles.length} files across ${sources.length} sources…`);
+  txn(allFiles);
+
+  for (const sp of existing.keys()) {
+    if (!seen.has(sp)) {
+      deleteEmbeddingsForFile.run(sp);
+      deleteChunks.run(sp);
+      db.prepare("DELETE FROM files WHERE source_path = ?").run(sp);
+      result.deleted++;
+    }
+  }
+
+  // ── Embeddings backfill — happens OUTSIDE the txn since it does network I/O ──
+  //
+  // Two correctness fixes vs. a naive "embed only newChunks" backfill:
+  //   M13: the file mtime is committed in the txn above BEFORE this network
+  //        step. A crash here used to leave chunks marked indexed-but-
+  //        unembedded; the next non-force run skipped the file (mtime
+  //        unchanged) and the chunks stayed unembedded forever. We now also
+  //        sweep ANY chunk lacking a current-(model,dim) embedding and embed
+  //        it, so an interrupted embed is completed on the next plain run.
+  //   M12: a re-index under a new embedModel leaves old-dim rows behind whose
+  //        cosine silently scores 0. We resolve the model/dim from a probe
+  //        embed and PURGE embeddings that don't match it, instead of relying
+  //        on that 0-similarity coincidence.
+  if (wantEmbeddings) {
+    try {
+      const embedMod = await import("./embed.mjs");
+      const _embedTexts = typeof opts._embedTexts === "function" ? opts._embedTexts : embedMod.embedTexts;
+
+      // Resolve the target (model, dim) via a single probe so we can both
+      // purge mismatches (M12) and identify already-correct embeddings (M13).
+      const probe = await _embedTexts([""], { model: embedModel, normalize: true });
+      const model = probe.model || embedModel;
+      const dim = probe.vectors[0]?.length || 0;
+
+      // M12: purge any embedding that isn't the current (model, dim).
+      const purged = db.prepare(
+        "DELETE FROM embeddings WHERE model != ? OR dim != ?"
+      ).run(model, dim);
+      result.purged = purged.changes || 0;
+      if (result.purged) {
+        log("info", `purged ${result.purged} stale-dim embeddings (now ${model}×${dim})`);
+      }
+
+      // M13: collect every chunk still missing a current-(model,dim)
+      // embedding — this is newChunks PLUS anything stranded by a prior
+      // interrupted run. No --force required.
+      const missing = db.prepare(
+        "SELECT c.rowid AS rowid, c.body AS body FROM chunks c " +
+        "LEFT JOIN embeddings e ON e.chunk_rowid = c.rowid AND e.model = ? AND e.dim = ? " +
+        "WHERE e.chunk_rowid IS NULL"
+      ).all(model, dim);
+
+      if (missing.length) {
+        const totalChars = missing.reduce((n, c) => n + (c.body?.length || 0), 0);
+        result.cost_usd = embedMod.embedCostEstimate(totalChars, embedModel);
+        log("info", `generating embeddings for ${missing.length} chunks via ${embedModel} (≈$${result.cost_usd})…`);
+
+        const { vectors, model: usedModel } = await _embedTexts(
+          missing.map((c) => c.body),
+          { model: embedModel, normalize: true }
+        );
+        const usedDim = vectors[0]?.length || dim;
+        const insertEmbedding = db.prepare(
+          "INSERT OR REPLACE INTO embeddings (chunk_rowid, model, dim, vector) VALUES (?, ?, ?, ?)"
+        );
+        const writeAll = db.transaction((rows) => {
+          for (const r of rows) insertEmbedding.run(r.rowid, usedModel || model, usedDim, embedMod.serializeVector(r.vector));
+        });
+        writeAll(missing.map((c, i) => ({ rowid: c.rowid, vector: vectors[i] })));
+        result.embeds = missing.length;
+        log("info", `embeddings written: ${result.embeds} chunks × ${usedDim} dims`);
+      }
+    } catch (err) {
+      log("warn", `embeddings step skipped — ${err.message}`);
+      // BM25-only fallback is fine; search.mjs handles missing embeddings.
+    }
+  }
+
+  if ((result.added + result.updated + result.deleted) > 50) {
+    try { db.exec("INSERT INTO chunks(chunks) VALUES('optimize')"); } catch { /* */ }
+  }
+
+  db.close();
+  result.took_ms = Date.now() - t0;
+  return result;
+}
+
+export async function indexedAt(agentRoot, sourcePath) {
+  const Database = await loadDatabase();
+  const dbPath = join(
+    resolve(agentRoot || process.env.AGENT_ROOT || process.cwd()),
+    "state/rag/index/agent.db"
+  );
+  if (!existsSync(dbPath)) return null;
+  const db = new Database(dbPath, { readonly: true });
+  try {
+    const row = db.prepare("SELECT indexed_at FROM files WHERE source_path = ?").get(sourcePath);
+    return row ? row.indexed_at : null;
+  } finally { db.close(); }
+}
+
+export default { buildIndex, indexedAt };