Created
March 2, 2026 20:49
-
-
Save badlogic/cca48c03ee63efd32c23b677dc3495b0 to your computer and use it in GitHub Desktop.
pi template accelerator extension - auto-fill prompt template args from conversation context
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| /** | |
| * Template Accelerator - use conversation context to auto-fill prompt template args. | |
| * | |
| * Type `$$ /template` (no args) and the extension will: | |
| * 1. Read the template file | |
| * 2. Ask the current LLM to extract arguments from conversation context | |
| * 3. Put the filled `/template "arg1" "arg2" ...` into the editor for review | |
| * | |
| * Press Enter to execute, or edit first. | |
| * | |
| * Usage: | |
| * pi -e ./examples/extensions/template-accelerator.ts | |
| * | |
| * Then in a conversation: | |
| * > We need a math worksheet for grade 8, easy difficulty | |
| * > $$ /arbeitsblatt | |
| * (editor gets: /arbeitsblatt "mathematik" "8" "leicht") | |
| */ | |
| import { readFileSync } from "node:fs"; | |
| import { streamSimple } from "@mariozechner/pi-ai"; | |
| import type { ExtensionAPI, SlashCommandInfo } from "@mariozechner/pi-coding-agent"; | |
| const PREFIX = "$$"; | |
| export default function (pi: ExtensionAPI) { | |
| pi.on("input", async (event, ctx) => { | |
| if (event.source === "extension") return { action: "continue" }; | |
| // Check for $$ prefix | |
| const text = event.text.trim(); | |
| if (!text.startsWith(PREFIX)) return { action: "continue" }; | |
| const afterPrefix = text.slice(PREFIX.length).trim(); | |
| if (!afterPrefix.startsWith("/")) { | |
| ctx.ui.notify(`Usage: ${PREFIX} /template`, "warning"); | |
| return { action: "handled" }; | |
| } | |
| // Parse command name (ignore any args after it, we generate them) | |
| const spaceIdx = afterPrefix.indexOf(" "); | |
| const cmdName = spaceIdx === -1 ? afterPrefix.slice(1) : afterPrefix.slice(1, spaceIdx); | |
| // Find matching prompt template | |
| const commands = pi.getCommands(); | |
| const cmd = commands.find((c: SlashCommandInfo) => c.source === "prompt" && c.name === cmdName); | |
| if (!cmd) { | |
| ctx.ui.notify(`Unknown prompt template: /${cmdName}`, "warning"); | |
| return { action: "handled" }; | |
| } | |
| if (!cmd.path) { | |
| ctx.ui.notify(`No file path for template: /${cmdName}`, "warning"); | |
| return { action: "handled" }; | |
| } | |
| // Read template content | |
| let templateContent: string; | |
| try { | |
| templateContent = readFileSync(cmd.path, "utf-8"); | |
| } catch { | |
| ctx.ui.notify(`Cannot read template: ${cmd.path}`, "error"); | |
| return { action: "handled" }; | |
| } | |
| // Parse placeholders from template | |
| const argSpec = parseArgSpec(templateContent); | |
| if (argSpec.highestPositional === 0 && !argSpec.usesAllArgs && argSpec.slices.length === 0) { | |
| // Template takes no args, just pass through | |
| ctx.ui.setEditorText(`/${cmdName}`); | |
| ctx.ui.notify(`/${cmdName} takes no arguments. Inserted into editor.`, "info"); | |
| return { action: "handled" }; | |
| } | |
| // Build description of what args the template expects | |
| const argDescription = describeArgs(argSpec); | |
| // Get conversation context from session | |
| const entries = ctx.sessionManager.getBranch(); | |
| const contextSnippets: string[] = []; | |
| for (let i = entries.length - 1; i >= 0 && contextSnippets.length < 10; i--) { | |
| const entry = entries[i]; | |
| if (entry.type !== "message") continue; | |
| const msg = entry.message; | |
| if (msg.role === "user" || msg.role === "assistant") { | |
| const text = | |
| typeof msg.content === "string" | |
| ? msg.content | |
| : msg.content | |
| .filter((c: { type: string }) => c.type === "text") | |
| .map((c: { type: string; text?: string }) => c.text ?? "") | |
| .join(""); | |
| if (text.trim()) contextSnippets.unshift(text.slice(0, 500)); | |
| } | |
| } | |
| if (contextSnippets.length === 0) { | |
| ctx.ui.notify("No conversation context to extract args from.", "warning"); | |
| ctx.ui.setEditorText(`/${cmdName} `); | |
| return { action: "handled" }; | |
| } | |
| // Ask LLM to extract args | |
| const model = ctx.model; | |
| if (!model) { | |
| ctx.ui.notify("No model selected.", "error"); | |
| return { action: "handled" }; | |
| } | |
| const apiKey = await ctx.modelRegistry.getApiKeyForProvider(model.provider); | |
| if (!apiKey) { | |
| ctx.ui.notify(`No API key for ${model.provider}.`, "error"); | |
| return { action: "handled" }; | |
| } | |
| ctx.ui.notify(`Extracting args for /${cmdName}...`, "info"); | |
| const extractionPrompt = [ | |
| "Extract arguments for a prompt template from the conversation context below.", | |
| "", | |
| `Template name: /${cmdName}`, | |
| "", | |
| "Template content:", | |
| "```", | |
| templateContent, | |
| "```", | |
| "", | |
| `This template expects: ${argDescription}`, | |
| "", | |
| "Conversation context (most recent last):", | |
| ...contextSnippets.map((s, i) => `--- message ${i + 1} ---\n${s}`), | |
| "", | |
| "Respond with ONLY a JSON array of argument strings, nothing else.", | |
| 'Example: ["math", "8", "easy"]', | |
| "If you cannot determine an argument, use an empty string.", | |
| ].join("\n"); | |
| try { | |
| const stream = streamSimple( | |
| model, | |
| { | |
| systemPrompt: | |
| "You extract structured arguments from conversation context. Respond only with valid JSON.", | |
| messages: [{ role: "user", content: [{ type: "text", text: extractionPrompt }], timestamp: Date.now() }], | |
| }, | |
| { apiKey }, | |
| ); | |
| let response = ""; | |
| for await (const event of stream) { | |
| if (event.type === "text_delta") response += event.delta; | |
| } | |
| // Parse JSON array from response | |
| const jsonMatch = response.match(/\[[\s\S]*\]/); | |
| if (!jsonMatch) { | |
| ctx.ui.notify("LLM did not return a valid JSON array.", "warning"); | |
| ctx.ui.setEditorText(`/${cmdName} `); | |
| return { action: "handled" }; | |
| } | |
| const args: string[] = JSON.parse(jsonMatch[0]); | |
| if (!Array.isArray(args)) { | |
| ctx.ui.notify("LLM returned non-array JSON.", "warning"); | |
| ctx.ui.setEditorText(`/${cmdName} `); | |
| return { action: "handled" }; | |
| } | |
| // Build the command string | |
| const quotedArgs = args.map((a) => `"${String(a).replace(/"/g, '\\"')}"`).join(" "); | |
| const filled = `/${cmdName} ${quotedArgs}`; | |
| ctx.ui.setEditorText(filled); | |
| ctx.ui.notify( | |
| `Suggested: ${filled.slice(0, 80)}${filled.length > 80 ? "..." : ""}. Review and press Enter.`, | |
| "info", | |
| ); | |
| } catch (err) { | |
| const msg = err instanceof Error ? err.message : String(err); | |
| ctx.ui.notify(`Arg extraction failed: ${msg}`, "error"); | |
| ctx.ui.setEditorText(`/${cmdName} `); | |
| } | |
| return { action: "handled" }; | |
| }); | |
| } | |
| // ============================================================================ | |
| // Arg spec parsing - ~20 lines, scans template for placeholder patterns | |
| // ============================================================================ | |
| interface ArgSpec { | |
| highestPositional: number; | |
| positionalIndices: number[]; | |
| usesAllArgs: boolean; | |
| slices: Array<{ start: number; length?: number }>; | |
| } | |
| function parseArgSpec(content: string): ArgSpec { | |
| const positionalSet = new Set<number>(); | |
| let usesAllArgs = false; | |
| const slices: Array<{ start: number; length?: number }> = []; | |
| // $1, $2, ... (but not inside ${...}) | |
| for (const match of content.matchAll(/\$(\d+)/g)) { | |
| const idx = parseInt(match[1], 10); | |
| if (idx > 0) positionalSet.add(idx); | |
| } | |
| // ${@:N} and ${@:N:L} | |
| for (const match of content.matchAll(/\$\{@:(\d+)(?::(\d+))?\}/g)) { | |
| const start = parseInt(match[1], 10); | |
| const length = match[2] ? parseInt(match[2], 10) : undefined; | |
| slices.push({ start, length }); | |
| } | |
| // $@ or $ARGUMENTS | |
| if (/\$@/.test(content) || /\$ARGUMENTS/.test(content)) { | |
| usesAllArgs = true; | |
| } | |
| const positionalIndices = [...positionalSet].sort((a, b) => a - b); | |
| const highestPositional = positionalIndices.length > 0 ? positionalIndices[positionalIndices.length - 1] : 0; | |
| return { highestPositional, positionalIndices, usesAllArgs, slices }; | |
| } | |
| function describeArgs(spec: ArgSpec): string { | |
| const parts: string[] = []; | |
| if (spec.highestPositional > 0) { | |
| parts.push( | |
| `${spec.highestPositional} positional arg(s): ${spec.positionalIndices.map((i) => `$${i}`).join(", ")}`, | |
| ); | |
| } | |
| if (spec.usesAllArgs) { | |
| parts.push("uses all args ($@ or $ARGUMENTS)"); | |
| } | |
| for (const s of spec.slices) { | |
| if (s.length !== undefined) { | |
| parts.push(`slice: ${s.length} arg(s) starting at position ${s.start}`); | |
| } else { | |
| parts.push(`slice: all args from position ${s.start} onwards`); | |
| } | |
| } | |
| return parts.length > 0 ? parts.join("; ") : "unknown arg pattern"; | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment