Skip to main content

sparrow/memory/
fts.rs

1//! Full-Text Search (FTS5) for Sparrow sessions.
2//!
3//! Provides fast full-text search across all conversation sessions.
4//! Uses SQLite FTS5 for indexing and querying.
5
6use std::path::PathBuf;
7
8/// A search hit in the session database.
9#[derive(Debug, Clone)]
10pub struct SessionHit {
11    pub session_id: String,
12    pub title: Option<String>,
13    pub snippet: String,
14    pub started_at: String,
15    pub match_count: usize,
16}
17
18/// Session search engine backed by SQLite FTS5.
19pub struct SessionSearch {
20    db_path: PathBuf,
21}
22
23impl SessionSearch {
24    /// Open or create the session search database.
25    pub fn open(db_path: PathBuf) -> anyhow::Result<Self> {
26        let conn = rusqlite::Connection::open(&db_path)?;
27
28        // Enable FTS5
29        conn.execute_batch(
30            "
31            CREATE VIRTUAL TABLE IF NOT EXISTS sessions_fts USING fts5(
32                session_id,
33                title,
34                content,
35                tokenize='unicode61'
36            );
37        ",
38        )?;
39
40        Ok(Self { db_path })
41    }
42
43    /// Index a session's content for full-text search.
44    pub fn index_session(
45        &self,
46        session_id: &str,
47        title: Option<&str>,
48        content: &str,
49    ) -> anyhow::Result<()> {
50        let conn = rusqlite::Connection::open(&self.db_path)?;
51
52        // Delete existing entry for this session
53        conn.execute(
54            "DELETE FROM sessions_fts WHERE session_id = ?1",
55            rusqlite::params![session_id],
56        )?;
57
58        // Insert new entry
59        conn.execute(
60            "INSERT INTO sessions_fts (session_id, title, content) VALUES (?1, ?2, ?3)",
61            rusqlite::params![session_id, title.unwrap_or(""), content],
62        )?;
63
64        Ok(())
65    }
66
67    /// Search sessions by query.
68    pub fn search(&self, query: &str, limit: usize) -> anyhow::Result<Vec<SessionHit>> {
69        let conn = rusqlite::Connection::open(&self.db_path)?;
70
71        // Use FTS5 with snippet for highlighting
72        let mut stmt = conn.prepare(
73            "SELECT session_id, title, snippet(sessions_fts, 2, '<b>', '</b>', '...', 40) as snippet,
74                    sessions_fts.rank as rank
75             FROM sessions_fts
76             WHERE sessions_fts MATCH ?1
77             ORDER BY rank
78             LIMIT ?2"
79        )?;
80
81        let hits = stmt
82            .query_map(rusqlite::params![query, limit as i64], |row| {
83                Ok(SessionHit {
84                    session_id: row.get(0)?,
85                    title: row.get(1)?,
86                    snippet: row.get::<_, String>(2).unwrap_or_default(),
87                    started_at: String::new(),
88                    match_count: 1,
89                })
90            })?
91            .filter_map(|r| r.ok())
92            .collect();
93
94        Ok(hits)
95    }
96
97    /// Get recent sessions (no search query).
98    pub fn recent(&self, limit: usize) -> anyhow::Result<Vec<SessionHit>> {
99        let conn = rusqlite::Connection::open(&self.db_path)?;
100
101        let mut stmt = conn.prepare(
102            "SELECT session_id, title, substr(content, 1, 200) as snippet
103             FROM sessions_fts
104             ORDER BY rowid DESC
105             LIMIT ?1",
106        )?;
107
108        let hits = stmt
109            .query_map(rusqlite::params![limit as i64], |row| {
110                Ok(SessionHit {
111                    session_id: row.get(0)?,
112                    title: row.get(1)?,
113                    snippet: row.get::<_, String>(2).unwrap_or_default(),
114                    started_at: String::new(),
115                    match_count: 1,
116                })
117            })?
118            .filter_map(|r| r.ok())
119            .collect();
120
121        Ok(hits)
122    }
123
124    /// Delete a session from the index.
125    pub fn remove_session(&self, session_id: &str) -> anyhow::Result<bool> {
126        let conn = rusqlite::Connection::open(&self.db_path)?;
127        let count = conn.execute(
128            "DELETE FROM sessions_fts WHERE session_id = ?1",
129            rusqlite::params![session_id],
130        )?;
131        Ok(count > 0)
132    }
133
134    /// Get the total number of indexed sessions.
135    pub fn count(&self) -> anyhow::Result<usize> {
136        let conn = rusqlite::Connection::open(&self.db_path)?;
137        let count: usize =
138            conn.query_row("SELECT COUNT(*) FROM sessions_fts", [], |row| row.get(0))?;
139        Ok(count)
140    }
141
142    /// Rebuild the FTS index (useful after corruption or migration).
143    pub fn rebuild(&self) -> anyhow::Result<()> {
144        let conn = rusqlite::Connection::open(&self.db_path)?;
145        conn.execute_batch(
146            "
147            DELETE FROM sessions_fts;
148            INSERT INTO sessions_fts(sessions_fts) VALUES('rebuild');
149        ",
150        )?;
151        Ok(())
152    }
153}
154
155#[cfg(test)]
156mod tests {
157    use super::*;
158
159    #[test]
160    fn test_index_and_search() {
161        let tmp = std::env::temp_dir().join("sparrow_test_fts.db");
162        let search = SessionSearch::open(tmp.clone()).unwrap();
163
164        search
165            .index_session(
166                "sess1",
167                Some("Test Session"),
168                "This is about Rust programming",
169            )
170            .unwrap();
171        search
172            .index_session("sess2", Some("Another"), "Python is great for data science")
173            .unwrap();
174
175        let hits = search.search("Rust programming", 5).unwrap();
176        assert!(!hits.is_empty());
177
178        let _ = std::fs::remove_file(&tmp);
179    }
180}