pi-context-engine 0.4.0 → 0.5.0

package/README.md

@@ -37,7 +37,7 @@ runs fixed queries. Run it yourself; nothing is cherry-picked.
 
 ## What you get
 
-Ten tools the agent can call, plus three slash commands.
+Ten tools the agent can call, plus four slash commands.
 
 **Search**
 - **`context_search`**: hybrid vector + BM25 retrieval with confidence scores.
@@ -56,7 +56,8 @@ Ten tools the agent can call, plus three slash commands.
 - **`index_status`**: chunk count and measured savings.
 - **`reindex`**: refresh a path or the whole project.
 
-Commands: `/context-index [--full]`, `/context-status`, `/context-savings`.
+Commands: `/context-index [--full]`, `/context-status`, `/context-savings`,
+`/context-warmup`.
 
 **Shipped guidance.** The package also injects a ~900-token guidance section
 into the system prompt while `context_search` is active: how to formulate
@@ -106,6 +107,19 @@ unchanged repo took 237 ms.
   unchanged repo), so this removes the classic cold-start failure of searching
   an empty index, without ever blocking the UI.
 
+### Model warmup
+
+- **`/context-warmup`** downloads the embedding model up front, no project or
+  index needed. With `retrieval.rerank: local` it also fetches the reranker
+  model; with Ollama as the backend it forces the model pull through the same
+  command.
+- **Why it matters**: the embedder initializes lazily, only when a run has
+  chunks to embed. An empty repo or an already up-to-date index never touches
+  the network, so on a fresh machine or an ephemeral container the download
+  lands inside the first search. Warming turns it into an explicit
+  provisioning step (a devcontainer `postCreateCommand`, a CI cache setup),
+  keeping first searches fast.
+
 ## Requirements
 
 - **Node 22+** on macOS, Linux or Windows.
