1use std::fmt::{Display, Formatter};
7
8use index_core::{IndexDocument, IndexNode, Link, Redactor};
9use index_extract::{ExtractFormat, extract_document};
10
11pub const PROMPT_TEMPLATE_VERSION: &str = "index-ai-prompt-v1";
13
14const EXPLAIN_SYSTEM_PROMPT: &str = "Explain the Index document in concise terminal-native terms.";
15const SUMMARIZE_SYSTEM_PROMPT: &str =
16 "Summarize the Index document without adding unsupported claims.";
17const EXTRACT_SYSTEM_PROMPT: &str =
18 "Extract structured facts from the Index document as short bullet points.";
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum AiAction {
23 Explain,
25 Summarize,
27 Extract,
29}
30
31impl AiAction {
32 #[must_use]
34 pub fn parse(input: &str) -> Option<Self> {
35 match input.trim().to_ascii_lowercase().as_str() {
36 "explain" => Some(Self::Explain),
37 "summarize" | "summary" => Some(Self::Summarize),
38 "extract" => Some(Self::Extract),
39 _ => None,
40 }
41 }
42
43 #[must_use]
45 pub const fn as_str(&self) -> &'static str {
46 match self {
47 Self::Explain => "explain",
48 Self::Summarize => "summarize",
49 Self::Extract => "extract",
50 }
51 }
52}
53
54impl Display for AiAction {
55 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
56 f.write_str(self.as_str())
57 }
58}
59
60#[derive(Debug, Clone, PartialEq, Eq)]
62pub struct PromptTemplate {
63 pub version: &'static str,
65 pub action: AiAction,
67 pub system: &'static str,
69}
70
71#[must_use]
73pub const fn prompt_template(action: AiAction) -> PromptTemplate {
74 let system = match action {
75 AiAction::Explain => EXPLAIN_SYSTEM_PROMPT,
76 AiAction::Summarize => SUMMARIZE_SYSTEM_PROMPT,
77 AiAction::Extract => EXTRACT_SYSTEM_PROMPT,
78 };
79 PromptTemplate {
80 version: PROMPT_TEMPLATE_VERSION,
81 action,
82 system,
83 }
84}
85
86#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub enum PrivacyMode {
89 Redacted,
91 AllowPageContent,
93}
94
95#[derive(Debug, Clone, PartialEq, Eq)]
97pub struct AiPrompt {
98 pub template_version: String,
100 pub action: AiAction,
102 pub system: String,
104 pub user: String,
106 pub privacy_mode: PrivacyMode,
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct AiRequest {
113 pub prompt: AiPrompt,
115}
116
117#[derive(Debug, Clone, PartialEq, Eq)]
119pub struct AiResponse {
120 pub text: String,
122 pub deterministic_fallback: bool,
124}
125
126#[derive(Debug, Clone, PartialEq, Eq)]
128pub enum AiError {
129 Provider(String),
131 MissingMockResponse,
133}
134
135impl Display for AiError {
136 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
137 match self {
138 Self::Provider(message) => write!(f, "AI provider failed: {message}"),
139 Self::MissingMockResponse => f.write_str("AI mock provider has no response"),
140 }
141 }
142}
143
144impl std::error::Error for AiError {}
145
146pub trait AiProvider {
148 fn transform(&self, request: &AiRequest) -> Result<AiResponse, AiError>;
150}
151
152#[derive(Debug, Clone, Copy, Default)]
154pub struct OfflineProvider;
155
156impl AiProvider for OfflineProvider {
157 fn transform(&self, request: &AiRequest) -> Result<AiResponse, AiError> {
158 Ok(AiResponse {
159 text: deterministic_fallback(&request.prompt),
160 deterministic_fallback: true,
161 })
162 }
163}
164
165#[derive(Debug, Clone, PartialEq, Eq)]
167pub struct MockProvider {
168 response: Option<AiResponse>,
169}
170
171impl MockProvider {
172 #[must_use]
174 pub fn with_response(text: impl Into<String>) -> Self {
175 Self {
176 response: Some(AiResponse {
177 text: text.into(),
178 deterministic_fallback: false,
179 }),
180 }
181 }
182
183 #[must_use]
185 pub const fn empty() -> Self {
186 Self { response: None }
187 }
188}
189
190impl AiProvider for MockProvider {
191 fn transform(&self, _request: &AiRequest) -> Result<AiResponse, AiError> {
192 self.response.clone().ok_or(AiError::MissingMockResponse)
193 }
194}
195
196#[must_use]
198pub fn prepare_ai_request(
199 document: &IndexDocument,
200 action: AiAction,
201 privacy_mode: PrivacyMode,
202 redactor: &Redactor,
203) -> AiRequest {
204 let template = prompt_template(action);
205 let document_text = document_for_prompt(document, action);
206 let user = match privacy_mode {
207 PrivacyMode::Redacted => redactor.redact(&document_text),
208 PrivacyMode::AllowPageContent => document_text,
209 };
210
211 AiRequest {
212 prompt: AiPrompt {
213 template_version: template.version.to_owned(),
214 action,
215 system: template.system.to_owned(),
216 user,
217 privacy_mode,
218 },
219 }
220}
221
222pub fn run_ai_action<P: AiProvider>(
224 provider: &P,
225 document: &IndexDocument,
226 action: AiAction,
227 privacy_mode: PrivacyMode,
228 redactor: &Redactor,
229) -> Result<AiResponse, AiError> {
230 let request = prepare_ai_request(document, action, privacy_mode, redactor);
231 provider.transform(&request)
232}
233
234fn document_for_prompt(document: &IndexDocument, action: AiAction) -> String {
235 match action {
236 AiAction::Explain | AiAction::Summarize => {
237 extract_document(document, ExtractFormat::Markdown)
238 }
239 AiAction::Extract => extract_document(document, ExtractFormat::Json),
240 }
241}
242
243fn deterministic_fallback(prompt: &AiPrompt) -> String {
244 let title = prompt
245 .user
246 .lines()
247 .find_map(|line| line.strip_prefix("# "))
248 .unwrap_or("Untitled document");
249 let facts = prompt
250 .user
251 .lines()
252 .filter(|line| !line.trim().is_empty())
253 .take(4)
254 .collect::<Vec<_>>();
255
256 match prompt.action {
257 AiAction::Explain => format!(
258 "Offline explain: {title}. This document has {} visible prompt lines.",
259 facts.len()
260 ),
261 AiAction::Summarize => {
262 let summary = facts.join(" ");
263 format!("Offline summary: {summary}")
264 }
265 AiAction::Extract => format!(
266 "Offline extract:\n- template: {}\n- prompt_lines: {}",
267 prompt.template_version,
268 facts.len()
269 ),
270 }
271}
272
273#[must_use]
275pub fn document_links(document: &IndexDocument) -> Vec<&Link> {
276 document
277 .nodes
278 .iter()
279 .filter_map(|node| match node {
280 IndexNode::Link(link) => Some(link),
281 _ => None,
282 })
283 .collect()
284}
285
286#[cfg(test)]
287mod tests {
288 use index_core::{IndexDocument, IndexNode, Link, Redactor};
289
290 use super::{
291 AiAction, AiError, MockProvider, OfflineProvider, PrivacyMode, document_links,
292 prepare_ai_request, prompt_template, run_ai_action,
293 };
294
295 fn document() -> IndexDocument {
296 let mut document = IndexDocument::titled("AI Fixture");
297 document.push(IndexNode::Paragraph(
298 "Visible body token=secret Authorization: Bearer abc123".to_owned(),
299 ));
300 document.push(IndexNode::Link(Link::new(
301 "Docs",
302 "https://example.com/docs",
303 )));
304 document
305 }
306
307 #[test]
308 fn prompt_template_snapshot_is_versioned() {
309 let template = prompt_template(AiAction::Summarize);
310 assert_eq!(template.version, "index-ai-prompt-v1");
311 assert_eq!(
312 template.system,
313 "Summarize the Index document without adding unsupported claims."
314 );
315 }
316
317 #[test]
318 fn action_parser_accepts_supported_actions() {
319 assert_eq!(AiAction::parse("explain"), Some(AiAction::Explain));
320 assert_eq!(AiAction::parse("summary"), Some(AiAction::Summarize));
321 assert_eq!(AiAction::parse("extract"), Some(AiAction::Extract));
322 assert_eq!(AiAction::parse("chat"), None);
323 }
324
325 #[test]
326 fn redacted_prompt_hides_known_credential_fields() {
327 let mut redactor = Redactor::new();
328 redactor.add_secret("abc123");
329 let request = prepare_ai_request(
330 &document(),
331 AiAction::Summarize,
332 PrivacyMode::Redacted,
333 &redactor,
334 );
335
336 assert!(request.prompt.user.contains("[REDACTED]"));
337 assert!(!request.prompt.user.contains("abc123"));
338 assert!(!request.prompt.user.contains("token=secret"));
339 }
340
341 #[test]
342 fn explicit_page_content_mode_preserves_document_text() {
343 let redactor = Redactor::new();
344 let request = prepare_ai_request(
345 &document(),
346 AiAction::Summarize,
347 PrivacyMode::AllowPageContent,
348 &redactor,
349 );
350
351 assert!(request.prompt.user.contains("token=secret"));
352 }
353
354 #[test]
355 fn offline_provider_returns_deterministic_fallback() {
356 let redactor = Redactor::new();
357 let response = run_ai_action(
358 &OfflineProvider,
359 &document(),
360 AiAction::Explain,
361 PrivacyMode::Redacted,
362 &redactor,
363 );
364
365 assert!(
366 matches!(response, Ok(response) if response.deterministic_fallback && response.text.contains("Offline explain"))
367 );
368 }
369
370 #[test]
371 fn mock_provider_returns_configured_response() {
372 let provider = MockProvider::with_response("provider text");
373 let redactor = Redactor::new();
374 let response = run_ai_action(
375 &provider,
376 &document(),
377 AiAction::Extract,
378 PrivacyMode::Redacted,
379 &redactor,
380 );
381
382 assert!(matches!(response, Ok(response) if response.text == "provider text"));
383 }
384
385 #[test]
386 fn mock_provider_reports_missing_response() {
387 let redactor = Redactor::new();
388 let response = run_ai_action(
389 &MockProvider::empty(),
390 &document(),
391 AiAction::Extract,
392 PrivacyMode::Redacted,
393 &redactor,
394 );
395
396 assert_eq!(response, Err(AiError::MissingMockResponse));
397 }
398
399 #[test]
400 fn document_link_context_uses_document_model_links() {
401 let document = document();
402 let links = document_links(&document);
403 assert_eq!(links.len(), 1);
404 assert_eq!(links[0].href, "https://example.com/docs");
405 }
406}