pi-mega-compact 0.4.23 → 0.4.25

package/src/store/vectorIndex.ts

@@ -0,0 +1,243 @@
+/**
+ * vectorIndex.ts — Slice 2 async vector index (PGlite/pgvector HNSW).
+ *
+ * A REDUNDANT, additive, ASYNC index layered over the synchronous node:sqlite
+ * store (which remains the authoritative source of truth). The sync linear
+ * cosine scan over `embedding_blob` stays the DEFAULT recall path; this index
+ * exists only to provide real cross-repo / cross-session HNSW nearest-neighbor
+ * recall. It is best-effort and non-fatal: any init/write failure degrades to
+ * the sync scan and must NEVER break add(), compaction, or extension load.
+ *
+ * PREVENT-PI-004: PGlite is WASM Postgres — fully local, zero network.
+ *
+ * Index topology (decision 2026-07-15): ONE global PGlite DB, `repo_id` is a
+ * first-class column. `searchAsync(q, k, {repoId?})` → omit repoId for cross-repo
+ * NN, pass repoId to scope to a single repo. The sync store is per-repo (state
+ * dir); this global index is the thing that makes cross-repo recall possible.
+ */
+
+import { homedir } from "node:os";
+import { join } from "node:path";
+import { mkdirSync } from "node:fs";
+
+// PGlite + pgvector are script-free WASM (no native build) → survive pi's
+// install-script block. Imported lazily so a missing/broken package degrades
+// gracefully instead of crashing module load.
+import { PGlite, type PGlite as PGliteInstance } from "@electric-sql/pglite";
+import { vector } from "@electric-sql/pglite-pgvector";
+
+/** Vector dimension produced by the default TrigramEmbedder (src/embedder.ts). */
+export const EMBEDDING_DIM = 512;
+
+/** A single recall hit returned by the async index. */
+export interface VectorIndexHit {
+  repoId: string;
+  sessionId: string;
+  checkpointId: string;
+  /** Cosine similarity in [0,1] (1 = identical). */
+  score: number;
+}
+
+let db: PGliteInstance | undefined;
+let initPromise: Promise<PGliteInstance | undefined> | undefined;
+let disabled = false;
+let warned = false;
+
+function indexDir(): string {
+  const override = process.env.MEGACOMPACT_VECTOR_INDEX_DIR;
+  if (override && override.trim() !== "") return override;
+  try {
+    return join(homedir(), ".pi", "mega-compact-vector");
+  } catch {
+    return join("/tmp", ".mega-compact-vector");
+  }
+}
+
+function logWarn(msg: string): void {
+  // Never throw — degradation is the whole point. One warning per process.
+  if (warned) return;
+  warned = true;
+  try {
+    console.warn(`[mega-compact:vectorIndex] ${msg} (falling back to sync scan)`);
+  } catch {
+    /* ignore */
+  }
+}
+
+/** Honor the emergency kill-switch. When set, the index is fully disabled. */
+export function isVectorIndexDisabled(): boolean {
+  return (
+    disabled ||
+    process.env.MEGACOMPACT_PGLITE_DISABLED === "true" ||
+    process.env.MEGACOMPACT_PGLITE_DISABLED === "1"
+  );
+}
+
+/**
+ * Lazily open + schema-init the global PGlite DB. Idempotent and safe to call
+ * from many places. Returns undefined when disabled/unavailable so callers can
+ * fall back to the synchronous scan. Never throws.
+ */
+export function initVectorIndex(): Promise<PGliteInstance | undefined> {
+  if (isVectorIndexDisabled()) return Promise.resolve(undefined);
+  if (db) return Promise.resolve(db);
+  if (initPromise) return initPromise;
+  initPromise = (async (): Promise<PGliteInstance | undefined> => {
+    try {
+      const dir = indexDir();
+      mkdirSync(dir, { recursive: true });
+      const pg = await new PGlite({
+        dataDir: dir,
+        extensions: { vector },
+      });
+      await pg.exec("CREATE EXTENSION IF NOT EXISTS vector;");
+      await pg.exec(`
+        CREATE TABLE IF NOT EXISTS vector_index (
+          repo_id       TEXT NOT NULL,
+          session_id    TEXT NOT NULL,
+          checkpoint_id TEXT NOT NULL,
+          embedding     vector(${EMBEDDING_DIM}) NOT NULL,
+          PRIMARY KEY (repo_id, session_id, checkpoint_id)
+        );
+      `);
+      // HNSW index over cosine distance for fast NN. Created idempotently.
+      await pg.exec(
+        "CREATE INDEX IF NOT EXISTS vector_index_hnsw ON vector_index USING hnsw (embedding vector_cosine_ops);",
+      );
+      db = pg;
+      return pg;
+    } catch (err) {
+      disabled = true;
+      logWarn(`init failed: ${err instanceof Error ? err.message : String(err)}`);
+      return undefined;
+    }
+  })();
+  return initPromise;
+}
+
+function toVectorLiteral(v: number[]): string {
+  // pgvector text form: [a,b,c]. Guard against NaN/Inf for a clean literal.
+  const parts = v.map((x) => (Number.isFinite(x) ? x : 0));
+  return `[${parts.join(",")}]`;
+}
+
+/**
+ * Best-effort upsert of one checkpoint embedding into the global index.
+ * Dimension-mismatched vectors (e.g. a BYO embedder with dim ≠ 512) are skipped
+ * rather than corrupting the index. Fire-and-forget: resolved promise only;
+ * callers must NOT await this on the sync path. Never throws.
+ */
+export async function upsertEmbedding(
+  repoId: string,
+  sessionId: string,
+  checkpointId: string,
+  embedding: number[],
+): Promise<void> {
+  if (isVectorIndexDisabled()) return;
+  if (!embedding || embedding.length !== EMBEDDING_DIM) {
+    // Dimension guard: skip without corrupting the fixed-dim index.
+    return;
+  }
+  try {
+    const pg = await initVectorIndex();
+    if (!pg) return;
+    const lit = toVectorLiteral(embedding);
+    await pg.query(
+      `INSERT INTO vector_index (repo_id, session_id, checkpoint_id, embedding)
+       VALUES ($1, $2, $3, $4::vector)
+       ON CONFLICT (repo_id, session_id, checkpoint_id)
+       DO UPDATE SET embedding = EXCLUDED.embedding;`,
+      [repoId, sessionId, checkpointId, lit],
+    );
+  } catch (err) {
+    disabled = true;
+    logWarn(`upsert failed: ${err instanceof Error ? err.message : String(err)}`);
+  }
+}
+
+export interface SearchAsyncOpts {
+  /** When provided, scope the NN search to a single repo; omit for cross-repo. */
+  repoId?: string;
+  /** Max hits (default 3). */
+  k?: number;
+}
+
+/**
+ * Cross-repo (or single-repo) HNSW nearest-neighbor search. Returns hits sorted
+ * by descending similarity. Never throws — on any failure returns [].
+ */
+export async function searchAsync(
+  query: number[],
+  opts: SearchAsyncOpts = {},
+): Promise<VectorIndexHit[]> {
+  if (isVectorIndexDisabled() || !query || query.length !== EMBEDDING_DIM) return [];
+  const k = opts.k ?? 3;
+  const repoId = opts.repoId;
+  try {
+    const pg = await initVectorIndex();
+    if (!pg) return [];
+    const lit = toVectorLiteral(query);
+    const params: unknown[] = [lit, k];
+    let sql =
+      "SELECT repo_id, session_id, checkpoint_id, 1 - (embedding <=> $1::vector) AS score " +
+      "FROM vector_index";
+    if (repoId) {
+      sql += " WHERE repo_id = $3";
+      params.push(repoId);
+    }
+    sql += " ORDER BY embedding <=> $1::vector LIMIT $2";
+    const res = await pg.query(sql, params);
+    return res.rows.map((r: any) => ({
+      repoId: r.repo_id as string,
+      sessionId: r.session_id as string,
+      checkpointId: r.checkpoint_id as string,
+      score: r.score as number,
+    }));
+  } catch (err) {
+    disabled = true;
+    logWarn(`search failed: ${err instanceof Error ? err.message : String(err)}`);
+    return [];
+  }
+}
+
+/** Close the index (test teardown / shutdown). Safe to call when unopened. */
+export async function closeVectorIndex(): Promise<void> {
+  if (db) {
+    try {
+      await db.close();
+    } catch {
+      /* ignore */
+    }
+  }
+  db = undefined;
+  initPromise = undefined;
+}
+
+/**
+ * Rebuild the entire index from the authoritative node:sqlite store. Used for
+ * backfill + DR. `enumerateRepoStateDirs` yields each repo's state dir; we read
+ * its checkpoint embeddings and bulk upsert. Best-effort: counts successes and
+ * skips failures. Returns {upserted, errors}.
+ */
+export async function rebuildFromSqlite(
+  enumerateRepoStateDirs: () => Iterable<{ repoId: string; stateDir: string }>,
+  readCheckpoints: (
+    stateDir: string,
+  ) => Iterable<{ sessionId: string; checkpointId: string; embedding: number[] }>,
+): Promise<{ upserted: number; errors: number }> {
+  let upserted = 0;
+  let errors = 0;
+  const pg = await initVectorIndex();
+  if (!pg) return { upserted, errors: 1 };
+  for (const repo of enumerateRepoStateDirs()) {
+    for (const cp of readCheckpoints(repo.stateDir)) {
+      try {
+        await upsertEmbedding(repo.repoId, cp.sessionId, cp.checkpointId, cp.embedding);
+        upserted++;
+      } catch {
+        errors++;
+      }
+    }
+  }
+  return { upserted, errors };
+}

