Skip to main content

oxi/store/
memory_sqlite.rs

1//! SQLite-backed memory store implementing [`MemoryBackend`].
2//!
3//! Provides persistent storage with WAL journal mode for concurrent access
4//! safety. Embedding storage is reserved for a future cosine-search upgrade;
5//! search currently uses simple SQL `LIKE` matching.
6
7use oxi_agent::tools::{MemoryBackend, MemoryItem, ToolError};
8use rusqlite::{Connection, params};
9use std::path::Path;
10use std::pin::Pin;
11use tokio::sync::Mutex;
12
13/// SQLite-backed memory store.
14///
15/// Implements [`MemoryBackend`] for the `memory_*` agent tools.
16/// Each memory is stored with an auto-generated UUID, kind, content, and subject.
17/// The schema reserves an `embedding` column for future cosine-search support.
18#[derive(Debug)]
19pub struct SqliteMemoryStore {
20    db: Mutex<Connection>,
21}
22
23impl SqliteMemoryStore {
24    /// Open or create a SQLite memory store at `path`.
25    ///
26    /// Uses `:memory:` for an in-memory database. For persistent paths, WAL
27    /// journal mode is enabled for better concurrent read performance.
28    pub fn open(path: &Path) -> Result<Self, rusqlite::Error> {
29        let is_memory = path == Path::new(":memory:");
30        let conn = Connection::open(path)?;
31
32        conn.execute_batch("PRAGMA foreign_keys = ON;")?;
33        conn.execute_batch("PRAGMA busy_timeout = 5000;")?;
34
35        if !is_memory {
36            conn.execute_batch("PRAGMA journal_mode = WAL;")?;
37        }
38
39        conn.execute_batch(
40            "CREATE TABLE IF NOT EXISTS memories (
41                id          TEXT PRIMARY KEY,
42                subject     TEXT NOT NULL,
43                kind        TEXT NOT NULL,
44                content     TEXT NOT NULL,
45                embedding   BLOB,
46                created_at  TEXT NOT NULL DEFAULT (datetime('now')),
47                updated_at  TEXT NOT NULL DEFAULT (datetime('now')),
48                metadata    TEXT
49            );",
50        )?;
51
52        Ok(Self {
53            db: Mutex::new(conn),
54        })
55    }
56}
57
58impl MemoryBackend for SqliteMemoryStore {
59    fn put<'a>(
60        &'a self,
61        content: &'a str,
62        kind: &'a str,
63        subject: &'a str,
64    ) -> Pin<Box<dyn Future<Output = Result<String, ToolError>> + Send + 'a>> {
65        Box::pin(async move {
66            let id = uuid::Uuid::new_v4().to_string();
67            let db = self.db.lock().await;
68            db.execute(
69                "INSERT INTO memories (id, subject, kind, content)
70                 VALUES (?1, ?2, ?3, ?4)",
71                params![id, subject, kind, content],
72            )
73            .map_err(|e| format!("Failed to store memory: {e}"))?;
74            Ok(id)
75        })
76    }
77
78    fn search<'a>(
79        &'a self,
80        query: &'a str,
81        k: usize,
82    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
83        Box::pin(async move {
84            let db = self.db.lock().await;
85            let pattern = format!("%{}%", query.replace('%', "\\%").replace('_', "\\_"));
86            let mut stmt = db
87                .prepare(
88                    "SELECT id, kind, content, subject
89                     FROM memories
90                     WHERE content LIKE ?1 ESCAPE '\\'
91                     ORDER BY length(content) ASC
92                     LIMIT ?2",
93                )
94                .map_err(|e| format!("Failed to prepare search: {e}"))?;
95
96            let results: Vec<MemoryItem> = stmt
97                .query_map(params![pattern, k as i64], |row| {
98                    Ok(MemoryItem {
99                        id: row.get(0)?,
100                        kind: row.get(1)?,
101                        content: row.get(2)?,
102                        subject: row.get(3)?,
103                    })
104                })
105                .map_err(|e| format!("Failed to search memories: {e}"))?
106                .filter_map(|r| r.ok())
107                .collect();
108
109            Ok(results)
110        })
111    }
112
113    fn list<'a>(
114        &'a self,
115        subject: &'a str,
116    ) -> Pin<Box<dyn Future<Output = Result<Vec<MemoryItem>, ToolError>> + Send + 'a>> {
117        Box::pin(async move {
118            let db = self.db.lock().await;
119            let mut stmt = db
120                .prepare(
121                    "SELECT id, kind, content, subject
122                     FROM memories
123                     WHERE subject = ?1
124                     ORDER BY updated_at DESC",
125                )
126                .map_err(|e| format!("Failed to prepare list: {e}"))?;
127
128            let results: Vec<MemoryItem> = stmt
129                .query_map(params![subject], |row| {
130                    Ok(MemoryItem {
131                        id: row.get(0)?,
132                        kind: row.get(1)?,
133                        content: row.get(2)?,
134                        subject: row.get(3)?,
135                    })
136                })
137                .map_err(|e| format!("Failed to list memories: {e}"))?
138                .filter_map(|r| r.ok())
139                .collect();
140
141            Ok(results)
142        })
143    }
144
145    fn delete<'a>(
146        &'a self,
147        id: &'a str,
148    ) -> Pin<Box<dyn Future<Output = Result<(), ToolError>> + Send + 'a>> {
149        Box::pin(async move {
150            let db = self.db.lock().await;
151            db.execute("DELETE FROM memories WHERE id = ?1", params![id])
152                .map_err(|e| format!("Failed to delete memory: {e}"))?;
153            Ok(())
154        })
155    }
156}