pi-context-engine 0.1.4 → 0.3.0

package/src/storage/graph-store.ts

@@ -24,6 +24,10 @@ CREATE TABLE IF NOT EXISTS edges (
 CREATE INDEX IF NOT EXISTS idx_edges_source ON edges (source_id);
 CREATE INDEX IF NOT EXISTS idx_edges_target ON edges (target_id);
 CREATE INDEX IF NOT EXISTS idx_nodes_file   ON nodes  (file_path);
+CREATE TABLE IF NOT EXISTS pagerank (
+	file_path TEXT PRIMARY KEY,
+	score     REAL NOT NULL
+);
 `;
 
 type Sqlite = SqliteHandle;
@@ -85,6 +89,64 @@ export class GraphStore {
 		return rows.map(rowToNode);
 	}
 
+	/**
+	 * PR #7 review fix (incremental call graph): every function/class node
+	 * across ALL indexed files, so calls-edge derivation on an incremental
+	 * run still sees symbols owned by files that did not change this run.
+	 */
+	listSymbolNodes(): GraphNode[] {
+		const rows = this.db
+			.prepare(
+				"SELECT id, node_type, name, file_path, properties FROM nodes WHERE node_type IN ('function','class')",
+			)
+			.all() as Record<string, unknown>[];
+		return rows.map(rowToNode);
+	}
+
+	/**
+	 * R5: structured symbol lookup for the multi-path retriever. Matches a
+	 * chunk-level node whose name EQUALS a queried identifier, ENDS with
+	 * `.identifier` (so `charge` finds `PaymentGateway.charge`), or STARTS with
+	 * `identifier.` (so `PaymentGateway` surfaces its methods). Exact matches
+	 * rank first; results are deterministic (name length as the tiebreak) so
+	 * RRF fusion is reproducible.
+	 */
+	findSymbols(names: string[], limit = 20): GraphNode[] {
+		if (names.length === 0) return [];
+		const conds: string[] = [];
+		const params: unknown[] = [];
+		for (const raw of names) {
+			const name = raw.trim().toLowerCase();
+			if (name.length < 3) continue;
+			conds.push("(lower(name) = ? OR lower(name) LIKE ? OR lower(name) LIKE ?)");
+			params.push(name, `%.${name}`, `${name}.%`);
+		}
+		if (conds.length === 0) return [];
+		params.push(Math.max(limit * 3, 30)); // overfetch: JS ranks exact-first, then cuts
+		let rows: Record<string, unknown>[];
+		try {
+			rows = this.db
+				.prepare(
+					`SELECT id, node_type, name, file_path, properties FROM nodes
+					 WHERE node_type IN ('function', 'class', 'module', 'doc')
+					 AND (${conds.join(" OR ")})
+					 LIMIT ?`,
+				)
+				.all(...params) as Record<string, unknown>[];
+		} catch {
+			return [];
+		}
+		const exact = new Set(names.map((n) => n.trim().toLowerCase()));
+		const nodes = rows.map(rowToNode);
+		nodes.sort((a, b) => {
+			const ax = exact.has(a.name.toLowerCase()) ? 0 : 1;
+			const bx = exact.has(b.name.toLowerCase()) ? 0 : 1;
+			if (ax !== bx) return ax - bx;
+			return a.name.length - b.name.length;
+		});
+		return nodes.slice(0, limit);
+	}
+
 	neighborsForFiles(filePaths: string[], edgeTypes: EdgeType[], nodeTypes?: NodeType[]): GraphNode[] {
 		if (filePaths.length === 0 || edgeTypes.length === 0) return [];
 		const fpPh = filePaths.map(() => "?").join(",");
@@ -107,6 +169,86 @@ export class GraphStore {
 		return rows.map(rowToNode);
 	}
 
+	/**
+	 * R7: PageRank over the file-level projection of the graph (imports +
+	 * calls edges collapsed to file→file). Scores are normalised so the
+	 * best-connected file is 1.0; the retriever applies them as a bounded
+	 * multiplicative prior. Hub files genuinely central to a codebase —
+	 * types, shared utils, the engine — consistently rank above leaves, which
+	 * is exactly the tie-break signal fusion needs when lexical and vector
+	 * evidence disagree.
+	 */
+	computePageRank(): void {
+		const nodeFile = new Map<string, string>();
+		for (const row of this.db.prepare("SELECT id, file_path FROM nodes").all() as {
+			id: string;
+			file_path: string;
+		}[]) {
+			nodeFile.set(row.id, row.file_path);
+		}
+		const edgeRows = this.db
+			.prepare("SELECT source_id, target_id FROM edges WHERE edge_type IN ('imports', 'calls')")
+			.all() as { source_id: string; target_id: string }[];
+
+		const nodes = new Set<string>();
+		for (const fp of nodeFile.values()) if (fp) nodes.add(fp);
+		const out = new Map<string, Set<string>>();
+		for (const e of edgeRows) {
+			const s = nodeFile.get(e.source_id);
+			const tgt = nodeFile.get(e.target_id);
+			if (!s || !tgt || s === tgt) continue;
+			let set = out.get(s);
+			if (!set) out.set(s, (set = new Set()));
+			set.add(tgt);
+		}
+
+		const run = this.db.transaction(() => {
+			this.db.exec("DELETE FROM pagerank");
+			if (nodes.size === 0) return;
+			const d = 0.85;
+			const n = nodes.size;
+			const incoming = new Map<string, string[]>();
+			for (const [s, targets] of out) {
+				for (const tgt of targets) {
+					let arr = incoming.get(tgt);
+					if (!arr) incoming.set(tgt, (arr = []));
+					arr.push(s);
+				}
+			}
+			let rank = new Map<string, number>();
+			for (const node of nodes) rank.set(node, 1 / n);
+			for (let iter = 0; iter < 25; iter++) {
+				const next = new Map<string, number>();
+				let dangling = 0;
+				for (const node of nodes) {
+					const o = out.get(node);
+					if (!o || o.size === 0) dangling += rank.get(node) ?? 0;
+				}
+				for (const node of nodes) {
+					let s = (1 - d) / n + (d * dangling) / n;
+					for (const src of incoming.get(node) ?? []) {
+						s += (d * (rank.get(src) ?? 0)) / (out.get(src)?.size ?? 1);
+					}
+					next.set(node, s);
+				}
+				rank = next;
+			}
+			const max = Math.max(...rank.values(), 1e-12);
+			const ins = this.db.prepare("INSERT OR REPLACE INTO pagerank (file_path, score) VALUES (?, ?)");
+			for (const [fp, s] of rank) ins.run(fp, s / max);
+		});
+		run();
+	}
+
+	/** Read of the PageRank table (file_path → score in [0,1]). */
+	pageRanks(): Map<string, number> {
+		const rows = this.db.prepare("SELECT file_path, score FROM pagerank").all() as {
+			file_path: string;
+			score: number;
+		}[];
+		return new Map(rows.map((r) => [r.file_path, r.score]));
+	}
+
 	getNodesByType(nodeType: NodeType): GraphNode[] {
 		const rows = this.db
 			.prepare("SELECT id, node_type, name, file_path, properties FROM nodes WHERE node_type = ?")
@@ -126,15 +268,28 @@ export class GraphStore {
 				const nodeIds = (
 					this.db.prepare(`SELECT id FROM nodes WHERE file_path IN (${ph})`).all(...batch) as {
 						id: string;
-					}[]
+				}[]
 				).map((r) => r.id);
+				// R7: file_* node ids are STABLE across re-ingest of the same
+				// file, so inbound imports edges from other files stay valid.
+				// Only delete edges pointing at this file's chunk-level nodes
+				// (their ids change when content changes). Without this,
+				// re-indexing b.ts silently deleted every a.ts→b.ts imports
+				// edge unless a.ts changed too.
+				const chunkNodeIds = nodeIds.filter((id) => !id.startsWith("file_"));
 				for (const idBatch of batchedParams(nodeIds)) {
 					const idPh = idBatch.map(() => "?").join(",");
 					this.db
-						.prepare(`DELETE FROM edges WHERE source_id IN (${idPh}) OR target_id IN (${idPh})`)
-						.run(...idBatch, ...idBatch);
-					this.db.prepare(`DELETE FROM nodes WHERE id IN (${idPh})`).run(...idBatch);
+						.prepare(`DELETE FROM nodes WHERE id IN (${idPh})`).run(...idBatch);
+				}
+				const srcPh = nodeIds.map(() => "?").join(",");
+				this.db.prepare("DELETE FROM edges WHERE source_id IN (" + srcPh + ")").run(...nodeIds);
+				if (chunkNodeIds.length > 0) {
+					const tgtPh = chunkNodeIds.map(() => "?").join(",");
+					this.db.prepare("DELETE FROM edges WHERE target_id IN (" + tgtPh + ")").run(...chunkNodeIds);
 				}
+				const filePh = batch.map(() => "?").join(",");
+				this.db.prepare("DELETE FROM pagerank WHERE file_path IN (" + filePh + ")").run(...batch);
 			}
 		});
 		run([...batchedParams(filePaths)]);
@@ -143,6 +298,7 @@ export class GraphStore {
 	clear(): void {
 		this.db.exec("DELETE FROM edges");
 		this.db.exec("DELETE FROM nodes");
+		this.db.exec("DELETE FROM pagerank");
 	}
 
 	close(): void {

package/src/storage/local-backend.ts

@@ -11,6 +11,8 @@ export class LocalBackend {
 	private vectorStore: VectorStore;
 	private ftsStore: FTSStore;
 	private graphStore: GraphStore;
+	/** R7: cached PageRank table; refreshed by recomputePriors(). */
+	private pageRankCache: Map<string, number> | null = null;
 
 	constructor(basePath: string) {
 		this.vectorStore = new VectorStore(path.join(basePath, "vectors"));
@@ -18,10 +20,6 @@ export class LocalBackend {
 		this.graphStore = new GraphStore(path.join(basePath, "graph"));
 	}
 
-	get metricRebuilt(): boolean {
-		return this.vectorStore.metricRebuilt;
-	}
-
 	ingest(chunks: Chunk[], nodes: GraphNode[], edges: GraphEdge[]): void {
 		this.vectorStore.ingest(chunks);
 		this.ftsStore.ingest(chunks);
@@ -56,6 +54,38 @@ export class LocalBackend {
 		return this.vectorStore.getById(chunkId);
 	}
 
+	/** R7: recompute file PageRank after indexing; retrieval reads it as a prior. */
+	recomputePriors(): void {
+		this.graphStore.computePageRank();
+		this.pageRankCache = null;
+	}
+
+	/** R7: cached file→[0,1] PageRank (empty until first recomputePriors). */
+	pageRanks(): Map<string, number> {
+		if (!this.pageRankCache) this.pageRankCache = this.graphStore.pageRanks();
+		return this.pageRankCache;
+	}
+
+	/** R5: symbol lookup over graph node names (exact / `.suffix`). */
+	findSymbols(names: string[], limit = 20): GraphNode[] {
+		return this.graphStore.findSymbols(names, limit);
+	}
+
+	/** R9: graph nodes of one file (sibling-symbol hints). */
+	getNodesByFile(filePath: string): GraphNode[] {
+		return this.graphStore.getNodesByFile(filePath);
+	}
+
+	/** PR #7 review: all function/class nodes, for incremental call-graph derivation. */
+	listSymbolNodes(): GraphNode[] {
+		return this.graphStore.listSymbolNodes();
+	}
+
+	/** R5: chunks of one file in source order. */
+	chunksByFile(filePath: string, limit = 10): Chunk[] {
+		return this.vectorStore.byFile(filePath, limit);
+	}
+
 	getChunksByIds(chunkIds: string[]): Chunk[] {
 		return this.vectorStore.getChunksByIds(chunkIds);
 	}
@@ -69,6 +99,7 @@ export class LocalBackend {
 		this.vectorStore.deleteByFiles(filePaths);
 		this.ftsStore.deleteByFiles(filePaths);
 		this.graphStore.deleteByFiles(filePaths);
+		this.pageRankCache = null;
 	}
 
 	countChunks(): number {
@@ -91,6 +122,7 @@ export class LocalBackend {
 		this.vectorStore.clear();
 		this.ftsStore.clear();
 		this.graphStore.clear();
+		this.pageRankCache = null;
 	}
 
 	close(): void {

package/src/storage/vector-store.ts

@@ -13,6 +13,7 @@ import fs from "node:fs";
 import path from "node:path";
 import { openDatabase, type SqliteHandle } from "../sqlite.js";
 import { batchedParams } from "../utils.js";
+import { parseContextHeaderUnit } from "../indexer/context-header.js";
 import type { Chunk, ChunkType } from "../types.js";
 
 const MAX_CONTENT_CHARS = 5_000;
@@ -26,8 +27,6 @@ type Sqlite = SqliteHandle;
 export class VectorStore {
 	private db: Sqlite;
 	private dim: number | null = null;
-	/** Set when a legacy L2 vec table was wiped — pipeline forces a full reindex. */
-	metricRebuilt = false;
 	constructor(dbPath: string) {
 		fs.mkdirSync(dbPath, { recursive: true });
 		// Without sqlite-vec the store still keeps chunk metadata; search()
@@ -36,10 +35,6 @@ export class VectorStore {
 		this.ensureTables();
 	}
 
-	private get vecAvailable(): boolean {
-		return this.dim !== null || this.vecTableExists();
-	}
-
 	private vecTableExists(): boolean {
 		const row = this.db
 			.prepare("SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'")
@@ -60,10 +55,6 @@ export class VectorStore {
 				modified_ts REAL
 			)
 		`);
