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", "tell", "than", "that", "the",
21    "their", "them", "then", "there", "these", "they", "this", "those", "through",
22    "to", "too", "under", "up", "use", "used", "using", "versus", "very", "via", "vs",
23    "want", "wants", "was", "we", "were", "what", "when", "where", "which", "who",
24    "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/// Build a safe FTS5 MATCH string and ranking tokens from raw user input.
90///
91/// Strategy:
92/// 1. Tokenize (punctuation-safe)
93/// 2. Prefer significant (non-stopword) tokens
94/// 3. Single token → exact quoted match
95/// 4. Multiple tokens → `OR` between them for natural-language recall
96///    (documents rarely contain every filler word from a full question)
97///
98/// # Errors
99///
100/// Returns [`BrainError::FtsQuery`] when no usable token remains.
101pub fn prepare_search_query(raw: &str) -> Result<PreparedQuery> {
102    let all = tokenize_query(raw);
103    if all.is_empty() {
104        return Err(BrainError::FtsQuery(
105            "query must contain at least one alphanumeric token".into(),
106        ));
107    }
108    let tokens = significant_tokens(&all);
109    if tokens.is_empty() {
110        return Err(BrainError::FtsQuery(
111            "query must contain at least one non-stopword token".into(),
112        ));
113    }
114
115    let used_or = tokens.len() > 1;
116    let fts_match = if used_or {
117        tokens
118            .iter()
119            .map(|t| quote_fts_token(t))
120            .collect::<Vec<_>>()
121            .join(" OR ")
122    } else {
123        quote_fts_token(&tokens[0])
124    };
125
126    Ok(PreparedQuery {
127        fts_match,
128        tokens,
129        used_or,
130    })
131}
132
133/// Escape a user query for SQLite FTS5 `MATCH` (legacy helper).
134///
135/// Prefer [`prepare_search_query`] for search/context paths — this keeps
136/// whitespace-split quoting without stopword handling for callers that want
137/// literal AND of every token.
138pub fn escape_fts5_query(raw: &str) -> Result<String> {
139    let tokens: Vec<&str> = raw.split_whitespace().filter(|t| !t.is_empty()).collect();
140    if tokens.is_empty() {
141        return Err(BrainError::FtsQuery(
142            "query must contain at least one non-whitespace token".into(),
143        ));
144    }
145
146    let escaped = tokens
147        .into_iter()
148        .map(|t| {
149            let q = t.replace('"', "\"\"");
150            format!("\"{q}\"")
151        })
152        .collect::<Vec<_>>()
153        .join(" ");
154
155    Ok(escaped)
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161
162    #[test]
163    fn stopwords_sorted() {
164        let mut sorted = STOPWORDS.to_vec();
165        sorted.sort_unstable();
166        assert_eq!(
167            STOPWORDS,
168            sorted.as_slice(),
169            "STOPWORDS must be sorted for binary_search"
170        );
171    }
172
173    #[test]
174    fn escapes_and_quotes_tokens() {
175        let q = escape_fts5_query(r#"raft "consensus""#).unwrap();
176        assert_eq!(q, r#""raft" """consensus""""#);
177    }
178
179    #[test]
180    fn rejects_empty() {
181        assert!(escape_fts5_query("   ").is_err());
182        assert!(prepare_search_query("   ").is_err());
183        assert!(prepare_search_query("!!!").is_err());
184    }
185
186    #[test]
187    fn operators_become_literals() {
188        let q = escape_fts5_query("foo AND bar").unwrap();
189        assert_eq!(q, r#""foo" "AND" "bar""#);
190    }
191
192    #[test]
193    fn prepare_strips_stopwords_and_ors() {
194        let p = prepare_search_query("why egui not tauri").unwrap();
195        assert_eq!(p.tokens, vec!["egui".to_string(), "tauri".to_string()]);
196        assert!(p.used_or);
197        assert_eq!(p.fts_match, r#""egui" OR "tauri""#);
198    }
199
200    #[test]
201    fn prepare_single_significant() {
202        let p = prepare_search_query("explain raft").unwrap();
203        assert_eq!(p.tokens, vec!["raft".to_string()]);
204        assert!(!p.used_or);
205        assert_eq!(p.fts_match, r#""raft""#);
206    }
207
208    #[test]
209    fn tokenize_strips_punctuation() {
210        let t = tokenize_query("egui/tauri + duckdb?");
211        assert!(t.contains(&"egui".to_string()));
212        assert!(t.contains(&"tauri".to_string()));
213        assert!(t.contains(&"duckdb".to_string()));
214    }
215}