Skip to main content

index_ai/
lib.rs

1//! Optional AI-assisted transformation boundary.
2//!
3//! This crate defines provider traits, prompt templates, privacy preparation,
4//! and deterministic local fallbacks. It performs no network IO.
5
6use std::fmt::{Display, Formatter};
7
8use index_core::{IndexDocument, IndexNode, Link, Redactor};
9use index_extract::{ExtractFormat, extract_document};
10
11/// Version identifier for prompt templates shipped in this crate.
12pub 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/// AI-assisted action requested by a user.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum AiAction {
23    /// Explain the document.
24    Explain,
25    /// Summarize the document.
26    Summarize,
27    /// Extract facts from the document.
28    Extract,
29}
30
31impl AiAction {
32    /// Parses an AI action name.
33    #[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    /// Returns the canonical action name.
44    #[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/// Prompt template metadata.
61#[derive(Debug, Clone, PartialEq, Eq)]
62pub struct PromptTemplate {
63    /// Stable template version.
64    pub version: &'static str,
65    /// Action this template is for.
66    pub action: AiAction,
67    /// System instruction text.
68    pub system: &'static str,
69}
70
71/// Returns the prompt template for an AI action.
72#[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/// Privacy mode for prompt preparation.
87#[derive(Debug, Clone, Copy, PartialEq, Eq)]
88pub enum PrivacyMode {
89    /// Redact credential-shaped content before a provider sees it.
90    Redacted,
91    /// Permit page text after the user explicitly invokes an AI action.
92    AllowPageContent,
93}
94
95/// Prepared prompt sent to an AI provider.
96#[derive(Debug, Clone, PartialEq, Eq)]
97pub struct AiPrompt {
98    /// Stable prompt template version.
99    pub template_version: String,
100    /// Requested action.
101    pub action: AiAction,
102    /// System prompt.
103    pub system: String,
104    /// User prompt.
105    pub user: String,
106    /// Privacy mode used while preparing the prompt.
107    pub privacy_mode: PrivacyMode,
108}
109
110/// AI provider request.
111#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct AiRequest {
113    /// Prepared prompt.
114    pub prompt: AiPrompt,
115}
116
117/// AI provider response.
118#[derive(Debug, Clone, PartialEq, Eq)]
119pub struct AiResponse {
120    /// Provider response text.
121    pub text: String,
122    /// Whether the response came from deterministic local fallback logic.
123    pub deterministic_fallback: bool,
124}
125
126/// Errors returned by AI providers.
127#[derive(Debug, Clone, PartialEq, Eq)]
128pub enum AiError {
129    /// Provider failed.
130    Provider(String),
131    /// No mock response was configured.
132    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
146/// AI provider abstraction.
147pub trait AiProvider {
148    /// Transforms a prepared prompt into a response.
149    fn transform(&self, request: &AiRequest) -> Result<AiResponse, AiError>;
150}
151
152/// Deterministic local/offline provider.
153#[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/// Mock provider for tests and adapter integration.
166#[derive(Debug, Clone, PartialEq, Eq)]
167pub struct MockProvider {
168    response: Option<AiResponse>,
169}
170
171impl MockProvider {
172    /// Creates a mock provider with a response.
173    #[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    /// Creates a mock provider with no response.
184    #[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/// Prepares an AI request from a document.
197#[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
222/// Runs an action through a provider.
223pub 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/// Collects links from a document for provider-side context.
274#[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}