-		const cols = (this.db.pragma("table_info(chunks)") as { name: string }[]).map((c) => c.name);
-		if (!cols.includes("modified_ts")) {
-			this.db.exec("ALTER TABLE chunks ADD COLUMN modified_ts REAL");
-		}
 		this.db.exec("CREATE INDEX IF NOT EXISTS idx_chunks_file_path ON chunks(file_path)");
 		this.db.exec(`
 			CREATE TABLE IF NOT EXISTS chunk_compressions (
@@ -74,25 +65,19 @@ export class VectorStore {
 			)
 		`);
 
-		const row = this.db
-			.prepare("SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'")
-			.get() as { sql: string | null } | undefined;
-		if (row) {
-			if (!(row.sql ?? "").includes("distance_metric=cosine")) {
-				this.db.exec("DROP TABLE IF EXISTS chunks_vec");
-				this.db.exec("DELETE FROM chunks");
-				this.db.exec("DELETE FROM chunk_compressions");
-				this.metricRebuilt = true;
-			} else {
-				const r = this.db.prepare("SELECT rowid FROM chunks_vec LIMIT 1").get() as
-					| { rowid: number }
-					| undefined;
-				if (r) {
-					const len = this.db
-						.prepare("SELECT vec_length(embedding) AS len FROM chunks_vec LIMIT 1")
-						.get() as { len: number };
-					this.dim = len.len;
-				}
+		if (this.vecTableExists()) {
+			// The table exists — read its dimension. There is no migration path
+			// for tables this package did not write (older metric/column shapes):
+			// an unrecognized shape cannot occur from our own DDL, and foreign
+			// stores are unsupported by design.
+			const r = this.db.prepare("SELECT rowid FROM chunks_vec LIMIT 1").get() as
+				| { rowid: number }
+				| undefined;
+			if (r) {
+				const len = this.db
+					.prepare("SELECT vec_length(embedding) AS len FROM chunks_vec LIMIT 1")
+					.get() as { len: number };
+				this.dim = len.len;
 			}
 		}
 	}
@@ -144,6 +129,12 @@ export class VectorStore {
 		};
 		if (row.modified_ts != null) chunk.metadata["modified_ts"] = row.modified_ts;
 		if (distance !== undefined) chunk.metadata["_distance"] = distance;
+		// Restore `unit` from the R1 context header line so served chunks carry
+		// their symbol (refine hints filter siblings by it; the header is the
+		// single source of truth for indexed content). No header → no unit,
+		// same as before.
+		const unit = parseContextHeaderUnit(chunk.content);
+		if (unit) chunk.metadata["unit"] = unit;
 		return chunk;
 	}
 
