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