pi-context-engine 0.1.3 → 0.2.0

package/package.json

@@ -1,81 +1,85 @@
 {
-	"name": "pi-context-engine",
-	"version": "0.1.3",
-	"description": "Semantic code context for pi: tree-sitter index + hybrid vector/BM25 retrieval + cross-session memory, all local. Port of code-context-engine (MIT, elara-labs) as a native pi extension.",
-	"keywords": [
-		"pi-package",
-		"pi-extension",
-		"context-engine",
-		"code-search",
-		"semantic-search",
-		"vector-search",
-		"tree-sitter",
-		"token-savings",
-		"memory"
-	],
-	"license": "MIT",
-	"type": "module",
-	"main": "src/extension.ts",
-	"files": [
-		"src",
-		"docs",
-		"README.md",
-		"LICENSE"
-	],
-	"scripts": {
-		"test": "vitest run",
-		"test:watch": "vitest watch",
-		"typecheck": "tsc --noEmit",
-		"check": "npm run typecheck && npm run test",
-		"smoke:env": "node scripts/smoke-env.mjs",
-		"smoke:loader": "node scripts/smoke-loader.mjs",
-		"smoke:real-world": "tsx scripts/smoke-real-world.mjs",
-		"prepublishOnly": "node -e \"if(!process.env.CI)throw new Error('publish via release.yml CI: push a v* tag. Manual npm publish is blocked to guarantee provenance + version sync.')\""
-	},
-	"pi": {
-		"extensions": [
-			"./src/extension.ts"
-		]
-	},
-	"engines": {
-		"node": ">=22.0.0"
-	},
-	"allowScripts": {
-		"better-sqlite3": true
-	},
-	"dependencies": {
-		"better-sqlite3": "^13.0.3",
-		"chokidar": "^5.0.0",
-		"sqlite-vec": "^0.1.9",
-		"tree-sitter-wasms": "^0.1.13",
-		"web-tree-sitter": "^0.22.6",
-		"yaml": "^2.9.1"
-	},
-	"peerDependencies": {
-		"@earendil-works/pi-ai": "*",
-		"@earendil-works/pi-coding-agent": "*",
-		"typebox": "*"
-	},
-	"devDependencies": {
-		"@earendil-works/pi-ai": ">=0.80.0",
-		"@earendil-works/pi-coding-agent": ">=0.99.2",
-		"@types/node": "^22.0.0",
-		"@huggingface/transformers": "^4.3.1",
-		"typebox": "^1.0.0",
-		"typescript": "^5.9.0",
-		"vitest": "^3.1.0"
-	},
-	"publishConfig": {
-		"access": "public",
-		"provenance": true
-	},
-	"author": "alessandro@pungitore.ch",
-	"repository": {
-		"type": "git",
-		"url": "git+https://github.com/pungggi/pi-context-engine.git"
-	},
-	"bugs": {
-		"url": "https://github.com/pungggi/pi-context-engine/issues"
-	},
-	"homepage": "https://github.com/pungggi/pi-context-engine#readme"
+  "name": "pi-context-engine",
+  "version": "0.2.0",
+  "description": "Semantic code context for pi: tree-sitter index + hybrid vector/BM25 retrieval + cross-session memory, all local. Port of code-context-engine (MIT, elara-labs) as a native pi extension.",
+  "keywords": [
+    "pi-package",
+    "pi-extension",
+    "context-engine",
+    "code-search",
+    "semantic-search",
+    "vector-search",
+    "tree-sitter",
+    "token-savings",
+    "memory"
+  ],
+  "license": "MIT",
+  "type": "module",
+  "main": "src/extension.ts",
+  "files": [
+    "src",
+    "docs",
+    "README.md",
+    "LICENSE"
+  ],
+  "scripts": {
+    "test": "vitest run",
+    "test:watch": "vitest watch",
+    "typecheck": "tsc --noEmit",
+    "check": "npm run typecheck && npm run test",
+    "smoke:env": "node scripts/smoke-env.mjs",
+    "smoke:loader": "node scripts/smoke-loader.mjs",
+    "smoke:real-world": "tsx scripts/smoke-real-world.mjs",
+    "eval": "tsx scripts/eval-retrieval.mts",
+    "prepublishOnly": "node -e \"if(!process.env.CI)throw new Error('publish via release.yml CI: push a v* tag. Manual npm publish is blocked to guarantee provenance + version sync.')\""
+  },
+  "pi": {
+    "extensions": [
+      "./src/extension.ts"
+    ]
+  },
+  "engines": {
+    "node": ">=22.0.0"
+  },
+  "allowScripts": {
+    "better-sqlite3": true
+  },
+  "dependencies": {
+    "better-sqlite3": "^13.0.3",
+    "chokidar": "^5.0.0",
+    "sqlite-vec": "^0.1.9",
+    "tree-sitter-wasms": "^0.1.13",
+    "web-tree-sitter": "^0.22.6",
+    "yaml": "^2.9.1"
+  },
+  "optionalDependencies": {
+    "@huggingface/transformers": "^4.3.1"
+  },
+  "peerDependencies": {
+    "@earendil-works/pi-ai": "*",
+    "@earendil-works/pi-coding-agent": "*",
+    "typebox": "*"
+  },
+  "devDependencies": {
+    "@earendil-works/pi-ai": ">=0.80.0",
+    "@earendil-works/pi-coding-agent": ">=0.99.2",
+    "@types/node": "^22.0.0",
+    "typebox": "^1.0.0",
+    "tsx": "^4.23.15",
+    "typescript": "^5.9.0",
+    "vitest": "^3.1.0"
+  },
+  "publishConfig": {
+    "access": "public",
+    "provenance": true
+  },
+  "author": "alessandro@pungitore.ch",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/pungggi/pi-context-engine.git"
+  },
+  "bugs": {
+    "url": "https://github.com/pungggi/pi-context-engine/issues"
+  },
+  "homepage": "https://github.com/pungggi/pi-context-engine#readme"
 }

