Skip to main content

tuillem_db/
search.rs

1use crate::{Db, DbError};
2
3#[derive(Debug, Clone)]
4pub struct SearchResult {
5    pub session_id: String,
6    pub session_title: String,
7    pub message_id: String,
8    pub content_snippet: String,
9    pub role: String,
10}
11
12impl Db {
13    pub fn search_messages(&self, query: &str) -> Result<Vec<SearchResult>, DbError> {
14        let mut stmt = self.conn.prepare(
15            "SELECT m.session_id, s.title, m.id, snippet(messages_fts, 0, '**', '**', '...', 32), m.role
16             FROM messages_fts
17             JOIN messages m ON m.rowid = messages_fts.rowid
18             JOIN sessions s ON s.id = m.session_id
19             WHERE messages_fts MATCH ?1
20             LIMIT 50",
21        )?;
22
23        let rows = stmt.query_map(rusqlite::params![query], |row| {
24            Ok(SearchResult {
25                session_id: row.get(0)?,
26                session_title: row.get(1)?,
27                message_id: row.get(2)?,
28                content_snippet: row.get(3)?,
29                role: row.get(4)?,
30            })
31        })?;
32
33        let mut results = Vec::new();
34        for row in rows {
35            results.push(row?);
36        }
37
38        Ok(results)
39    }
40
41    pub fn search_sessions_by_tag(&self, tag: &str) -> Result<Vec<String>, DbError> {
42        let mut stmt = self
43            .conn
44            .prepare("SELECT session_id FROM session_tags WHERE tag = ?1")?;
45
46        let rows = stmt.query_map(rusqlite::params![tag], |row| row.get(0))?;
47
48        let mut ids = Vec::new();
49        for row in rows {
50            ids.push(row?);
51        }
52
53        Ok(ids)
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use crate::Db;
60    use crate::messages::{NewBlock, NewMessage};
61
62    fn setup_with_data() -> Db {
63        let db = Db::open_in_memory().unwrap();
64        let session = db.create_session("Rust Chat").unwrap();
65        db.add_session_tag(&session.id, "rust").unwrap();
66
67        let msg = NewMessage {
68            session_id: &session.id,
69            role: "user",
70            content: Some("How do I use iterators in Rust?"),
71            model_id: None,
72            provider_name: None,
73            parent_message_id: None,
74        };
75        db.create_message(&msg, &[]).unwrap();
76
77        let msg2 = NewMessage {
78            session_id: &session.id,
79            role: "assistant",
80            content: Some("Iterators in Rust are lazy and composable."),
81            model_id: None,
82            provider_name: None,
83            parent_message_id: None,
84        };
85        let blocks = vec![NewBlock {
86            block_type: "text",
87            content: "Iterators in Rust are lazy and composable.",
88            sequence: 0,
89        }];
90        db.create_message(&msg2, &blocks).unwrap();
91
92        db
93    }
94
95    #[test]
96    fn test_fts_search() {
97        let db = setup_with_data();
98
99        let results = db.search_messages("iterators").unwrap();
100        assert!(!results.is_empty());
101        assert!(
102            results
103                .iter()
104                .any(|r| r.content_snippet.to_lowercase().contains("iterator"))
105        );
106    }
107
108    #[test]
109    fn test_search_no_results() {
110        let db = setup_with_data();
111
112        let results = db.search_messages("quantum computing").unwrap();
113        assert!(results.is_empty());
114    }
115
116    #[test]
117    fn test_search_by_tag() {
118        let db = setup_with_data();
119
120        let ids = db.search_sessions_by_tag("rust").unwrap();
121        assert_eq!(ids.len(), 1);
122
123        let ids = db.search_sessions_by_tag("python").unwrap();
124        assert!(ids.is_empty());
125    }
126}