Skip to main content

context_forge/storage/
searcher.rs

1use std::sync::Arc;
2
3use r2d2::Pool;
4use r2d2_sqlite::SqliteConnectionManager;
5
6use crate::engine::MATCH_ALL_QUERY;
7use crate::entry::ScoredEntry;
8use crate::storage::schema::row_to_entry;
9use crate::traits::Searcher;
10use crate::Result;
11
12/// FTS5-backed full-text search over stored context entries.
13pub struct SqliteSearcher {
14    pool: Arc<Pool<SqliteConnectionManager>>,
15}
16
17impl SqliteSearcher {
18    /// Create a new searcher sharing the given connection pool.
19    #[must_use]
20    pub fn new(pool: Arc<Pool<SqliteConnectionManager>>) -> Self {
21        Self { pool }
22    }
23
24    /// Return all entries (optionally restricted to `scope`) with a fixed
25    /// score of `1.0`, ordered by timestamp descending.
26    ///
27    /// This implements the `MATCH_ALL_QUERY` contract without relying on FTS5
28    /// MATCH, which does not support a bare `*` glob. A fixed score keeps the
29    /// scale consistent with the BM25 path's "higher = better" convention;
30    /// the engine's recency decay (monotonic in age) preserves ordering.
31    fn search_all(&self, scope: Option<&str>, limit: usize) -> Result<Vec<ScoredEntry>> {
32        let conn = self.pool.get()?;
33
34        let mut stmt = conn.prepare(
35            "SELECT id, content, timestamp, kind, scope, session_id, token_count, metadata \
36             FROM entries \
37             WHERE (?1 IS NULL OR scope = ?1) \
38             ORDER BY timestamp DESC \
39             LIMIT ?2",
40        )?;
41
42        let results = stmt
43            .query_map(
44                rusqlite::params![scope, i64::try_from(limit).unwrap_or(i64::MAX)],
45                |row| {
46                    let entry = row_to_entry(row)?;
47                    Ok(ScoredEntry { score: 1.0, entry })
48                },
49            )?
50            .collect::<std::result::Result<Vec<_>, _>>()?;
51
52        Ok(results)
53    }
54}
55
56/// Convert an arbitrary natural-language string into a safe FTS5 `MATCH`
57/// expression.
58///
59/// Splits `query` on any character that is not Unicode-alphanumeric,
60/// drops empty fragments, double-quotes each remaining term (escaping any
61/// internal `"` as `""`), and joins the terms with `" OR "`. Quoting
62/// ensures FTS5 operator characters in the input (`.`, `"`, `*`, `:`, `-`,
63/// etc.) are never interpreted as query syntax. Returns `None` when `query`
64/// contains no alphanumeric terms (e.g. empty or punctuation-only input).
65fn to_fts5_match(query: &str) -> Option<String> {
66    let terms: Vec<String> = query
67        .split(|c: char| !c.is_alphanumeric())
68        .filter(|term| !term.is_empty())
69        .map(|term| format!("\"{}\"", term.replace('"', "\"\"")))
70        .collect();
71
72    if terms.is_empty() {
73        None
74    } else {
75        Some(terms.join(" OR "))
76    }
77}
78
79impl Searcher for SqliteSearcher {
80    fn search(&self, query: &str, scope: Option<&str>, limit: usize) -> Result<Vec<ScoredEntry>> {
81        if query == MATCH_ALL_QUERY {
82            return self.search_all(scope, limit);
83        }
84
85        let Some(match_expr) = to_fts5_match(query) else {
86            return Ok(Vec::new());
87        };
88
89        let conn = self.pool.get()?;
90
91        let mut stmt = conn.prepare(
92            "SELECT e.id, e.content, e.timestamp, e.kind, e.scope, e.session_id, \
93                    e.token_count, e.metadata, bm25(entries_fts) AS score \
94             FROM entries_fts f \
95             JOIN entries e ON e.rowid = f.rowid \
96             WHERE entries_fts MATCH ?1 AND (?2 IS NULL OR e.scope = ?2) \
97             ORDER BY score \
98             LIMIT ?3",
99        )?;
100
101        let results = stmt
102            .query_map(
103                rusqlite::params![match_expr, scope, i64::try_from(limit).unwrap_or(i64::MAX)],
104                |row| {
105                    let entry = row_to_entry(row)?;
106                    let raw_score: f64 = row.get("score")?;
107
108                    Ok(ScoredEntry {
109                        entry,
110                        // bm25() returns negative values; negate so higher = more relevant.
111                        score: -raw_score,
112                    })
113                },
114            )?
115            .collect::<std::result::Result<Vec<_>, _>>()?;
116
117        Ok(results)
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::to_fts5_match;
124
125    #[test]
126    fn single_term() {
127        assert_eq!(to_fts5_match("marco"), Some("\"marco\"".to_owned()));
128    }
129
130    #[test]
131    fn multi_word_with_punctuation() {
132        let input = r#"if I say "marco" you say "polo"."#;
133        let expected = "\"if\" OR \"I\" OR \"say\" OR \"marco\" OR \"you\" OR \"say\" OR \"polo\"";
134        assert_eq!(to_fts5_match(input), Some(expected.to_owned()));
135    }
136
137    #[test]
138    fn punctuation_only_returns_none() {
139        assert_eq!(to_fts5_match("...:-*"), None);
140    }
141
142    #[test]
143    fn empty_returns_none() {
144        assert_eq!(to_fts5_match(""), None);
145    }
146
147    #[test]
148    fn hyphen_and_colon_split_terms() {
149        assert_eq!(
150            to_fts5_match("well-known scope:test"),
151            Some("\"well\" OR \"known\" OR \"scope\" OR \"test\"".to_owned())
152        );
153    }
154
155    #[test]
156    fn unicode_term_survives_as_single_term() {
157        assert_eq!(to_fts5_match("café"), Some("\"café\"".to_owned()));
158    }
159}