Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
3.7 kB · 122 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123/** * Q&A extraction extension - extracts questions from assistant responses * * Demonstrates the "prompt generator" pattern: * 1. /qna command gets the last assistant message * 2. Shows a spinner while extracting (hides editor) * 3. Loads the result into the editor for user to fill in answers */
import { complete, type UserMessage } from "@earendil-works/pi-ai";import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";import { BorderedLoader } from "@earendil-works/pi-coding-agent";
const SYSTEM_PROMPT = `You are a question extractor. Given text from a conversation, extract any questions that need answering and format them for the user to fill in.
Output format:- List each question on its own line, prefixed with "Q: "- After each question, add a blank line for the answer prefixed with "A: "- If no questions are found, output "No questions found in the last message."
Example output:Q: What is your preferred database?A:
Q: Should we use TypeScript or JavaScript?A:
Keep questions in the order they appeared. Be concise.`;
export default function (pi: ExtensionAPI) { pi.registerCommand("qna", { description: "Extract questions from last assistant message into editor", handler: async (_args, ctx) => { if (!ctx.hasUI) { ctx.ui.notify("qna requires interactive mode", "error"); return; }
if (!ctx.model) { ctx.ui.notify("No model selected", "error"); return; }
// Find the last assistant message on the current branch const branch = ctx.sessionManager.getBranch(); let lastAssistantText: string | undefined;
for (let i = branch.length - 1; i >= 0; i--) { const entry = branch[i]; if (entry.type === "message") { const msg = entry.message; if ("role" in msg && msg.role === "assistant") { if (msg.stopReason !== "stop") { ctx.ui.notify(`Last assistant message incomplete (${msg.stopReason})`, "error"); return; } const textParts = msg.content .filter((c): c is { type: "text"; text: string } => c.type === "text") .map((c) => c.text); if (textParts.length > 0) { lastAssistantText = textParts.join("\n"); break; } } } }
if (!lastAssistantText) { ctx.ui.notify("No assistant messages found", "error"); return; }
// Run extraction with loader UI const result = await ctx.ui.custom<string | null>((tui, theme, _kb, done) => { const loader = new BorderedLoader(tui, theme, `Extracting questions using ${ctx.model!.id}...`); loader.onAbort = () => done(null);
// Do the work const doExtract = async () => { const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model!); if (!auth.ok || !auth.apiKey) { throw new Error(auth.ok ? `No API key for ${ctx.model!.provider}` : auth.error); } const userMessage: UserMessage = { role: "user", content: [{ type: "text", text: lastAssistantText! }], timestamp: Date.now(), };
const response = await complete( ctx.model!, { systemPrompt: SYSTEM_PROMPT, messages: [userMessage] }, { apiKey: auth.apiKey, headers: auth.headers, signal: loader.signal }, );
if (response.stopReason === "aborted") { return null; }
return response.content .filter((c): c is { type: "text"; text: string } => c.type === "text") .map((c) => c.text) .join("\n"); };
doExtract() .then(done) .catch(() => done(null));
return loader; });
if (result === null) { ctx.ui.notify("Cancelled", "info"); return; }
ctx.ui.setEditorText(result); ctx.ui.notify("Questions loaded. Edit and submit when ready.", "info"); }, });}