package/src/compression/compressor.ts

@@ -25,7 +25,16 @@ const PROMPT_MAP: Partial<Record<ChunkType, string>> = {
 };
 
 const TRUNCATION_LIMITS: Record<string, number> = { minimal: 100, standard: 300, full: 800 };
-const FULL_PASSTHROUGH_THRESHOLD = 0.8;
+// Query-aware selection: keep matched lines within the level budget, never more.
+const MAX_QUERY_LINES = 40;
+/** Words that must not win the line budget just because they recur. */
+const QUERY_NOISE_TERMS = new Set([
+	"the", "and", "for", "with", "that", "this", "from", "are", "was", "not",
+	"you", "all", "can", "has", "its", "but", "how", "what", "when", "where",
+	"which", "who", "why", "does", "did", "should", "would", "into", "than",
+	"then", "them", "they", "use", "used", "using", "show", "find", "call",
+	"calls", "make", "made", "need", "needs", "get", "set", "returns", "return",
+]);
 // TTL for the Ollama liveness probe so compress() calls don't all pay it.
 const OLLAMA_PROBE_TTL_MS = 30_000;
 
@@ -153,35 +162,58 @@ export class Compressor {
 		return this.probeInFlight;
 	}
 
-	async compress(chunks: Chunk[], level = "standard"): Promise<Chunk[]> {
+	/**
+	 * Compress chunks for serving.
+	 *
+	 * `query` enables R3 query-aware selection: when a truncation would cut
+	 * lines (minimal/standard without the LLM path), keep the lines that
+	 * mention the query's terms within the same character budget instead of
+	 * the first N lines. That stops the compressor from deleting the only
+	 * line that answers the question, which fixed head-cuts did routinely.
+	 *
+	 * `queryAware=false` preserves the old head-truncation behaviour. `full`
+	 * never truncates: the retrieval token budget decides (plan R3), so
+	 * there is nothing for the query to protect at that level.
+	 */
+	async compress(chunks: Chunk[], level = "standard", query?: string, queryAware = true): Promise<Chunk[]> {
 		const ollamaAvailable = await this.isOllamaAvailable();
+		const queryText = (query ?? "").trim();
+		// Query selection only applies where truncation would otherwise cut.
+		// An LLM summary is query-independent (and therefore cacheable), so it
+		// keeps precedence when Ollama is up at standard/full.
+		const querySelection =
+			queryAware && queryText.length > 0 && level !== "full" && !(ollamaAvailable && level !== "minimal");
 		for (const chunk of chunks) {
-			if (level === "full" && this.isFullPassthrough(chunk)) {
+			if (level === "full") {
+				// Budget-driven: the engine's maxTokens packing trims, not us.
 				chunk.compressedContent = chunk.content;
 				continue;
 			}
-			const cached = this.cache?.getCachedCompression(chunk.id, level);
-			if (cached !== null && cached !== undefined) {
-				chunk.compressedContent = cached;
-				continue;
+			if (!querySelection) {
+				const cached = this.cache?.getCachedCompression(chunk.id, level);
+				if (cached !== null && cached !== undefined) {
+					chunk.compressedContent = cached;
+					continue;
+				}
 			}
-			if (ollamaAvailable && level !== "minimal") {
-				chunk.compressedContent = await this.llmCompress(chunk, level);
+			let compressed: string;
+			if (querySelection) {
+				// Never read or write the cache here: its key is (chunk, level)
+				// and would leak a stale selection from another query.
+				compressed = this.queryAwareFallback(chunk, level, queryText);
+			} else if (ollamaAvailable && level !== "minimal") {
+				compressed = await this.llmCompress(chunk, level);
 			} else {
-				chunk.compressedContent = this.fallbackCompress(chunk, level);
+				compressed = this.fallbackCompress(chunk, level);
 			}
-			if (this.cache && chunk.compressedContent) {
-				this.cache.putCachedCompression(chunk.id, level, chunk.compressedContent);
+			chunk.compressedContent = compressed;
+			if (this.cache && compressed && !querySelection) {
+				this.cache.putCachedCompression(chunk.id, level, compressed);
 			}
 		}
 		return chunks;
 	}
 
-	private isFullPassthrough(chunk: Chunk): boolean {
-		if (chunk.confidenceScore > FULL_PASSTHROUGH_THRESHOLD) return true;
-		return chunk.embedding === null;
-	}
-
 	private async llmCompress(chunk: Chunk, level: string): Promise<string> {
 		const prompt = PROMPT_MAP[chunk.chunkType] ?? CODE_PROMPT;
 		try {
@@ -195,11 +227,57 @@ export class Compressor {
 
 	private fallbackCompress(chunk: Chunk, level: string): string {
 		const limit = TRUNCATION_LIMITS[level] ?? 300;
+		// The R1 header is retrieval context, not serving context: the result
+		// formatter already prints `[path:start]`, so paying for it again inside
+		// a 100-char signature would waste most of the budget.
+		const content = stripContextHeader(chunk.content);
 		if (chunk.chunkType === "function" || chunk.chunkType === "class") {
-			return this.extractSignature(chunk.content, limit);
+			return this.extractSignature(content, limit);
+		}
+		if (content.length <= limit) return content;
+		return content.slice(0, limit) + "...";
+	}
+
+	/**
+	 * R3: keep the lines that answer the query instead of the first N lines.
+	 *
+	 * The level's character cap still bounds the output, so serving size never
+	 * grows: matched lines simply win the budget over leading boilerplate, then
+	 * the selection is restored to source order for readability. When nothing
+	 * matches (a vague query), plain truncation is returned — a weak query must
+	 * not blank the result.
+	 */
+	private queryAwareFallback(chunk: Chunk, level: string, query: string): string {
+		const limit = TRUNCATION_LIMITS[level] ?? 300;
+		const terms = queryTerms(query);
+		const body = stripContextHeader(chunk.content);
+		if (terms.length === 0) return this.fallbackCompress(chunk, level);
+
+		const lines = body.split("\n");
+		const matches: { line: string; idx: number; hits: number }[] = [];
+		for (let i = 0; i < lines.length; i++) {
+			const line = lines[i]!;
+			if (line.trim().length === 0) continue;
+			const lower = line.toLowerCase();
+			let hits = 0;
+			for (const t of terms) if (lower.includes(t)) hits++;
+			if (hits > 0) matches.push({ line, idx: i, hits });
+		}
+		if (matches.length === 0) return this.fallbackCompress(chunk, level);
+
+		matches.sort((a, b) => b.hits - a.hits || a.idx - b.idx);
+		let budget = limit;
+		const selected: typeof matches = [];
+		let used = 0;
+		for (const m of matches) {
+			if (used + m.line.length + 1 > budget) continue;
+			selected.push(m);
+			used += m.line.length + 1;
+			if (selected.length >= MAX_QUERY_LINES) break;
 		}
-		if (chunk.content.length <= limit) return chunk.content;
-		return chunk.content.slice(0, limit) + "...";
+		if (selected.length === 0) return this.fallbackCompress(chunk, level);
+		selected.sort((a, b) => a.idx - b.idx);
+		return selected.map((m) => m.line).join("\n");
 	}
 
 	private extractSignature(content: string, limit: number): string {
@@ -220,3 +298,24 @@ export class Compressor {
 		return resultLines.join("\n");
 	}
 }
+
+/** Drop the R1 `// file: …` header line (serving shows `[path:start]` already). */
+function stripContextHeader(content: string): string {
+	if (!content.startsWith("// file: ")) return content;
+	const nl = content.indexOf("\n");
+	return nl < 0 ? content : content.slice(nl + 1);
+}
+
+/** Distinct query terms (>=3 chars, non-boilerplate) for line matching. */
+function queryTerms(query: string): string[] {
+	const tokens = query.toLowerCase().match(/[a-z0-9_]+/g) ?? [];
+	const seen = new Set<string>();
+	const out: string[] = [];
+	for (const t of tokens) {
+		if (t.length < 3 || QUERY_NOISE_TERMS.has(t) || seen.has(t)) continue;
+		seen.add(t);
+		out.push(t);
+		if (out.length >= 40) break;
+	}
+	return out;
+}

package/src/config.ts

@@ -19,7 +19,7 @@ export const DEFAULT_IGNORE: string[] = [
 	// Dependencies (JS, PHP, Python, Ruby, Go, Rust, Java, .NET)
 	"node_modules", "vendor", "bower_components",
 	".pnpm-store", ".pnpm", ".yarn",
-	".venv", "venv", "env", ".env",
+	".venv", "venv", ".env",
 	".tox", ".nox", ".mypy_cache", ".pytest_cache",
 	".ruff_cache", ".cache",
 	"Pods",
@@ -41,13 +41,23 @@ export const DEFAULT_IGNORE: string[] = [
 	// Package locks (huge, not useful)
 	"package-lock.json", "yarn.lock", "pnpm-lock.yaml",
 	"composer.lock", "poetry.lock",
-	// Storage / logs
-	"storage", "logs", "tmp", "temp",
 ];
 
+/**
+ * Directory names that are ambiguous: equally plausible as junk at the project
+ * root (`./storage`, `./logs`, `./tmp`) as they are as real source trees
+ * nested below it (`src/storage/`, `app/logs/`). Matching them at any depth
+ * silently hid whole source trees from the index — including this package's
+ * own `src/storage/*.ts`. They are therefore ignored ONLY at the project root;
+ * nested dirs with these names are indexed.
+ */
+export const ROOT_ONLY_IGNORE: string[] = ["storage", "logs", "tmp", "temp", "env"];
+
 export interface Config {
 	compressionLevel: "minimal" | "standard" | "full";
 	compressionModel: string;
+	/** R3: keep query-matching lines when truncating instead of the head. */
+	compressionQueryAware: boolean;
 	ollamaUrl: string;
 	embeddingModel: string;
 	ollamaEmbedModel: string;
@@ -60,7 +70,29 @@ export interface Config {
 	retrievalConfidenceThreshold: number;
 	retrievalTopK: number;
 	retrievalMarginalRatio: number;
+	/** R6: attach a hit's module parent + adjacent chunk (small-to-big serving). */
+	retrievalExpandContext: boolean;
+	/** R8: composite quality score at/above which a query reports `high`. */
+	retrievalConfidenceHigh: number;
+	/** R8: composite quality score below which a query reports `low` + hints. */
+	retrievalConfidenceLow: number;
+	/** R8: deterministic query elaboration fused as an extra candidate source. */
+	retrievalAutoElaborate: boolean;
+	/** R9: append a `[refine: ...]` hint line to search results. */
+	retrievalRefineHints: boolean;
+	/** R10: opt-in reranking of the fused list. `off` (default) | `local`. */
+	retrievalRerank: "off" | "local";
+	/** R10: HF model id for the local cross-encoder reranker. */
+	retrievalRerankModel: string;
+	/** R10: how many fused candidates the reranker may reorder. */
+	retrievalRerankTopn: number;
+	/** R12: chunking strategy — `legacy` function/class walk (default) or `ast-density` split-merge. */
+	indexerChunking: "legacy" | "ast-density";
+	/** R12: non-whitespace char budget per ast-density chunk. */
+	indexerChunkBudgetChars: number;
 	indexerIgnore: string[];
+	/** Prepend the R1 context header (path/lang/unit/imports) to chunk content. */
+	indexerContextHeaders: boolean;
 	indexerRedactSecrets: boolean;
 	memoryRedactPii: boolean;
 	pricingModel: string;
@@ -73,6 +105,7 @@ export function defaultConfig(): Config {
 	return {
 		compressionLevel: "standard",
 		compressionModel: "phi3:mini",
+		compressionQueryAware: true,
 		ollamaUrl: "http://localhost:11434",
 		embeddingModel: "local",
 		ollamaEmbedModel: "nomic-embed-text",
@@ -80,7 +113,18 @@ export function defaultConfig(): Config {
 		retrievalConfidenceThreshold: 0.2,
 		retrievalTopK: 20,
 		retrievalMarginalRatio: 0.75,
+		retrievalExpandContext: true,
+		retrievalConfidenceHigh: 0.55,
+		retrievalConfidenceLow: 0.45,
+		retrievalAutoElaborate: false,
+		retrievalRefineHints: true,
+		retrievalRerank: "off",
+		retrievalRerankModel: "Xenova/ms-marco-MiniLM-L-6-v2",
+		retrievalRerankTopn: 30,
+		indexerChunking: "legacy",
+		indexerChunkBudgetChars: 1200,
 		indexerIgnore: [...DEFAULT_IGNORE],
+		indexerContextHeaders: true,
 		indexerRedactSecrets: true,
 		memoryRedactPii: true,
 		pricingModel: "opus",
@@ -94,6 +138,7 @@ export function defaultConfig(): Config {
 const MAPPING: Record<string, keyof Config> = {
 	"compression.level": "compressionLevel",
 	"compression.model": "compressionModel",
+	"compression.query_aware": "compressionQueryAware",
 	"compression.ollama_url": "ollamaUrl",
 	"guidance.enabled": "guidanceEnabled",
 	"embedding.model": "embeddingModel",
@@ -101,7 +146,18 @@ const MAPPING: Record<string, keyof Config> = {
 	"retrieval.confidence_threshold": "retrievalConfidenceThreshold",
 	"retrieval.top_k": "retrievalTopK",
 	"retrieval.marginal_ratio": "retrievalMarginalRatio",
+	"retrieval.expand_context": "retrievalExpandContext",
+	"retrieval.confidence_high": "retrievalConfidenceHigh",
+	"retrieval.confidence_low": "retrievalConfidenceLow",
+	"retrieval.auto_elaborate": "retrievalAutoElaborate",
+	"retrieval.refine_hints": "retrievalRefineHints",
+	"retrieval.rerank": "retrievalRerank",
+	"retrieval.rerank_model": "retrievalRerankModel",
+	"retrieval.rerank_topn": "retrievalRerankTopn",
+	"indexer.chunking": "indexerChunking",
+	"indexer.chunk_budget_chars": "indexerChunkBudgetChars",
 	"indexer.ignore": "indexerIgnore",
+	"indexer.context_headers": "indexerContextHeaders",
 	"indexer.redact_secrets": "indexerRedactSecrets",
 	"memory.redact_pii": "memoryRedactPii",
 	"storage.path": "storagePath",