1// SPDX-FileCopyrightText: Amolith <amolith@secluded.site>
2// SPDX-FileCopyrightText: Armin Ronacher <armin.ronacher@active-4.com>
3//
4// SPDX-License-Identifier: Apache-2.0
5
6/**
7 * Q&A extraction hook - extracts questions from assistant responses
8 *
9 * Custom interactive TUI for answering questions.
10 *
11 * Demonstrates the "prompt generator" pattern with custom TUI:
12 * 1. /answer command gets the last assistant message
13 * 2. Shows a spinner while extracting questions as structured JSON
14 * 3. Presents an interactive TUI to navigate and answer questions
15 * 4. Submits the compiled answers when done
16 */
17
18import { complete, type UserMessage } from "@mariozechner/pi-ai";
19import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
20import { BorderedLoader } from "@mariozechner/pi-coding-agent";
21import { parseToolCallResult, resolveExtractionModel } from "./extract.js";
22import { type ExtractionResult, QUESTION_EXTRACTION_TOOL, SYSTEM_PROMPT } from "./prompt.js";
23import { QnAComponent } from "./QnAComponent.js";
24
25export default function (pi: ExtensionAPI) {
26 const answerHandler = async (ctx: ExtensionContext, instruction?: string) => {
27 if (!ctx.hasUI) {
28 ctx.ui.notify("answer requires interactive mode", "error");
29 return;
30 }
31
32 if (!ctx.model) {
33 ctx.ui.notify("No model selected", "error");
34 return;
35 }
36
37 // Find the last assistant message on the current branch
38 const branch = ctx.sessionManager.getBranch();
39 let lastAssistantText: string | undefined;
40
41 for (let i = branch.length - 1; i >= 0; i--) {
42 const entry = branch[i];
43 if (entry.type === "message") {
44 const msg = entry.message;
45 if ("role" in msg && msg.role === "assistant") {
46 if (msg.stopReason !== "stop") {
47 ctx.ui.notify(`Last assistant message incomplete (${msg.stopReason})`, "error");
48 return;
49 }
50 const textParts = msg.content
51 .filter((c): c is { type: "text"; text: string } => c.type === "text")
52 .map((c) => c.text);
53 if (textParts.length > 0) {
54 lastAssistantText = textParts.join("\n");
55 break;
56 }
57 }
58 }
59 }
60
61 if (!lastAssistantText) {
62 ctx.ui.notify("No assistant messages found", "error");
63 return;
64 }
65
66 // Capture narrowed values so closures can use them without non-null assertions.
67 const sessionModel = ctx.model;
68 const assistantText = lastAssistantText;
69
70 // Run extraction with loader UI.
71 // The result distinguishes user cancellation from errors so we can
72 // show the right message after the custom UI closes.
73 type ExtractionOutcome =
74 | { kind: "ok"; result: ExtractionResult }
75 | { kind: "cancelled" }
76 | { kind: "error"; message: string };
77
78 const outcome = await ctx.ui.custom<ExtractionOutcome>((tui, theme, _kb, done) => {
79 const extractionModel = resolveExtractionModel(ctx, sessionModel);
80 const loader = new BorderedLoader(tui, theme, `Extracting questions using ${extractionModel.name}...`);
81
82 // Guard against double-completion: loader.onAbort fires on user
83 // cancel, but the in-flight promise may also resolve/reject after
84 // the abort. Only the first call to finish() takes effect.
85 let finished = false;
86 const finish = (result: ExtractionOutcome) => {
87 if (finished) return;
88 finished = true;
89 done(result);
90 };
91
92 loader.onAbort = () => finish({ kind: "cancelled" });
93
94 const tryExtract = async (model: ReturnType<typeof resolveExtractionModel>) => {
95 const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
96 if (!auth.ok) {
97 return { kind: "model_error" as const, model };
98 }
99 const instructionBlock = instruction
100 ? `<user_instruction>\n${instruction}\n</user_instruction>\n\nExtract questions from the assistant's message based on the user's instruction above.`
101 : "Extract questions from the assistant's message for the user to fill out.";
102
103 const userMessage: UserMessage = {
104 role: "user",
105 content: [
106 {
107 type: "text",
108 text: `<last_assistant_message>\n${assistantText}\n</last_assistant_message>\n\n${instructionBlock}`,
109 },
110 ],
111 timestamp: Date.now(),
112 };
113
114 const response = await complete(
115 model,
116 {
117 systemPrompt: SYSTEM_PROMPT,
118 messages: [userMessage],
119 tools: [QUESTION_EXTRACTION_TOOL],
120 },
121 { apiKey: auth.apiKey, headers: auth.headers, signal: loader.signal },
122 );
123
124 if (response.stopReason === "aborted") {
125 return { kind: "cancelled" as const };
126 }
127
128 if (response.stopReason === "error" || response.content.length === 0) {
129 return { kind: "model_error" as const, model };
130 }
131
132 const parsed = parseToolCallResult(response);
133 if (!parsed) {
134 return {
135 kind: "parse_error" as const,
136 message: `${model.name} did not call extract_questions tool`,
137 };
138 }
139 return { kind: "ok" as const, result: parsed };
140 };
141
142 const doExtract = async () => {
143 let result = await tryExtract(extractionModel);
144
145 // If the preferred model errored and it's not already the
146 // session model, fall back to the session model.
147 if (result.kind === "model_error" && extractionModel.id !== sessionModel.id) {
148 ctx.ui.notify(`${extractionModel.name} unavailable, falling back to ${sessionModel.name}...`, "warning");
149 result = await tryExtract(sessionModel);
150 }
151
152 switch (result.kind) {
153 case "ok":
154 return finish({ kind: "ok", result: result.result });
155 case "cancelled":
156 return finish({ kind: "cancelled" });
157 case "model_error":
158 return finish({
159 kind: "error",
160 message: `${result.model.name} returned an error with no content`,
161 });
162 case "parse_error":
163 return finish({ kind: "error", message: result.message });
164 }
165 };
166
167 doExtract().catch((err) =>
168 finish({
169 kind: "error",
170 message: `${err?.message ?? err}`,
171 }),
172 );
173
174 return loader;
175 });
176
177 if (outcome.kind === "cancelled") {
178 ctx.ui.notify("Cancelled", "info");
179 return;
180 }
181 if (outcome.kind === "error") {
182 ctx.ui.notify(`Extraction failed: ${outcome.message}`, "error");
183 return;
184 }
185
186 const extractionResult = outcome.result;
187
188 if (extractionResult.questions.length === 0) {
189 ctx.ui.notify("No questions found in the last message", "info");
190 return;
191 }
192
193 // Show the Q&A component
194 const answersResult = await ctx.ui.custom<string | null>((tui, _theme, _kb, done) => {
195 return new QnAComponent(extractionResult.questions, tui, done);
196 });
197
198 if (answersResult === null) {
199 ctx.ui.notify("Cancelled", "info");
200 return;
201 }
202
203 // Send the answers directly as a message and trigger a turn
204 pi.sendMessage(
205 {
206 customType: "answers",
207 content: `I answered your questions in the following way:\n\n${answersResult}`,
208 display: true,
209 },
210 { triggerTurn: true },
211 );
212 };
213
214 pi.registerCommand("answer", {
215 description:
216 "Extract questions from last assistant message into interactive Q&A. Optional: provide instructions for what to extract.",
217 handler: (args, ctx) => answerHandler(ctx, args?.trim() || undefined),
218 });
219
220 pi.registerShortcut("ctrl+.", {
221 description: "Extract and answer questions",
222 handler: answerHandler,
223 });
224}