agentdb/
db.rs

1use crate::audit::AuditStore;
2use crate::context::ContextStore;
3use crate::conversations::ConversationStore;
4use crate::error::Result;
5use crate::fts::FullTextStore;
6use crate::hybrid::{HybridQuery, HybridResult, HybridStore};
7use crate::labels::LabelStore;
8use crate::memory::MemoryGraph;
9use crate::prompts::PromptStore;
10use crate::schema;
11use crate::tools::ToolStore;
12use crate::traces::TraceStore;
13use crate::vectors::VectorStore;
14use crate::workflows::WorkflowStore;
15use rusqlite::Connection;
16use std::sync::atomic::{AtomicBool, Ordering};
17use std::sync::{Arc, Mutex};
18
19/// The main AgentDB connection — your single-file AI database.
20#[derive(Clone)]
21pub struct AgentDB {
22    conn: Arc<Mutex<Connection>>,
23    closed: Arc<AtomicBool>,
24}
25
26impl AgentDB {
27    /// Open or create an AgentDB database. Use `":memory:"` for tests.
28    pub fn open(path: &str) -> Result<Self> {
29        let conn = Connection::open(path)?;
30        schema::bootstrap(&conn)?;
31        schema::check_version(&conn)?;
32        Ok(Self {
33            conn: Arc::new(Mutex::new(conn)),
34            closed: Arc::new(AtomicBool::new(false)),
35        })
36    }
37
38    /// Access the vector store layer
39    pub fn vectors(&self) -> VectorStore {
40        VectorStore::new(Arc::clone(&self.conn))
41    }
42
43    /// Access the memory graph layer
44    pub fn memory(&self) -> MemoryGraph {
45        MemoryGraph::new(Arc::clone(&self.conn))
46    }
47
48    /// Access the full-text search layer
49    pub fn fts(&self) -> FullTextStore {
50        FullTextStore::new(Arc::clone(&self.conn))
51    }
52
53    /// Access the conversation / message-threading layer
54    pub fn conversations(&self) -> ConversationStore {
55        ConversationStore::new(Arc::clone(&self.conn))
56    }
57
58    /// Access the workflow persistence layer
59    pub fn workflows(&self) -> WorkflowStore {
60        WorkflowStore::new(Arc::clone(&self.conn))
61    }
62
63    /// Access the reasoning-trace layer
64    pub fn traces(&self) -> TraceStore {
65        TraceStore::new(Arc::clone(&self.conn))
66    }
67
68    /// Access the tool registry and call log
69    pub fn tools(&self) -> ToolStore {
70        ToolStore::new(Arc::clone(&self.conn))
71    }
72
73    /// Access the immutable audit log
74    pub fn audit(&self) -> AuditStore {
75        AuditStore::new(Arc::clone(&self.conn))
76    }
77
78    /// Access the token-budgeted context window manager
79    pub fn context(&self) -> ContextStore {
80        ContextStore::new(Arc::clone(&self.conn))
81    }
82
83    /// Access the versioned prompt template store
84    pub fn prompts(&self) -> PromptStore {
85        PromptStore::new(Arc::clone(&self.conn))
86    }
87
88    /// Access the data classification / privacy label store
89    pub fn labels(&self) -> LabelStore {
90        LabelStore::new(Arc::clone(&self.conn))
91    }
92
93    /// Run a hybrid graph + vector query
94    pub fn hybrid_query(&self, q: HybridQuery) -> Result<Vec<HybridResult>> {
95        let dim: usize = {
96            let conn = self.conn.lock().unwrap();
97            conn.query_row(
98                "SELECT dim FROM _adb_collections WHERE name = ?1",
99                rusqlite::params![q.collection],
100                |r| r.get::<_, i64>(0).map(|v| v as usize),
101            )
102            .unwrap_or(q.embedding.len())
103        };
104        let col = self.vectors().collection(q.collection, dim)?;
105        let store = HybridStore::new(Arc::clone(&self.conn));
106        store.query(q, &col)
107    }
108
109    /// Execute a raw SQL statement
110    pub fn execute(&self, sql: &str) -> Result<usize> {
111        let conn = self.conn.lock().unwrap();
112        Ok(conn.execute(sql, [])?)
113    }
114
115    /// Execute a parameterized SQL statement
116    pub fn execute_params(&self, sql: &str, params: &[&dyn rusqlite::ToSql]) -> Result<usize> {
117        let conn = self.conn.lock().unwrap();
118        Ok(conn.execute(sql, params)?)
119    }
120
121    /// Run multiple operations atomically inside a single SQLite transaction.
122    ///
123    /// The closure receives a [`rusqlite::Transaction`] and may perform any
124    /// number of reads or writes.  If the closure returns `Ok`, the transaction
125    /// is committed; if it returns `Err` (or panics), the transaction is rolled
126    /// back automatically.
127    ///
128    /// # Example
129    /// ```rust,no_run
130    /// # use agentdb::AgentDB;
131    /// let db = AgentDB::open(":memory:").unwrap();
132    /// db.transaction(|tx| {
133    ///     tx.execute("INSERT INTO _adb_nodes (id, kind, data, created_at, updated_at) VALUES ('x','tag','{}',0,0)", [])?;
134    ///     tx.execute("INSERT INTO _adb_nodes (id, kind, data, created_at, updated_at) VALUES ('y','tag','{}',0,0)", [])?;
135    ///     Ok(())
136    /// }).unwrap();
137    /// ```
138    pub fn transaction<F, T>(&self, f: F) -> Result<T>
139    where
140        F: FnOnce(&rusqlite::Transaction) -> Result<T>,
141    {
142        let mut conn = self.conn.lock().unwrap();
143        let tx = conn.transaction()?;
144        let result = f(&tx)?;
145        tx.commit()?;
146        Ok(result)
147    }
148
149    /// Execute one or more semicolon-separated SQL statements as a single
150    /// atomic batch.  This is a convenience wrapper around
151    /// [`execute_batch`](rusqlite::Connection::execute_batch) that wraps the
152    /// statements in an explicit transaction so partial execution is never
153    /// visible to other threads.
154    ///
155    /// # Example
156    /// ```rust,no_run
157    /// # use agentdb::AgentDB;
158    /// let db = AgentDB::open(":memory:").unwrap();
159    /// db.execute_batch(
160    ///     "INSERT INTO _adb_nodes (id,kind,data,created_at,updated_at) VALUES ('a','t','{}',0,0);
161    ///      INSERT INTO _adb_nodes (id,kind,data,created_at,updated_at) VALUES ('b','t','{}',0,0);"
162    /// ).unwrap();
163    /// ```
164    pub fn execute_batch(&self, sql: &str) -> Result<()> {
165        let mut conn = self.conn.lock().unwrap();
166        let tx = conn.transaction()?;
167        tx.execute_batch(sql)?;
168        tx.commit()?;
169        Ok(())
170    }
171
172    /// Query and return rows as JSON values
173    pub fn query_json(&self, sql: &str) -> Result<Vec<serde_json::Value>> {
174        let conn = self.conn.lock().unwrap();
175        let mut stmt = conn.prepare(sql)?;
176        let col_names: Vec<String> = stmt.column_names().iter().map(|s| s.to_string()).collect();
177        let rows = stmt.query_map([], |row| {
178            let mut map = serde_json::Map::new();
179            for (i, name) in col_names.iter().enumerate() {
180                let val: rusqlite::types::Value = row.get(i)?;
181                map.insert(name.clone(), rusqlite_value_to_json(val));
182            }
183            Ok(serde_json::Value::Object(map))
184        })?;
185        rows.map(|r| r.map_err(crate::error::AgentDbError::Sqlite))
186            .collect()
187    }
188
189    /// Query with parameters and return rows as JSON values.
190    ///
191    /// Use this for any query involving user-supplied values to prevent SQL injection.
192    ///
193    /// # Example
194    /// ```rust,no_run
195    /// # use agentdb::AgentDB;
196    /// let db = AgentDB::open(":memory:").unwrap();
197    /// let rows = db.query_json_params(
198    ///     "SELECT * FROM _adb_nodes WHERE kind = ?1",
199    ///     &[&"session" as &dyn rusqlite::ToSql],
200    /// ).unwrap();
201    /// ```
202    pub fn query_json_params(
203        &self,
204        sql: &str,
205        params: &[&dyn rusqlite::ToSql],
206    ) -> Result<Vec<serde_json::Value>> {
207        let conn = self.conn.lock().unwrap();
208        let mut stmt = conn.prepare(sql)?;
209        let col_names: Vec<String> = stmt.column_names().iter().map(|s| s.to_string()).collect();
210        let rows = stmt.query_map(params, |row| {
211            let mut map = serde_json::Map::new();
212            for (i, name) in col_names.iter().enumerate() {
213                let val: rusqlite::types::Value = row.get(i)?;
214                map.insert(name.clone(), rusqlite_value_to_json(val));
215            }
216            Ok(serde_json::Value::Object(map))
217        })?;
218        rows.map(|r| r.map_err(crate::error::AgentDbError::Sqlite))
219            .collect()
220    }
221
222    /// Flush dirty HNSW indexes and close gracefully.
223    ///
224    /// After this returns, the subsequent `Drop` is a no-op (no double flush).
225    pub fn close(self) -> Result<()> {
226        let collections = self.vectors().list_collections()?;
227        for (name, dim, _) in collections {
228            let col = self.vectors().collection(&name, dim)?;
229            let is_dirty: i64 = {
230                let conn = self.conn.lock().unwrap();
231                conn.query_row(
232                    "SELECT COALESCE(
233                        (SELECT is_dirty FROM _adb_hnsw_index
234                         WHERE collection_id =
235                           (SELECT id FROM _adb_collections WHERE name = ?1)
236                        ), 0)",
237                    rusqlite::params![name],
238                    |r| r.get(0),
239                )
240                .unwrap_or(0)
241            };
242            if is_dirty == 1 {
243                col.reindex()?;
244            }
245        }
246        self.closed.store(true, Ordering::Release);
247        Ok(())
248    }
249
250    /// Return database-wide statistics (single-query implementation).
251    pub fn stats(&self) -> Result<DbStats> {
252        let conn = self.conn.lock().unwrap();
253        conn.query_row(
254            "SELECT
255                 (SELECT COUNT(*)                FROM _adb_collections)     AS collections,
256                 (SELECT COALESCE(SUM(count),0)  FROM _adb_collections)     AS vectors,
257                 (SELECT COUNT(*)                FROM _adb_nodes)            AS nodes,
258                 (SELECT COUNT(*)                FROM _adb_edges)            AS edges,
259                 (SELECT COUNT(*)                FROM _adb_conversations)    AS conversations,
260                 (SELECT COUNT(*)                FROM _adb_messages)         AS messages,
261                 (SELECT COUNT(*)                FROM _adb_workflows)        AS workflows,
262                 (SELECT COUNT(*)                FROM _adb_workflow_steps)   AS workflow_steps,
263                 (SELECT COUNT(*)                FROM _adb_traces)           AS traces,
264                 (SELECT COUNT(*)                FROM _adb_tools)            AS tools,
265                 (SELECT COUNT(*)                FROM _adb_tool_calls)       AS tool_calls,
266                 (SELECT COUNT(*)                FROM _adb_audit_log)        AS audit_entries,
267                 (SELECT COUNT(*)                FROM _adb_prompt_templates) AS prompt_templates",
268            [],
269            |r| {
270                Ok(DbStats {
271                    collections: r.get(0)?,
272                    vectors: r.get(1)?,
273                    nodes: r.get(2)?,
274                    edges: r.get(3)?,
275                    conversations: r.get(4)?,
276                    messages: r.get(5)?,
277                    workflows: r.get(6)?,
278                    workflow_steps: r.get(7)?,
279                    traces: r.get(8)?,
280                    tools: r.get(9)?,
281                    tool_calls: r.get(10)?,
282                    audit_entries: r.get(11)?,
283                    prompt_templates: r.get(12)?,
284                })
285            },
286        )
287        .map_err(crate::error::AgentDbError::Sqlite)
288    }
289}
290
291/// Database-wide statistics returned by [`AgentDB::stats`].
292#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
293pub struct DbStats {
294    /// Number of named vector collections.
295    pub collections: i64,
296    /// Total number of vectors across all collections.
297    pub vectors: i64,
298    /// Number of nodes in the memory graph.
299    pub nodes: i64,
300    /// Number of directed edges in the memory graph.
301    pub edges: i64,
302    /// Number of conversation threads.
303    pub conversations: i64,
304    /// Total number of messages across all conversations.
305    pub messages: i64,
306    /// Number of workflow records.
307    pub workflows: i64,
308    /// Total number of workflow steps across all workflows.
309    pub workflow_steps: i64,
310    /// Total number of reasoning trace entries.
311    pub traces: i64,
312    /// Number of registered tool definitions.
313    pub tools: i64,
314    /// Total number of logged tool calls.
315    pub tool_calls: i64,
316    /// Total number of audit log entries.
317    pub audit_entries: i64,
318    /// Total number of prompt template versions.
319    pub prompt_templates: i64,
320}
321
322impl Drop for AgentDB {
323    fn drop(&mut self) {
324        if self.closed.load(Ordering::Acquire) {
325            return;
326        }
327        if let Ok(collections) = self.vectors().list_collections() {
328            for (name, dim, _) in collections {
329                if let Ok(col) = self.vectors().collection(&name, dim) {
330                    let is_dirty: bool = {
331                        let conn = self.conn.lock().unwrap();
332                        conn.query_row(
333                            "SELECT COALESCE(
334                                (SELECT is_dirty FROM _adb_hnsw_index
335                                 WHERE collection_id =
336                                   (SELECT id FROM _adb_collections WHERE name = ?1)
337                                ), 0)",
338                            rusqlite::params![name],
339                            |r| r.get::<_, i64>(0),
340                        )
341                        .unwrap_or(0)
342                            == 1
343                    };
344                    if is_dirty {
345                        let _ = col.reindex();
346                    }
347                }
348            }
349        }
350    }
351}
352
353fn rusqlite_value_to_json(val: rusqlite::types::Value) -> serde_json::Value {
354    match val {
355        rusqlite::types::Value::Null => serde_json::Value::Null,
356        rusqlite::types::Value::Integer(i) => serde_json::Value::Number(i.into()),
357        rusqlite::types::Value::Real(f) => serde_json::Number::from_f64(f)
358            .map(serde_json::Value::Number)
359            .unwrap_or(serde_json::Value::Null),
360        rusqlite::types::Value::Text(s) => serde_json::Value::String(s),
361        rusqlite::types::Value::Blob(b) => {
362            serde_json::Value::String(format!("<blob {} bytes>", b.len()))
363        }
364    }
365}
agentdb/db.rs

agentdb/
db.rs
