oxi/store/
memory_sqlite.rs1use oxi_agent::tools::{MemoryBackend, MemoryItem, ToolError};
8use rusqlite::{Connection, params};
9use std::path::Path;
10use std::pin::Pin;
11use tokio::sync::Mutex;
12
13#[derive(Debug)]
19pub struct SqliteMemoryStore {
20 db: Mutex<Connection>,
21}
22
23impl SqliteMemoryStore {
24 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}