2515a81663d99c95cd1498fde8474a43e85c1737

Author
TheEdgeOfRage <git@theedgeofrage.com>
Committer
TheEdgeOfRage <git@theedgeofrage.com>
Date

Message

Make model autoloading on select configurable

Diff

  1diff --git a/README.md b/README.md
  2index b3e058812775ab905697ef29c73587272b596454..eebc531869308291c084e9fa8caa4d28fd45cfe3 100644
  3--- a/README.md
  4+++ b/README.md
  5@@ -32,6 +32,24 @@ pi -e ~/code/pi-llama/index.ts
  6 `-e` loads the extension only for the current session, useful while
  7 developing.
  8 
  9+## Configuration
 10+
 11+The extension does not autoload a model when you select it by default. It does
 12+autoload the selected model when you send a message.
 13+
 14+Use `~/.pi/agent/extensions/pi-llama.json` for global settings or
 15+`.pi/pi-llama.json` for project settings:
 16+
 17+```json
 18+{
 19+  "autoloadOnSelect": false
 20+}
 21+```
 22+
 23+Project settings override global settings. Set `autoloadOnSelect` to `true` to
 24+restore loading on model selection. Restart pi or run `/reload` after changing
 25+configuration.
 26+
 27 ## Environment Variables
 28 
 29 This extension supports the following environment variables:
 30diff --git a/index.ts b/index.ts
 31index fd5e90c88f54b57f01e84db4b00fd864ac05311f..63de5dafe5a0409e8377b13ee95cba903c6f7bfb 100644
 32--- a/index.ts
 33+++ b/index.ts
 34@@ -7,7 +7,9 @@
 35  * Usage: `pi install github.com/huggingface/pi-llama`
 36  */
 37 
 38-import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 39+import { CONFIG_DIR_NAME, getAgentDir, type ExtensionAPI } from "@earendil-works/pi-coding-agent";
 40+import { existsSync, readFileSync } from "node:fs";
 41+import { join } from "node:path";
 42 import { Type } from "typebox";
 43 import { Compile } from "typebox/compile";
 44 import { Loader, truncateToWidth, visibleWidth } from "@earendil-works/pi-tui";
 45@@ -21,6 +23,43 @@ const DEFAULT_CONTEXT_WINDOW = 8192;
 46 // maxTokens (see model-registry.ts parseModels).
 47 const DEFAULT_MAX_TOKENS = 16384;
 48 const PROPS_TIMEOUT_MS = 120_000;
 49+const CONFIG_FILE_NAME = "pi-llama.json";
 50+
 51+type LlamaConfig = {
 52+	autoloadOnSelect?: boolean;
 53+};
 54+
 55+const DEFAULT_CONFIG: Required<LlamaConfig> = {
 56+	autoloadOnSelect: false,
 57+};
 58+
 59+function readConfig(path: string): LlamaConfig {
 60+	if (!existsSync(path)) {
 61+		return {};
 62+	}
 63+
 64+	try {
 65+		const value: unknown = JSON.parse(readFileSync(path, "utf8"));
 66+		if (!value || typeof value !== "object" || Array.isArray(value)) {
 67+			throw new Error("expected a JSON object");
 68+		}
 69+
 70+		const autoloadOnSelect = (value as Record<string, unknown>).autoloadOnSelect;
 71+		return typeof autoloadOnSelect === "boolean" ? { autoloadOnSelect } : {};
 72+	} catch (error) {
 73+		console.warn(`[llama-cpp] failed to read config ${path}: ${(error as Error).message}`);
 74+		return {};
 75+	}
 76+}
 77+
 78+function loadConfig(cwd: string, projectIsTrusted: boolean): Required<LlamaConfig> {
 79+	const globalConfig = readConfig(join(getAgentDir(), "extensions", CONFIG_FILE_NAME));
 80+	const projectConfig = projectIsTrusted
 81+		? readConfig(join(cwd, CONFIG_DIR_NAME, CONFIG_FILE_NAME))
 82+		: {};
 83+
 84+	return { ...DEFAULT_CONFIG, ...globalConfig, ...projectConfig };
 85+}
 86 
 87 const ModelsResponseSchema = Type.Object({
 88 	data: Type.Optional(
 89@@ -139,6 +178,7 @@ function isStaleContextError(error: unknown): boolean {
 90 
 91 export default async function (pi: ExtensionAPI) {
 92 	let currentModels: LlamaModel[] = [];
 93+	let autoloadOnSelect = DEFAULT_CONFIG.autoloadOnSelect;
 94 
 95 	pi.registerCommand("llama-version", {
 96 		description: "Get build info of llama.cpp server",
 97@@ -551,6 +591,10 @@ export default async function (pi: ExtensionAPI) {
 98 		}
 99 	}
100 
101+	pi.on("session_start", (_event, ctx) => {
102+		autoloadOnSelect = loadConfig(ctx.cwd, ctx.isProjectTrusted()).autoloadOnSelect;
103+	});
104+
105 	await refreshProvider();
106 
107 	pi.on("input", async (event) => {
108@@ -564,7 +608,7 @@ export default async function (pi: ExtensionAPI) {
109 		if (event.model.provider !== PROVIDER_ID) {
110 			return;
111 		}
112-		void discoverModelMetadata(event.model.id, ctx, true, PROPS_TIMEOUT_MS, event.model);
113+		void discoverModelMetadata(event.model.id, ctx, autoloadOnSelect, PROPS_TIMEOUT_MS, event.model);
114 	});
115 
116 	// Discover /props for already-active models because re-selecting them does not emit model_select.