Skip to main content

rustbrain_core/
fts.rs

1//! FTS5 query helpers: escaping, stopword stripping, and safe `MATCH` construction.
2//!
3//! User input must never be passed raw to SQLite FTS5: operators like `AND`,
4//! `OR`, `NEAR`, and bare `"` can change query structure or cause errors.
5//! Natural-language prompts also contain stopwords (`why`, `not`, `the`) that
6//! would force AND-matches to fail when those tokens are absent from notes.
7
8use crate::error::{BrainError, Result};
9
10/// English stopwords + question fillers stripped for retrieval (not for display).
11const STOPWORDS: &[&str] = &[
12    "a", "about", "above", "after", "again", "against", "all", "also", "an", "and",
13    "are", "as", "at", "be", "been", "before", "being", "below", "between", "both",
14    "but", "by", "can", "could", "describe", "did", "do", "does", "down", "during",
15    "each", "else", "explain", "few", "find", "for", "from", "further", "get", "give",
16    "had", "has", "have", "help", "here", "how", "i", "if", "in", "into", "is", "it",
17    "its", "just", "look", "looking", "make", "may", "me", "might", "more", "most",
18    "must", "my", "need", "needs", "no", "nor", "not", "of", "off", "on", "once",
19    "only", "or", "other", "our", "out", "over", "own", "per", "please", "same",
20    "shall", "should", "show", "so", "some", "such", "summarize", "tell", "than",
21    "that", "the", "their", "them", "then", "there", "these", "they", "this", "those",
22    "through", "to", "too", "under", "up", "use", "used", "using", "versus", "very",
23    "via", "vs", "want", "wants", "was", "we", "were", "what", "when", "where",
24    "which", "who", "whom", "whose", "why", "will", "with", "would", "you", "your",
25];
26
27/// Prepared FTS query plus the significant tokens used for ranking boosts.
28#[derive(Debug, Clone)]
29pub struct PreparedQuery {
30    /// Escaped FTS5 `MATCH` expression (quoted tokens, optional OR).
31    pub fts_match: String,
32    /// Lowercase significant tokens (stopwords removed when possible).
33    pub tokens: Vec<String>,
34    /// True when multiple significant tokens were OR-joined for recall.
35    pub used_or: bool,
36}
37
38/// Tokenize raw user text: split on whitespace/punctuation, lowercase, keep alnum/_/-/:/.
39pub fn tokenize_query(raw: &str) -> Vec<String> {
40    let mut out = Vec::new();
41    let mut cur = String::new();
42    for ch in raw.chars() {
43        if ch.is_alphanumeric() || ch == '_' || ch == '-' || ch == ':' || ch == '.' {
44            cur.push(ch.to_ascii_lowercase());
45        } else if !cur.is_empty() {
46            out.push(std::mem::take(&mut cur));
47        }
48    }
49    if !cur.is_empty() {
50        out.push(cur);
51    }
52    out.retain(|t| !t.is_empty() && t != "-" && t != "_" && t != ":" && t != ".");
53    out
54}
55
56fn is_stopword(t: &str) -> bool {
57    STOPWORDS.binary_search(&t).is_ok()
58}
59
60/// Strip stopwords; if nothing remains, fall back to tokens longer than 1 char.
61pub fn significant_tokens(tokens: &[String]) -> Vec<String> {
62    let mut seen = std::collections::HashSet::new();
63    let mut kept: Vec<String> = tokens
64        .iter()
65        .filter(|t| !is_stopword(t) && t.len() > 1)
66        .filter(|t| seen.insert((*t).clone()))
67        .cloned()
68        .collect();
69    if kept.is_empty() {
70        seen.clear();
71        kept = tokens
72            .iter()
73            .filter(|t| t.len() > 1)
74            .filter(|t| seen.insert((*t).clone()))
75            .cloned()
76            .collect();
77    }
78    if kept.is_empty() {
79        kept = tokens.to_vec();
80    }
81    kept
82}
83
84fn quote_fts_token(t: &str) -> String {
85    let q = t.replace('"', "\"\"");
86    format!("\"{q}\"")
87}
88
89/// Tokens that are real English words but too generic to retrieve alone.
90const GENERIC_TOPIC_TOKENS: &[&str] = &[
91    "architecture", "code", "codebase", "design", "overview", "project", "repo",
92    "repository", "software", "stack", "structure", "system", "tool", "tools",
93];
94
95/// True when every token is a generic overview word (or the list is empty).
96pub fn is_generic_topic(tokens: &[String]) -> bool {
97    if tokens.is_empty() {
98        return true;
99    }
100    tokens.iter().all(|t| {
101        GENERIC_TOPIC_TOKENS.binary_search(&t.as_str()).is_ok()
102    })
103}
104
105/// Build a safe FTS5 MATCH string and ranking tokens from raw user input.
106///
107/// Strategy:
108/// 1. Tokenize (punctuation-safe)
109/// 2. Prefer significant (non-stopword) tokens
110/// 3. Single token → exact quoted match
111/// 4. Multiple tokens → `OR` between them for natural-language recall
112///    (documents rarely contain every filler word from a full question)
113///
114/// # Errors
115///
116/// Returns [`BrainError::FtsQuery`] when no usable token remains.
117pub fn prepare_search_query(raw: &str) -> Result<PreparedQuery> {
118    let all = tokenize_query(raw);
119    if all.is_empty() {
120        return Err(BrainError::FtsQuery(
121            "query must contain at least one alphanumeric token".into(),
122        ));
123    }
124    let tokens = significant_tokens(&all);
125    if tokens.is_empty() {
126        return Err(BrainError::FtsQuery(
127            "query must contain at least one non-stopword token".into(),
128        ));
129    }
130
131    let used_or = tokens.len() > 1;
132    let fts_match = if used_or {
133        tokens
134            .iter()
135            .map(|t| quote_fts_token(t))
136            .collect::<Vec<_>>()
137            .join(" OR ")
138    } else {
139        quote_fts_token(&tokens[0])
140    };
141
142    Ok(PreparedQuery {
143        fts_match,
144        tokens,
145        used_or,
146    })
147}
148
149/// Escape a user query for SQLite FTS5 `MATCH` (legacy helper).
150///
151/// Prefer [`prepare_search_query`] for search/context paths — this keeps
152/// whitespace-split quoting without stopword handling for callers that want
153/// literal AND of every token.
154pub fn escape_fts5_query(raw: &str) -> Result<String> {
155    let tokens: Vec<&str> = raw.split_whitespace().filter(|t| !t.is_empty()).collect();
156    if tokens.is_empty() {
157        return Err(BrainError::FtsQuery(
158            "query must contain at least one non-whitespace token".into(),
159        ));
160    }
161
162    let escaped = tokens
163        .into_iter()
164        .map(|t| {
165            let q = t.replace('"', "\"\"");
166            format!("\"{q}\"")
167        })
168        .collect::<Vec<_>>()
169        .join(" ");
170
171    Ok(escaped)
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    #[test]
179    fn stopwords_sorted() {
180        let mut sorted = STOPWORDS.to_vec();
181        sorted.sort_unstable();
182        assert_eq!(
183            STOPWORDS,
184            sorted.as_slice(),
185            "STOPWORDS must be sorted for binary_search"
186        );
187    }
188
189    #[test]
190    fn escapes_and_quotes_tokens() {
191        let q = escape_fts5_query(r#"raft "consensus""#).unwrap();
192        assert_eq!(q, r#""raft" """consensus""""#);
193    }
194
195    #[test]
196    fn rejects_empty() {
197        assert!(escape_fts5_query("   ").is_err());
198        assert!(prepare_search_query("   ").is_err());
199        assert!(prepare_search_query("!!!").is_err());
200    }
201
202    #[test]
203    fn operators_become_literals() {
204        let q = escape_fts5_query("foo AND bar").unwrap();
205        assert_eq!(q, r#""foo" "AND" "bar""#);
206    }
207
208    #[test]
209    fn prepare_strips_stopwords_and_ors() {
210        let p = prepare_search_query("why egui not tauri").unwrap();
211        assert_eq!(p.tokens, vec!["egui".to_string(), "tauri".to_string()]);
212        assert!(p.used_or);
213        assert_eq!(p.fts_match, r#""egui" OR "tauri""#);
214    }
215
216    #[test]
217    fn prepare_single_significant() {
218        let p = prepare_search_query("explain raft").unwrap();
219        assert_eq!(p.tokens, vec!["raft".to_string()]);
220        assert!(!p.used_or);
221        assert_eq!(p.fts_match, r#""raft""#);
222    }
223
224    #[test]
225    fn tokenize_strips_punctuation() {
226        let t = tokenize_query("egui/tauri + duckdb?");
227        assert!(t.contains(&"egui".to_string()));
228        assert!(t.contains(&"tauri".to_string()));
229        assert!(t.contains(&"duckdb".to_string()));
230    }
231
232    #[test]
233    fn generic_topic_detection() {
234        assert!(is_generic_topic(&["architecture".into(), "overview".into()]));
235        assert!(!is_generic_topic(&["egui".into(), "architecture".into()]));
236        let p = prepare_search_query("summarize architecture").unwrap();
237        // "summarize" is a stopword → only architecture remains → generic
238        assert!(is_generic_topic(&p.tokens), "tokens={:?}", p.tokens);
239    }
240
241    #[test]
242    fn generic_tokens_sorted() {
243        let mut s = GENERIC_TOPIC_TOKENS.to_vec();
244        s.sort_unstable();
245        assert_eq!(GENERIC_TOPIC_TOKENS, s.as_slice());
246    }
247}