1use std::path::PathBuf;
2
3use rusqlite::Connection;
4use serde::{Deserialize, Serialize};
5
6use crate::error::Error;
7use crate::util;
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct SessionMeta {
15 pub source: String,
16 pub session: String,
17 pub command: String,
18 pub timestamp: i64,
19}
20
21#[derive(Debug)]
22#[allow(dead_code)]
23pub struct SearchResult {
24 pub id: String,
25 pub content: String,
26 pub meta: Option<SessionMeta>,
27 pub similarity: Option<f64>,
28}
29
30pub trait Store {
35 fn index(
36 &mut self,
37 project_id: &str,
38 content: &str,
39 meta: &SessionMeta,
40 ) -> Result<String, Error>;
41
42 fn search(
43 &mut self,
44 project_id: &str,
45 query: &str,
46 limit: usize,
47 ) -> Result<Vec<SearchResult>, Error>;
48
49 fn delete_by_session(&mut self, project_id: &str, session_id: &str) -> Result<usize, Error>;
50
51 fn cleanup_stale(&mut self, project_id: &str, max_age_secs: i64) -> Result<usize, Error>;
52}
53
54pub struct SqliteStore {
59 conn: Connection,
60}
61
62fn db_path() -> PathBuf {
63 dirs::data_dir()
64 .or_else(dirs::home_dir)
65 .unwrap_or_else(|| PathBuf::from("/tmp"))
66 .join(".oo")
67 .join("oo.db")
68}
69
70fn map_err(e: rusqlite::Error) -> Error {
71 Error::Store(e.to_string())
72}
73
74impl SqliteStore {
75 pub fn open() -> Result<Self, Error> {
76 Self::open_at(&db_path())
77 }
78
79 pub fn open_at(path: &std::path::Path) -> Result<Self, Error> {
80 if let Some(parent) = path.parent() {
81 std::fs::create_dir_all(parent).map_err(|e| Error::Store(e.to_string()))?;
82 }
83 let conn = Connection::open(path).map_err(map_err)?;
84 conn.execute_batch(
85 "CREATE TABLE IF NOT EXISTS entries (
86 id TEXT PRIMARY KEY,
87 project TEXT NOT NULL,
88 content TEXT NOT NULL,
89 metadata TEXT,
90 created INTEGER NOT NULL
91 );
92 CREATE VIRTUAL TABLE IF NOT EXISTS entries_fts USING fts5(
93 content,
94 content='entries',
95 content_rowid='rowid'
96 );
97 CREATE TRIGGER IF NOT EXISTS entries_ai AFTER INSERT ON entries BEGIN
98 INSERT INTO entries_fts(rowid, content)
99 VALUES (new.rowid, new.content);
100 END;
101 CREATE TRIGGER IF NOT EXISTS entries_ad AFTER DELETE ON entries BEGIN
102 INSERT INTO entries_fts(entries_fts, rowid, content)
103 VALUES ('delete', old.rowid, old.content);
104 END;
105 CREATE TRIGGER IF NOT EXISTS entries_au AFTER UPDATE ON entries BEGIN
106 INSERT INTO entries_fts(entries_fts, rowid, content)
107 VALUES ('delete', old.rowid, old.content);
108 INSERT INTO entries_fts(rowid, content)
109 VALUES (new.rowid, new.content);
110 END;",
111 )
112 .map_err(map_err)?;
113 Ok(Self { conn })
114 }
115}
116
117impl Store for SqliteStore {
118 fn index(
119 &mut self,
120 project_id: &str,
121 content: &str,
122 meta: &SessionMeta,
123 ) -> Result<String, Error> {
124 let id = uuid::Uuid::new_v4().to_string();
125 let meta_json = serde_json::to_string(meta).map_err(|e| Error::Store(e.to_string()))?;
126 self.conn
127 .execute(
128 "INSERT INTO entries (id, project, content, metadata, created)
129 VALUES (?1, ?2, ?3, ?4, ?5)",
130 rusqlite::params![id, project_id, content, meta_json, meta.timestamp],
131 )
132 .map_err(map_err)?;
133 Ok(id)
134 }
135
136 fn search(
137 &mut self,
138 project_id: &str,
139 query: &str,
140 limit: usize,
141 ) -> Result<Vec<SearchResult>, Error> {
142 let results = if query.len() >= 2 {
144 let mut stmt = self
145 .conn
146 .prepare(
147 "SELECT e.id, e.content, e.metadata, rank
148 FROM entries_fts f
149 JOIN entries e ON e.rowid = f.rowid
150 WHERE entries_fts MATCH ?1 AND e.project = ?2
151 ORDER BY rank
152 LIMIT ?3",
153 )
154 .map_err(map_err)?;
155
156 let fts_query = query
158 .split_whitespace()
159 .map(|w| format!("\"{w}\""))
160 .collect::<Vec<_>>()
161 .join(" ");
162
163 stmt.query_map(rusqlite::params![fts_query, project_id, limit], |row| {
164 let id: String = row.get(0)?;
165 let content: String = row.get(1)?;
166 let meta_json: Option<String> = row.get(2)?;
167 let rank: f64 = row.get(3)?;
168 Ok(SearchResult {
169 id,
170 content,
171 meta: meta_json.as_deref().and_then(parse_meta),
172 similarity: Some(-rank), })
174 })
175 .map_err(map_err)?
176 .filter_map(|r| r.ok())
177 .collect()
178 } else {
179 let mut stmt = self
180 .conn
181 .prepare(
182 "SELECT id, content, metadata
183 FROM entries
184 WHERE project = ?1 AND content LIKE ?2
185 ORDER BY created DESC
186 LIMIT ?3",
187 )
188 .map_err(map_err)?;
189
190 let like = format!("%{query}%");
191 stmt.query_map(rusqlite::params![project_id, like, limit], |row| {
192 let id: String = row.get(0)?;
193 let content: String = row.get(1)?;
194 let meta_json: Option<String> = row.get(2)?;
195 Ok(SearchResult {
196 id,
197 content,
198 meta: meta_json.as_deref().and_then(parse_meta),
199 similarity: None,
200 })
201 })
202 .map_err(map_err)?
203 .filter_map(|r| r.ok())
204 .collect()
205 };
206
207 Ok(results)
208 }
209
210 fn delete_by_session(&mut self, project_id: &str, session_id: &str) -> Result<usize, Error> {
211 let ids: Vec<String> = {
213 let mut stmt = self
214 .conn
215 .prepare("SELECT id, metadata FROM entries WHERE project = ?1")
216 .map_err(map_err)?;
217 stmt.query_map(rusqlite::params![project_id], |row| {
218 let id: String = row.get(0)?;
219 let meta_json: Option<String> = row.get(1)?;
220 Ok((id, meta_json))
221 })
222 .map_err(map_err)?
223 .filter_map(|r| r.ok())
224 .filter(|(_, meta_json)| {
225 meta_json
226 .as_deref()
227 .and_then(parse_meta)
228 .is_some_and(|m| m.source == "oo" && m.session == session_id)
229 })
230 .map(|(id, _)| id)
231 .collect()
232 };
233
234 let count = ids.len();
235 for id in &ids {
236 self.conn
237 .execute("DELETE FROM entries WHERE id = ?1", rusqlite::params![id])
238 .map_err(map_err)?;
239 }
240 Ok(count)
241 }
242
243 fn cleanup_stale(&mut self, project_id: &str, max_age_secs: i64) -> Result<usize, Error> {
244 let now = util::now_epoch();
245 let ids: Vec<String> = {
246 let mut stmt = self
247 .conn
248 .prepare("SELECT id, metadata FROM entries WHERE project = ?1")
249 .map_err(map_err)?;
250 stmt.query_map(rusqlite::params![project_id], |row| {
251 let id: String = row.get(0)?;
252 let meta_json: Option<String> = row.get(1)?;
253 Ok((id, meta_json))
254 })
255 .map_err(map_err)?
256 .filter_map(|r| r.ok())
257 .filter(|(_, meta_json)| {
258 meta_json
259 .as_deref()
260 .and_then(parse_meta)
261 .is_some_and(|m| m.source == "oo" && (now - m.timestamp) > max_age_secs)
262 })
263 .map(|(id, _)| id)
264 .collect()
265 };
266
267 let count = ids.len();
268 for id in &ids {
269 self.conn
270 .execute("DELETE FROM entries WHERE id = ?1", rusqlite::params![id])
271 .map_err(map_err)?;
272 }
273 Ok(count)
274 }
275}
276
277#[cfg(feature = "vipune-store")]
282pub struct VipuneStore {
283 store: vipune::MemoryStore,
284}
285
286#[cfg(feature = "vipune-store")]
287impl VipuneStore {
288 pub fn open() -> Result<Self, Error> {
289 let config = vipune::Config::load().map_err(|e| Error::Store(e.to_string()))?;
290 let store =
291 vipune::MemoryStore::new(&config.database_path, &config.embedding_model, config)
292 .map_err(|e| Error::Store(e.to_string()))?;
293 Ok(Self { store })
294 }
295}
296
297#[cfg(feature = "vipune-store")]
298impl Store for VipuneStore {
299 fn index(
300 &mut self,
301 project_id: &str,
302 content: &str,
303 meta: &SessionMeta,
304 ) -> Result<String, Error> {
305 let meta_json = serde_json::to_string(meta).map_err(|e| Error::Store(e.to_string()))?;
306 match self
307 .store
308 .add_with_conflict(project_id, content, Some(&meta_json), true)
309 {
310 Ok(vipune::AddResult::Added { id }) => Ok(id),
311 Ok(vipune::AddResult::Conflicts { .. }) => Ok(String::new()),
312 Err(e) => Err(Error::Store(e.to_string())),
313 }
314 }
315
316 fn search(
317 &mut self,
318 project_id: &str,
319 query: &str,
320 limit: usize,
321 ) -> Result<Vec<SearchResult>, Error> {
322 let memories = self
323 .store
324 .search_hybrid(project_id, query, limit, 0.3)
325 .map_err(|e| Error::Store(e.to_string()))?;
326 Ok(memories
327 .into_iter()
328 .map(|m| SearchResult {
329 id: m.id,
330 meta: m.metadata.as_deref().and_then(parse_meta),
331 content: m.content,
332 similarity: m.similarity,
333 })
334 .collect())
335 }
336
337 fn delete_by_session(&mut self, project_id: &str, session_id: &str) -> Result<usize, Error> {
338 let entries = self
339 .store
340 .list(project_id, 10_000)
341 .map_err(|e| Error::Store(e.to_string()))?;
342 let mut count = 0;
343 for entry in entries {
344 if let Some(meta) = entry.metadata.as_deref().and_then(parse_meta) {
345 if meta.source == "oo" && meta.session == session_id {
346 self.store
347 .delete(&entry.id)
348 .map_err(|e| Error::Store(e.to_string()))?;
349 count += 1;
350 }
351 }
352 }
353 Ok(count)
354 }
355
356 fn cleanup_stale(&mut self, project_id: &str, max_age_secs: i64) -> Result<usize, Error> {
357 let now = util::now_epoch();
358 let entries = self
359 .store
360 .list(project_id, 10_000)
361 .map_err(|e| Error::Store(e.to_string()))?;
362 let mut count = 0;
363 for entry in entries {
364 if let Some(meta) = entry.metadata.as_deref().and_then(parse_meta) {
365 if meta.source == "oo" && (now - meta.timestamp) > max_age_secs {
366 self.store
367 .delete(&entry.id)
368 .map_err(|e| Error::Store(e.to_string()))?;
369 count += 1;
370 }
371 }
372 }
373 Ok(count)
374 }
375}
376
377fn parse_meta(json: &str) -> Option<SessionMeta> {
382 serde_json::from_str(json).ok()
383}
384
385pub fn open() -> Result<Box<dyn Store>, Error> {
387 #[cfg(feature = "vipune-store")]
388 {
389 return Ok(Box::new(VipuneStore::open()?));
390 }
391 #[cfg(not(feature = "vipune-store"))]
392 {
393 Ok(Box::new(SqliteStore::open()?))
394 }
395}
396
397#[cfg(test)]
402#[path = "store_tests.rs"]
403mod tests;