@@ -253,7 +267,7 @@ bootstrapping for non-pi editors.
 
 ```bash
 npm install
-npm run check          # typecheck + 97 unit tests
+npm run check          # typecheck + 251 unit tests
 npm run smoke:env      # tree-sitter + sqlite-vec capability check
 npm run smoke:loader   # load the extension through pi's jiti loader
 npm run smoke:real-world  # index this repo with the local embedder

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-context-engine",
-  "version": "0.4.0",
+  "version": "0.5.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",

package/src/engine.ts

@@ -125,8 +125,19 @@ export interface SessionStats {
 export interface EngineOptions {
 	projectDir: string;
 	config?: Config;
-	/** Test seam: reuse an existing backend/embedder instead of auto-detecting. */
-	runtime?: { backend: LocalBackend; embedder: Embedder };
+	/** Test seam: reuse an existing backend/embedder/reranker instead of auto-detecting. */
+	runtime?: { backend: LocalBackend; embedder: Embedder; reranker?: LocalReranker };
+}
+
+/** Issue #11 (upstream warmup): what /context-warmup resolved + cached. */
+export interface WarmupReport {
+	backend: string;
+	model: string;
+	dimension: number;
+	/** Present only when retrieval.rerank is "local". */
+	reranker?: string;
+	rerankerReady?: boolean;
+	rerankerError?: string;
 }
 
 /** R13: cheap keyword extraction for decision tags (no LLM): distinctive
@@ -150,10 +161,13 @@ export class ContextEngine {
 	readonly projectName: string;
 	readonly storageBase: string;
 	config: Config;
-	private readonly injectedRuntime: { backend: LocalBackend; embedder: Embedder } | undefined;
+	private readonly injectedRuntime:
+		| { backend: LocalBackend; embedder: Embedder; reranker?: LocalReranker }
+		| undefined;
 	private backendInstance: LocalBackend | null = null;
 	private embedderInstance: Embedder | null = null;
 	private retrieverInstance: HybridRetriever | null = null;
+	private rerankerInstance: LocalReranker | null = null;
 	private compressorInstance: Compressor | null = null;
 	private memoryInstance: MemoryDb | null = null;
 	private memoryFailed = false;
@@ -210,15 +224,48 @@ export class ContextEngine {
 			this.retrieverInstance = new HybridRetriever(this.backend, await this.getEmbedder(), {
 				expandContext: this.config.retrievalExpandContext,
 				autoElaborate: this.config.retrievalAutoElaborate,
-				reranker:
-					this.config.retrievalRerank === "local"
-						? new LocalReranker(this.config.retrievalRerankModel)
-						: null,
+				reranker: this.getReranker(),
 			});
 		}
 		return this.retrieverInstance;
 	}
 
+	/**
+	 * The opt-in local reranker, or null when reranking is off. Memoized so
+	 * /context-warmup loads the SAME instance the retriever later searches
+	 * with — the model stays warm across the command boundary.
+	 */
+	private getReranker(): LocalReranker | null {
+		if (this.config.retrievalRerank !== "local") return null;
+		this.rerankerInstance ??= this.injectedRuntime?.reranker ?? new LocalReranker(this.config.retrievalRerankModel);
+		return this.rerankerInstance;
+	}
+
+	/**
+	 * Warm the model caches (issue #11, upstream elara-labs#168): resolve +
+	 * download + dimension-probe the embedding backend, and — when the opt-in
+	 * local reranker is configured — its cross-encoder too. Unlike a reindex
+	 * on an empty or fully up-to-date repo (which never initializes the
+	 * embedder and therefore never touches the network), this guarantees the
+	 * models are cached, so the first search or index doesn't pay the fetch
+	 * inline. Ollama pulls ride the existing ensureModel() path.
+	 */
+	async warmup(): Promise<WarmupReport> {
+		const embedder = await this.getEmbedder();
+		const report: WarmupReport = {
+			backend: embedder.backendName,
+			model: embedder.modelName,
+			dimension: embedder.dimension,
+		};
+		const reranker = this.getReranker();
+		if (reranker) {
+			report.reranker = reranker.model;
+			report.rerankerReady = await reranker.warm();
+			if (!report.rerankerReady) report.rerankerError = reranker.loadError ?? "unknown error";
+		}
+		return report;
+	}
+
 	async getCompressor(): Promise<Compressor> {
 		if (!this.compressorInstance) {
 			this.compressorInstance = new Compressor(

package/src/extension.ts

@@ -461,6 +461,22 @@ export default function contextEngineExtension(pi: ExtensionAPI) {
 		},
 	});
 
+	pi.registerCommand("context-warmup", {
+		description: "Pre-download the embedding (and opt-in reranker) model so the first search doesn't fetch inline",
+		handler: async (_args, ctx) => {
+			ctx.ui.notify("Warming up embedding model…", "info");
+			const report = await getEngine(ctx.cwd).warmup();
+			const parts = [`${report.backend} ${report.model} (${report.dimension} dims)`];
+			if (report.reranker) {
+				parts.push(report.rerankerReady ? `reranker ${report.reranker} ready` : `reranker ${report.reranker} FAILED: ${report.rerankerError}`);
+			}
+			ctx.ui.notify(
+				`✓ Warmup complete: ${parts.join(" · ")}`,
+				report.reranker && !report.rerankerReady ? "warning" : "info",
+			);
+		},
+	});
+
 	/* ── Lifecycle hooks (replace upstream's curl-based hooks) ───────────── */
 
 	pi.on("session_start", async (_event, ctx) => {

package/src/retrieval/reranker.ts

@@ -56,6 +56,8 @@ export class LocalReranker implements RerankerLike {
 	private loading: Promise<LoadedReranker | null> | null = null;
 	/** Set after a load failure; every later call short-circuits to no-op. */
 	failed = false;
+	/** Why the load failed (null before any attempt). Surfaced by /context-warmup. */
+	loadError: string | null = null;
 
 	constructor(
 		readonly model: string,
@@ -88,6 +90,17 @@ export class LocalReranker implements RerankerLike {
 		}
 	}
 
+	/**
+	 * Force the model load now (issue #11, upstream warmup): downloads and
+	 * initializes the cross-encoder without needing a query. Returns false
+	 * when the load failed — the reranker then stays a silent no-op, per its
+	 * degrade-never-break design; `loadError` carries the reason.
+	 */
+	async warm(): Promise<boolean> {
+		if (this.failed) return false;
+		return (await this.ensureLoaded()) !== null;
+	}
+
 	private ensureLoaded(): Promise<LoadedReranker | null> {
 		if (this.loaded) return Promise.resolve(this.loaded);
 		if (!this.loading) {
@@ -102,8 +115,9 @@ export class LocalReranker implements RerankerLike {
 				});
 				this.loaded = { tokenizer, model };
 				return this.loaded;
-			})().catch(() => {
+			})().catch((err: unknown) => {
 				this.failed = true;
+				this.loadError = err instanceof Error ? err.message : String(err);
 				return null;
 			});
 		}