a307df2e23b9ad213bd925c1db9a12d540146d4e

Author
célina <hanouticelina@gmail.com>
Committer
GitHub <noreply@github.com>
Date

Message

Set maxTokens on registered models (#17)

Diff

 1diff --git a/index.ts b/index.ts
 2index 507e3c3dc3c5c38b33fbb19ec9b49630099849bd..42f1ddacdba79fe07ccb2bf253350f2c47742d11 100644
 3--- a/index.ts
 4+++ b/index.ts
 5@@ -15,6 +15,10 @@ const PROVIDER_ID = "llama-cpp";
 6 const DEFAULT_BASE_URL = "http://localhost:8080/v1";
 7 // Fallback for /v1/models entries missing meta.n_ctx.
 8 const DEFAULT_CONTEXT_WINDOW = 8192;
 9+// llama.cpp has no output-token cap (no endpoint reports one; generation is only
10+// bounded by the context window), so use Pi's own default for models that omit
11+// maxTokens (see model-registry.ts parseModels).
12+const DEFAULT_MAX_TOKENS = 16384;
13 const PROPS_TIMEOUT_MS = 120_000;
14 
15 const ModelsResponseSchema = Type.Object({
16@@ -77,14 +81,16 @@ const TEMPLATE_THINKING_LEVEL_MAP = {
17 } satisfies NonNullable<LlamaModel["thinkingLevelMap"]>;
18 
19 // Minimal shape needed to update both registered models and Pi's active model snapshot.
20-type MutableThinkingModel = {
21+type MutableModelMetadata = {
22 	reasoning: boolean;
23 	thinkingLevelMap?: LlamaModel["thinkingLevelMap"];
24 	compat?: LlamaModel["compat"];
25+	contextWindow: number;
26+	maxTokens: number;
27 };
28 
29 // Mark a model as using llama.cpp's chat_template_kwargs.enable_thinking control.
30-function applyTemplateThinkingSupport(model: MutableThinkingModel): void {
31+function applyTemplateThinkingSupport(model: MutableModelMetadata): void {
32 	model.reasoning = true;
33 	model.thinkingLevelMap = TEMPLATE_THINKING_LEVEL_MAP;
34 	model.compat = {
35@@ -162,6 +168,8 @@ export default async function (pi: ExtensionAPI) {
36 				if (isLoaded) {
37 					suffixes.push("(loaded)");
38 				}
39+				const contextWindow =
40+					model.meta?.n_ctx ?? previous?.contextWindow ?? DEFAULT_CONTEXT_WINDOW;
41 				return {
42 					id: model.id,
43 					name: suffixes.length > 0 ? `${model.id} ${suffixes.join(" ")}` : model.id,
44@@ -171,7 +179,8 @@ export default async function (pi: ExtensionAPI) {
45 					thinkingLevelMap: previous?.thinkingLevelMap,
46 					input,
47 					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
48-					contextWindow: model.meta?.n_ctx ?? previous?.contextWindow ?? DEFAULT_CONTEXT_WINDOW,
49+					contextWindow,
50+					maxTokens: Math.min(DEFAULT_MAX_TOKENS, contextWindow),
51 					compat: previous?.compat,
52 				} as LlamaModel;
53 			});
54@@ -209,7 +218,7 @@ export default async function (pi: ExtensionAPI) {
55 		ctx?: ExtensionCtx,
56 		autoload = true,
57 		timeoutMs = PROPS_TIMEOUT_MS,
58-		selectedModel?: MutableThinkingModel,
59+		selectedModel?: MutableModelMetadata,
60 	): Promise<void> {
61 		const model = currentModels.find((m) => m.id === modelId);
62 		if (!model) {
63@@ -217,11 +226,15 @@ export default async function (pi: ExtensionAPI) {
64 		}
65 		if (discoveredMetadata.has(modelId)) {
66 			// Provider re-registration does not update Pi's active model snapshot, so copy
67-			// already-discovered thinking metadata into the selected model when available.
68-			if (selectedModel && model.reasoning) {
69-				selectedModel.reasoning = model.reasoning;
70-				selectedModel.thinkingLevelMap = model.thinkingLevelMap;
71-				selectedModel.compat = model.compat;
72+			// already-discovered metadata into the selected model when available.
73+			if (selectedModel) {
74+				selectedModel.contextWindow = model.contextWindow;
75+				selectedModel.maxTokens = model.maxTokens;
76+				if (model.reasoning) {
77+					selectedModel.reasoning = model.reasoning;
78+					selectedModel.thinkingLevelMap = model.thinkingLevelMap;
79+					selectedModel.compat = model.compat;
80+				}
81 			}
82 			return;
83 		}
84@@ -266,9 +279,14 @@ export default async function (pi: ExtensionAPI) {
85 			let loadedFooterStatus = autoload ? `[llama.cpp] ${modelId} loaded` : undefined;
86 			if (typeof nCtx === "number" && nCtx > 0) {
87 				model.contextWindow = nCtx;
88+				model.maxTokens = Math.min(DEFAULT_MAX_TOKENS, nCtx);
89 				loadedFooterStatus = `[llama.cpp] ${modelId} loaded with ctx ${nCtx} tokens`;
90 				updated = true;
91 			}
92+			if (selectedModel) {
93+				selectedModel.contextWindow = model.contextWindow;
94+				selectedModel.maxTokens = model.maxTokens;
95+			}
96 			if (data.chat_template?.includes("enable_thinking") === true) {
97 				applyTemplateThinkingSupport(model);
98 				if (selectedModel) {