@@ -327,6 +318,21 @@ export class VectorStore {
 		}
 	}
 
+	/** R5: chunks of one file in source order (file-hint retrieval path). */
+	byFile(filePath: string, limit = 10): Chunk[] {
+		try {
+			const rows = this.db
+				.prepare(
+					`SELECT id, content, chunk_type, file_path, start_line, end_line, language, modified_ts
+					 FROM chunks WHERE file_path = ? ORDER BY start_line ASC LIMIT ?`,
+				)
+				.all(filePath, limit) as Record<string, unknown>[];
+			return rows.map((r) => this.rowToChunk(r));
+		} catch {
+			return [];
+		}
+	}
+
 	close(): void {
 		this.db.close();
 	}

package/src/utils.ts

@@ -79,6 +79,19 @@ export function epochSeconds(): number {
 	return Math.floor(Date.now() / 1000);
 }
 
+/**
+ * Cooperative macrotask yield for long-running indexing work. ONNX
+ * inference, tokenization and tree-sitter parsing all run on the main
+ * thread, so a long stretch of synchronous work starves pi's UI loop
+ * (measured by scripts/bench-event-loop.mts: ~0.5-0.9 s stalls per 64-text
+ * embed batch). setImmediate keeps the yield cheap — same loop turn's check
+ * phase — while still letting pending timers and I/O (keystrokes, renders)
+ * run between units of work.
+ */
+export function yieldToEventLoop(): Promise<void> {
+	return new Promise((resolve) => setImmediate(resolve));
+}
+
 /**
  * Fractional seconds. Session-capture ordering needs sub-second resolution:
  * several decisions are commonly recorded within the same wall-clock second,