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}