@cohortapp/agent-sdk 2.3.0

package/lib/rag/search.mjs

@@ -0,0 +1,475 @@
+/**
+ * Maestro — RAG search (sqlite-fts5)
+ *
+ * Query the indexed agent corpus and return top-N matching chunks with
+ * BM25 scores and snippet text. Used by:
+ *   - `rag_search` tool exposed to voice/Claude sessions
+ *   - voice context-loader for at-call pre-loading of relevant chunks
+ *   - daemon's classifier-side lookups when full-context is overkill
+ *
+ * Public API:
+ *   search({ query, topN?, kind?, agentRoot? }) →
+ *     [{ source_path, kind, title, heading, snippet, score, relevance }]
+ *
+ * BM25 score range (SQLite FTS5 `bm25()`): NEGATIVE, where a more-negative
+ * value means a better match (0 is the worst possible — no term overlap).
+ * `ORDER BY bm25(chunks)` ascending therefore puts the best hit first. We
+ * expose a normalized `relevance` field in [0,1] (1.0 = top hit in this
+ * result set, → 0 = weakest) so callers don't need to interpret the raw,
+ * sign-inverted BM25 value directly. See `ftsRelevance()` for the transform.
+ */
+
+import { existsSync } from "node:fs";
+import { join, resolve } from "node:path";
+
+let _Database = null;
+async function loadDatabase() {
+  if (_Database) return _Database;
+  const mod = await import("better-sqlite3");
+  _Database = mod.default || mod;
+  return _Database;
+}
+
+const FTS_DB_REL = "state/rag/index/agent.db";
+
+// Stop-word list: common English words that add noise to BM25 ranking
+// without carrying meaning. Anything < 3 chars is also dropped.
+const STOPWORDS = new Set([
+  "the", "and", "for", "are", "but", "not", "you", "all", "any", "can",
+  "had", "her", "was", "one", "our", "out", "day", "get", "has", "him",
+  "his", "how", "man", "new", "now", "old", "see", "two", "way", "who",
+  "boy", "did", "its", "let", "put", "say", "she", "too", "use",
+  "with", "have", "this", "that", "from", "will", "what", "your", "when",
+  "been", "were", "they", "them", "than", "then", "would", "could",
+  "should", "about", "which", "their", "there", "where", "while",
+  "into", "onto", "upon", "over", "under", "after", "before", "between",
+]);
+
+/**
+ * Tokenize + sanitize user input into FTS5-safe MATCH tokens. Strategy:
+ *   1. Lowercase + strip non-alphanumeric → bare tokens.
+ *   2. Drop stop-words and tokens < 3 chars (they're noise for BM25).
+ *   3. Take the top ~16 tokens by length (longer = more specific).
+ *   4. Each token is double-quoted (with any embedded quotes stripped) so we
+ *      never pass an FTS5 operator (NEAR, NOT, OR, AND, *, …) from user input.
+ *
+ * Returns the array of quoted tokens (e.g. `["\"alpha\"", "\"beta\""]`) so
+ * callers can choose the boolean joiner — see `buildMatch()` for AND-first /
+ * OR-fallback. Returns `[]` when nothing survives sanitization.
+ */
+function tokenizeQuery(q) {
+  const all = String(q || "")
+    .toLowerCase()
+    .replace(/[^\p{L}\p{N}\s]/gu, " ")
+    .split(/\s+/)
+    .filter((t) => t.length >= 3 && !STOPWORDS.has(t));
+  if (!all.length) return [];
+  // Dedup + prefer specific terms (longer first), cap at 16.
+  const ranked = Array.from(new Set(all)).sort((a, b) => b.length - a.length).slice(0, 16);
+  return ranked.map((t) => `"${t.replace(/"/g, "")}"`);
+}
+
+/**
+ * Build an FTS5 MATCH expression from sanitized tokens.
+ *   - mode "and": require every token (`"a" AND "b"`). High precision; used
+ *     first for short multi-term queries so a doc matching ALL terms ranks
+ *     above one matching just a common term (audit L14).
+ *   - mode "or": any token may match (`"a" OR "b"`). Used as the zero-hit
+ *     fallback so a strict AND never silently returns nothing.
+ * A single token is identical under both modes.
+ */
+function buildMatch(tokens, mode) {
+  if (!tokens.length) return null;
+  return tokens.join(mode === "or" ? " OR " : " AND ");
+}
+
+/**
+ * Back-compat shim: the original API returned an OR-joined MATCH string.
+ * Kept so any external caller importing the old behaviour still works;
+ * internal callers use `tokenizeQuery` + `buildMatch` directly.
+ */
+function escapeQuery(q) {
+  const tokens = tokenizeQuery(q);
+  return tokens.length ? buildMatch(tokens, "or") : null;
+}
+
+/**
+ * @typedef {object} RagHit
+ * @property {string} source_path
+ * @property {string} kind
+ * @property {string} title
+ * @property {string|null} heading
+ * @property {string} snippet
+ * @property {number} score          BM25 score (negative; lower = better) — kept for compat
+ * @property {number} relevance      Normalized 0..1 (1.0 = best)
+ * @property {string} [retrieval]    "fts" | "semantic" | "hybrid"
+ * @property {number} [bm25_rank]    Rank in FTS5 results (1-based)
+ * @property {number} [semantic_rank] Rank in semantic results (1-based)
+ * @property {number} [semantic_score] Raw cosine similarity (0..1)
+ */
+
+/**
+ * Search the indexed corpus.
+ *
+ * Modes (auto-detected unless `mode` is set):
+ *   - "fts"      Pure BM25 keyword search (default if no embeddings available
+ *                or `mode: "fts"` is set explicitly). Fast, free, exact terms.
+ *   - "semantic" Pure cosine similarity over embeddings. Best for paraphrase
+ *                queries ("anything about culture last quarter") where the
+ *                exact terms aren't in the corpus.
+ *   - "hybrid"   Run both, fuse via Reciprocal Rank Fusion. Default when
+ *                embeddings are available — best of both worlds.
+ *
+ * @param {object} opts
+ * @param {string} opts.query           User query (free-text)
+ * @param {number} [opts.topN]          Max results (default: 8)
+ * @param {string} [opts.kind]          Filter by chunk.kind
+ * @param {string} [opts.agentRoot]
+ * @param {boolean} [opts.raw]          For "fts" mode: trust caller's MATCH syntax
+ * @param {"fts"|"semantic"|"hybrid"} [opts.mode]
+ * @param {string} [opts.embedModel]    Override embedding model for query
+ * @returns {Promise<Array<RagHit>>}
+ */
+export async function search(opts = {}) {
+  if (!opts.query) return [];
+  const Database = await loadDatabase();
+  const root = resolve(opts.agentRoot || process.env.AGENT_ROOT || process.env.AGENT_DIR || process.cwd());
+  const dbPath = join(root, FTS_DB_REL);
+  if (!existsSync(dbPath)) return [];
+
+  const topN = Math.min(Math.max(parseInt(opts.topN, 10) || 8, 1), 50);
+  const db = new Database(dbPath, { readonly: true });
+  try {
+    // Detect whether embeddings are available + decide mode.
+    const hasEmbeds = db.prepare("SELECT COUNT(*) AS n FROM embeddings").get().n > 0;
+    let mode = opts.mode || (hasEmbeds && process.env.OPENAI_API_KEY ? "hybrid" : "fts");
+    if (mode !== "fts" && !hasEmbeds) mode = "fts";
+
+    if (mode === "fts") return runFts(db, opts, topN);
+    if (mode === "semantic") return await runSemantic(db, opts, topN);
+    return await runHybrid(db, opts, topN);
+  } finally { db.close(); }
+}
+
+// ---------------------------------------------------------------------------
+// FTS-only path
+// ---------------------------------------------------------------------------
+
+/**
+ * Map a row's raw FTS5 bm25() score to a faithful [0,1] relevance, given the
+ * best (most-negative) score in the same result set (audit L15).
+ *
+ * FTS5 bm25() is NEGATIVE — more negative = better — and 0 means "no overlap".
+ * The old `bestScore / r.score` ratio silently assumed a POSITIVE BM25, so it
+ * inverted the ordering for the (negative) real values and, when `r.score` was
+ * 0, produced Infinity before clamping. Here we anchor relevance at the best
+ * score: the top hit is exactly 1.0, weaker (less-negative, larger) scores
+ * decay toward 0, and a 0 score (no overlap) maps to 0. All paths are finite —
+ * there is no division that can hit 0 or produce Inf/NaN.
+ *
+ * Exported (additively) so the offline test suite can assert the zero-score
+ * and out-of-range guards directly without a contrived FTS row.
+ *
+ * @param {number} score      this row's bm25() value (≤ 0)
+ * @param {number} bestScore  most-negative bm25() in the set (≤ 0)
+ * @returns {number} relevance in [0,1]
+ */
+export function ftsRelevance(score, bestScore) {
+  const s = Number(score);
+  const best = Number(bestScore);
+  if (!Number.isFinite(s) || !Number.isFinite(best)) return 0;
+  // Best score 0 means even the top hit has no measurable overlap → no signal.
+  if (best === 0) return 0;
+  // A row's own score of 0 (no overlap) maps to 0.
+  if (s === 0) return 0;
+  // Both s and best are < 0 here, and best is the most-negative (best) score,
+  // so |s| ≤ |best| and ratio = s / best ∈ (0, 1]: the top hit (s === best) is
+  // exactly 1.0, and a weaker (less-negative, i.e. closer-to-0) score yields a
+  // smaller fraction. (Using best / s would invert this and make weak hits look
+  // perfect.) The clamp is a belt-and-braces guard against FP edge cases.
+  const ratio = s / best;
+  if (!Number.isFinite(ratio)) return 0;
+  return Math.max(0, Math.min(1, ratio));
+}
+
+function runFts(db, opts, topN) {
+  if (opts.raw) {
+    return queryFts(db, opts.query, opts, topN);
+  }
+  const tokens = tokenizeQuery(opts.query);
+  if (!tokens.length) return [];
+  // AND-first for precision on short multi-term queries; fall back to OR only
+  // when the strict conjunction returns nothing (audit L14).
+  const andMatch = buildMatch(tokens, "and");
+  let rows = queryFts(db, andMatch, opts, topN);
+  if (rows.length === 0 && tokens.length > 1) {
+    const orMatch = buildMatch(tokens, "or");
+    rows = queryFts(db, orMatch, opts, topN);
+  }
+  return rows;
+}
+
+/** Execute one FTS5 MATCH and shape the rows (incl. faithful relevance). */
+function queryFts(db, matchQuery, opts, topN) {
+  if (!matchQuery) return [];
+  let sql = `
+    SELECT
+      rowid,
+      body,
+      source_path,
+      kind,
+      title,
+      heading,
+      bm25(chunks) AS score,
+      snippet(chunks, 0, '«', '»', '…', 32) AS snippet
+    FROM chunks
+    WHERE chunks MATCH ?
+  `;
+  const params = [matchQuery];
+  if (opts.kind) { sql += " AND kind = ?"; params.push(opts.kind); }
+  sql += " ORDER BY bm25(chunks) LIMIT ?";
+  params.push(topN);
+
+  const rows = db.prepare(sql).all(...params);
+  if (rows.length === 0) return [];
+  const bestScore = rows[0].score; // most-negative (best) thanks to ORDER BY
+  return rows.map((r, i) => ({
+    source_path: r.source_path,
+    kind: r.kind,
+    title: r.title,
+    heading: r.heading || null,
+    snippet: r.snippet,
+    score: r.score,
+    relevance: ftsRelevance(r.score, bestScore),
+    retrieval: "fts",
+    bm25_rank: i + 1,
+  }));
+}
+
+// ---------------------------------------------------------------------------
+// Semantic path — cosine similarity over all stored embeddings
+// ---------------------------------------------------------------------------
+
+// Candidate ceiling for the brute-force cosine scan. Above this many
+// matching-(model,dim) embeddings we cap the scan and warn rather than
+// pulling the entire BLOB column into JS (audit H11). Override via
+// RAG_MAX_SEMANTIC_CANDIDATES. Normal-size corpora (< ceiling) are
+// scored in full and behave exactly as before.
+export const DEFAULT_MAX_SEMANTIC_CANDIDATES = 50000;
+
+function maxSemanticCandidates() {
+  const n = parseInt(process.env.RAG_MAX_SEMANTIC_CANDIDATES, 10);
+  return Number.isFinite(n) && n > 0 ? n : DEFAULT_MAX_SEMANTIC_CANDIDATES;
+}
+
+async function runSemantic(db, opts, topN) {
+  const embedMod = await import("./embed.mjs");
+  // Resolve the query embedding AND the model/dim it was produced with, so
+  // we can constrain scoring to compatible embeddings only (audit M12). A
+  // re-index under a new embedModel leaves old-dim rows whose cosine would
+  // be 0 on length mismatch — silently collapsing recall. Filtering on
+  // (model, dim) makes that impossible instead of relying on a 0 coincidence.
+  let queryVec, queryModel;
+  if (typeof opts._embedQuery === "function") {
+    // Offline/test injection seam: returns { vector, model } or a bare vector.
+    const r = await opts._embedQuery(opts.query, { model: opts.embedModel });
+    if (r && r.vector) { queryVec = r.vector; queryModel = r.model; }
+    else { queryVec = r; queryModel = opts.embedModel; }
+  } else {
+    const r = await embedMod.embedTexts([opts.query || ""], { model: opts.embedModel });
+    queryVec = r.vectors[0];
+    queryModel = r.model;
+  }
+  if (!queryVec) return [];
+  const queryDim = queryVec.length;
+  if (!queryModel) queryModel = opts.embedModel || "text-embedding-3-small";
+
+  // Bind the (model, dim) filter (+ optional kind) on every embeddings read.
+  const filterSql = " WHERE e.model = ? AND e.dim = ?" + (opts.kind ? " AND c.kind = ?" : "");
+  const filterParams = opts.kind ? [queryModel, queryDim, opts.kind] : [queryModel, queryDim];
+
+  // Enforce a documented candidate ceiling (audit H11): count first, and if
+  // the compatible-embedding count exceeds the ceiling, cap the scan with a
+  // LIMIT and warn rather than risk OOM on an unbounded BLOB pull.
+  const ceiling = maxSemanticCandidates();
+  const candidateCount = db.prepare(
+    "SELECT COUNT(*) AS n FROM chunks c JOIN embeddings e ON e.chunk_rowid = c.rowid" + filterSql
+  ).get(...filterParams).n;
+  if (candidateCount === 0) return [];
+  let limitClause = "";
+  const scanParams = [...filterParams];
+  if (candidateCount > ceiling) {
+    console.warn(
+      `[rag] semantic candidate set ${candidateCount} exceeds ceiling ${ceiling}; ` +
+      `capping scan (set RAG_MAX_SEMANTIC_CANDIDATES to raise, or add an ANN index).`
+    );
+    limitClause = " LIMIT ?";
+    scanParams.push(ceiling);
+  }
+
+  // Scoring SELECT deliberately omits c.body (audit H11): bodies are large
+  // and we only need them for the final topN, fetched by rowid below.
+  const rows = db.prepare(
+    "SELECT c.rowid, c.source_path, c.kind, c.title, c.heading, e.vector " +
+    "FROM chunks c JOIN embeddings e ON e.chunk_rowid = c.rowid" + filterSql + limitClause
+  ).all(...scanParams);
+  if (rows.length === 0) return [];
+
+  // Brute-force cosine over the (now body-free) candidate rows.
+  const scored = rows.map((r) => {
+    const v = embedMod.deserializeVector(r.vector);
+    return {
+      rowid: r.rowid,
+      source_path: r.source_path,
+      kind: r.kind,
+      title: r.title,
+      heading: r.heading || null,
+      similarity: embedMod.cosineSimilarity(queryVec, v),
+    };
+  });
+  scored.sort((a, b) => b.similarity - a.similarity);
+  const top = scored.slice(0, topN);
+
+  // Fetch bodies ONLY for the final topN (by rowid) to build snippets.
+  const bodyById = new Map();
+  if (top.length) {
+    const placeholders = top.map(() => "?").join(",");
+    for (const br of db.prepare(
+      `SELECT rowid, body FROM chunks WHERE rowid IN (${placeholders})`
+    ).all(...top.map((r) => r.rowid))) {
+      bodyById.set(br.rowid, br.body);
+    }
+  }
+
+  return top.map((r, i) => ({
+    source_path: r.source_path,
+    kind: r.kind,
+    title: r.title,
+    heading: r.heading,
+    snippet: bareSnippet(bodyById.get(r.rowid) || "", opts.query),
+    score: r.similarity, // for compat
+    relevance: Math.max(0.01, Math.min(1, (r.similarity - top[top.length - 1].similarity) /
+                                        (top[0].similarity - top[top.length - 1].similarity + 1e-9))),
+    retrieval: "semantic",
+    semantic_rank: i + 1,
+    semantic_score: r.similarity,
+  }));
+}
+
+// ---------------------------------------------------------------------------
+// Hybrid path — Reciprocal Rank Fusion of BM25 + cosine rankings
+// ---------------------------------------------------------------------------
+
+const RRF_K = 60; // standard constant in the RRF literature (Cormack et al. 2009)
+const FETCH_PER_LIST = 30; // pull more from each list than topN so the fusion has material to rank
+
+async function runHybrid(db, opts, topN) {
+  // Run FTS5 first (cheap, in-process).
+  const ftsResults = runFts(db, { ...opts, topN: FETCH_PER_LIST }, FETCH_PER_LIST);
+  // And semantic in parallel.
+  let semanticResults = [];
+  try {
+    semanticResults = await runSemantic(db, { ...opts, topN: FETCH_PER_LIST }, FETCH_PER_LIST);
+  } catch { /* embeddings unavailable; we fall back to FTS-only below */ }
+
+  if (semanticResults.length === 0) return ftsResults.slice(0, topN);
+  if (ftsResults.length === 0) return semanticResults.slice(0, topN);
+
+  // Reciprocal Rank Fusion: each result's contribution from each list is
+  // 1 / (k + rank). Sum across lists, sort, take top N.
+  const fused = new Map(); // source_path|heading → aggregator
+  const keyOf = (r) => `${r.source_path}|${r.heading || ""}`;
+
+  for (let i = 0; i < ftsResults.length; i++) {
+    const r = ftsResults[i];
+    const k = keyOf(r);
+    fused.set(k, {
+      ...r,
+      rrf_score: 1 / (RRF_K + (i + 1)),
+      bm25_rank: i + 1,
+      semantic_rank: null,
+      semantic_score: null,
+      retrieval: "hybrid",
+    });
+  }
+  for (let i = 0; i < semanticResults.length; i++) {
+    const r = semanticResults[i];
+    const k = keyOf(r);
+    if (fused.has(k)) {
+      const f = fused.get(k);
+      f.rrf_score += 1 / (RRF_K + (i + 1));
+      f.semantic_rank = i + 1;
+      f.semantic_score = r.semantic_score;
+    } else {
+      fused.set(k, {
+        source_path: r.source_path,
+        kind: r.kind,
+        title: r.title,
+        heading: r.heading,
+        snippet: r.snippet,
+        score: r.score,
+        rrf_score: 1 / (RRF_K + (i + 1)),
+        bm25_rank: null,
+        semantic_rank: i + 1,
+        semantic_score: r.semantic_score,
+        retrieval: "hybrid",
+      });
+    }
+  }
+
+  const arr = Array.from(fused.values()).sort((a, b) => b.rrf_score - a.rrf_score).slice(0, topN);
+  if (arr.length === 0) return [];
+  const bestRrf = arr[0].rrf_score;
+  return arr.map((r) => ({
+    source_path: r.source_path,
+    kind: r.kind,
+    title: r.title,
+    heading: r.heading,
+    snippet: r.snippet,
+    score: r.score,
+    relevance: bestRrf === 0 ? 1 : Math.max(0.01, Math.min(1, r.rrf_score / bestRrf)),
+    retrieval: "hybrid",
+    bm25_rank: r.bm25_rank,
+    semantic_rank: r.semantic_rank,
+    semantic_score: r.semantic_score,
+  }));
+}
+
+/** Produce a snippet centered on the first query term that appears in the body. */
+function bareSnippet(body, query) {
+  const text = String(body || "").replace(/\s+/g, " ").trim();
+  if (!text) return "";
+  const terms = String(query).toLowerCase().split(/\s+/).filter((t) => t.length > 3);
+  let idx = -1;
+  for (const t of terms) {
+    idx = text.toLowerCase().indexOf(t);
+    if (idx >= 0) break;
+  }
+  if (idx < 0) idx = 0;
+  const start = Math.max(0, idx - 60);
+  const end = Math.min(text.length, idx + 200);
+  return (start > 0 ? "…" : "") + text.slice(start, end) + (end < text.length ? "…" : "");
+}
+
+/**
+ * @returns {Promise<{ files: number, chunks: number, by_kind: object, db_size_bytes: number }>}
+ */
+export async function stats(opts = {}) {
+  const Database = await loadDatabase();
+  const root = resolve(opts.agentRoot || process.env.AGENT_ROOT || process.cwd());
+  const dbPath = join(root, FTS_DB_REL);
+  if (!existsSync(dbPath)) return { files: 0, chunks: 0, by_kind: {}, db_size_bytes: 0 };
+  const { statSync } = await import("node:fs");
+  const db = new Database(dbPath, { readonly: true });
+  try {
+    const files = db.prepare("SELECT COUNT(*) AS n FROM files").get().n;
+    const chunks = db.prepare("SELECT COUNT(*) AS n FROM chunks").get().n;
+    const byKind = {};
+    for (const row of db.prepare("SELECT kind, COUNT(*) AS n FROM files GROUP BY kind").all()) {
+      byKind[row.kind] = row.n;
+    }
+    return { files, chunks, by_kind: byKind, db_size_bytes: statSync(dbPath).size };
+  } finally { db.close(); }
+}
+
+export default { search, stats };