package/src/vectorStore.ts

@@ -26,6 +26,7 @@ import {
   listCheckpoints,
   nextCheckpointId,
   upsertCheckpoint,
+  getCheckpoint,
   loadSessionState,
   saveSessionState,
   upsertMinhashSignature,
@@ -38,6 +39,12 @@ import {
   repoStats as repoStatsFromStore,
   dataInvariantStats,
 } from "./store/sqlite.js";
+import {
+  upsertEmbedding as indexUpsertEmbedding,
+  initVectorIndex,
+  searchAsync as vectorIndexSearch,
+  type VectorIndexHit,
+} from "./store/vectorIndex.js";
 import { rehydrateRaptorTree } from "./dedup/raptor/index.js";
 import { stagedExpansion } from "./dedup/raptor/retrieval.js";
 import { migrateJsonToSqlite } from "./store/migrate.js";
@@ -98,6 +105,13 @@ export class VectorStore {
   private readonly cfg: DedupConfigShape;
   /** Optional monitoring target (Sprint 14). Undefined → no monitoring. */
   private readonly eventsPath?: string;
+  /**
+   * Repo key for the async PGlite vector index (Slice 2). We use the stateDir
+   * itself as the repo id — it is already unique per repo and available here
+   * without crossing into the pi-runtime layer (src/ stays pi-agnostic). The
+   * global index keys on repoId so recall can span repos.
+   */
+  private readonly repoId: string;
 
   constructor(
     opts: {
@@ -110,10 +124,13 @@ export class VectorStore {
       config?: DedupConfigShape;
       /** Optional events.log path for decision monitoring (Sprint 14). */
       eventsPath?: string;
+      /** Repo id for the async cross-repo vector index. Defaults to stateDir. */
+      repoId?: string;
     } = {},
   ) {
     this.embedder = opts.embedder ?? defaultEmbedder();
     this.stateDir = opts.stateDir ?? getStateDir();
+    this.repoId = opts.repoId ?? this.stateDir;
     // Sprint 14: all tier flags/thresholds flow from the single config source
     // (DedupConfig). The legacy opts.dedupSim / opts.l2Enabled remain accepted
     // for backward-compat callers but flags are authoritative via `cfg`.
@@ -371,6 +388,12 @@ export class VectorStore {
     // Cumulative store-wide dedup accounting (attempt, not collapsed).
     bumpDedupStats(false, this.stateDir);
     onTier?.({ tier: "new", status: "stored" });
+    // Slice 2: best-effort, fire-and-forget mirror of this new checkpoint into
+    // the async global PGlite/HNSW index. NEVER awaited — must not block or
+    // throw into the synchronous add() path. On failure the index degrades to
+    // the sync scan (handled inside vectorIndex). The node:sqlite store remains
+    // authoritative; the index is rebuildable from it at any time.
+    void indexUpsertEmbedding(this.repoId, sessionId, checkpointId, checkpoint.embedding);
     return { checkpoint, deduped: false };
   }
 
@@ -472,6 +495,60 @@ export class VectorStore {
     return ranked;
   }
 
+  /**
+   * Slice 2: async cross-repo (or single-repo) recall via the PGlite/HNSW index.
+   *
+   * This is the ONLY async recall surface and is a BONUS path — the synchronous
+   * `search()` above remains the default. `opts.repoId` scopes to one repo; omit
+   * it for cross-repo nearest-neighbor recall (the headline capability the sync
+   * per-session scan cannot provide).
+   *
+   * Best-effort: if the index is disabled/empty/failing, we fall back to the
+   * synchronous per-session `search()` for THIS repo so callers always get a
+   * sensible result. Hydrates each hit's StoredCheckpoint from the authoritative
+   * node:sqlite store (the hit's repoId doubles as that repo's stateDir), then
+   * MMR-dedupes the merged set.
+   */
+  async searchAsync(
+    sessionId: string,
+    query: string,
+    k = 3,
+    opts: { repoId?: string; crossRepo?: boolean } = {},
+  ): Promise<SearchHit[]> {
+    const sid = normalizeSessionId(sessionId);
+    const qv = this.embedder.embed(query);
+    // repoId filter: explicit opts.repoId wins; else this repo unless crossRepo.
+    const repoId = opts.repoId ?? (opts.crossRepo ? undefined : this.repoId);
+    let indexHits: VectorIndexHit[] = [];
+    try {
+      await initVectorIndex();
+      indexHits = await vectorIndexSearch(qv, { k: Math.max(k * 2, k), repoId });
+    } catch {
+      indexHits = [];
+    }
+    if (indexHits.length === 0) {
+      // Index empty/unavailable → synchronous per-session fallback (this repo).
+      return this.search(sid, query, k);
+    }
+    // Hydrate each index hit from the authoritative node:sqlite store. repoId is
+    // that repo's stateDir, so cross-repo hits resolve against their own store.
+    const hydrated: SearchHit[] = [];
+    for (const h of indexHits) {
+      const cp = getCheckpoint(h.sessionId, h.checkpointId, h.repoId);
+      if (cp && cp.dedupStatus !== "removed") {
+        hydrated.push({ checkpoint: cp, score: h.score });
+      }
+    }
+    if (hydrated.length === 0) return this.search(sid, query, k);
+    // MMR-dedupe the merged candidate set for diversity (mirrors sync search).
+    const mmrItems: MmrItem<SearchHit>[] = hydrated.map((h) => ({
+      item: h,
+      vector: h.checkpoint.embedding,
+      relevance: h.score,
+    }));
+    return mmrRerank(mmrItems, k, this.cfg.MMR_LAMBDA);
+  }
+
   /**
    * Serve the RAPTOR tree for a query (Fix D): rehydrate the persisted tree and
    * return its staged-expansion leaf hits as SearchHits. Returns [] when no tree