Skip to main content

agentdb/
schema.rs

1use crate::error::Result;
2use rusqlite::Connection;
3
4pub const SCHEMA_VERSION: &str = "2";
5
6pub fn bootstrap(conn: &Connection) -> Result<()> {
7    conn.execute_batch(
8        "
9        PRAGMA journal_mode=WAL;
10        PRAGMA foreign_keys=ON;
11        PRAGMA synchronous=NORMAL;
12
13        CREATE TABLE IF NOT EXISTS _adb_meta (
14            key   TEXT PRIMARY KEY,
15            value TEXT NOT NULL
16        );
17
18        CREATE TABLE IF NOT EXISTS _adb_collections (
19            id         TEXT PRIMARY KEY,
20            name       TEXT UNIQUE NOT NULL,
21            dim        INTEGER NOT NULL,
22            metric     TEXT NOT NULL DEFAULT 'cosine',
23            count      INTEGER NOT NULL DEFAULT 0,
24            created_at INTEGER NOT NULL
25        );
26
27        CREATE TABLE IF NOT EXISTS _adb_vectors (
28            id            TEXT NOT NULL,
29            collection_id TEXT NOT NULL,
30            vector        BLOB NOT NULL,
31            metadata      TEXT,
32            created_at    INTEGER NOT NULL,
33            PRIMARY KEY (id, collection_id),
34            FOREIGN KEY (collection_id)
35                REFERENCES _adb_collections(id) ON DELETE CASCADE
36        );
37
38        CREATE TABLE IF NOT EXISTS _adb_hnsw_index (
39            collection_id TEXT PRIMARY KEY,
40            index_blob    BLOB NOT NULL,
41            built_at      INTEGER NOT NULL,
42            is_dirty      INTEGER NOT NULL DEFAULT 0
43        );
44
45        CREATE TABLE IF NOT EXISTS _adb_nodes (
46            id         TEXT PRIMARY KEY,
47            kind       TEXT NOT NULL,
48            data       TEXT,
49            created_at INTEGER NOT NULL,
50            updated_at INTEGER NOT NULL
51        );
52
53        CREATE TABLE IF NOT EXISTS _adb_edges (
54            src        TEXT NOT NULL,
55            dst        TEXT NOT NULL,
56            relation   TEXT NOT NULL,
57            weight     REAL NOT NULL DEFAULT 1.0,
58            created_at INTEGER NOT NULL,
59            PRIMARY KEY (src, dst, relation),
60            FOREIGN KEY (src) REFERENCES _adb_nodes(id) ON DELETE CASCADE,
61            FOREIGN KEY (dst) REFERENCES _adb_nodes(id) ON DELETE CASCADE
62        );
63
64        CREATE INDEX IF NOT EXISTS idx_edges_src  ON _adb_edges(src);
65        CREATE INDEX IF NOT EXISTS idx_edges_dst  ON _adb_edges(dst);
66        CREATE INDEX IF NOT EXISTS idx_vectors_col ON _adb_vectors(collection_id);
67        CREATE INDEX IF NOT EXISTS idx_nodes_kind  ON _adb_nodes(kind);
68
69        -- Conversations / message threading
70        CREATE TABLE IF NOT EXISTS _adb_conversations (
71            id         TEXT PRIMARY KEY,
72            title      TEXT,
73            metadata   TEXT,
74            created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
75            updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
76        );
77
78        CREATE TABLE IF NOT EXISTS _adb_messages (
79            id              TEXT PRIMARY KEY,
80            conversation_id TEXT NOT NULL REFERENCES _adb_conversations(id) ON DELETE CASCADE,
81            role            TEXT NOT NULL,
82            content         TEXT NOT NULL,
83            metadata        TEXT,
84            created_at      INTEGER NOT NULL DEFAULT (strftime('%s','now'))
85        );
86        CREATE INDEX IF NOT EXISTS idx_messages_conv ON _adb_messages(conversation_id, created_at);
87
88        -- Workflow persistence
89        CREATE TABLE IF NOT EXISTS _adb_workflows (
90            id         TEXT PRIMARY KEY,
91            name       TEXT NOT NULL,
92            status     TEXT NOT NULL DEFAULT 'pending',
93            input      TEXT,
94            output     TEXT,
95            metadata   TEXT,
96            created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
97            updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
98        );
99
100        CREATE TABLE IF NOT EXISTS _adb_workflow_steps (
101            id           TEXT PRIMARY KEY,
102            workflow_id  TEXT NOT NULL REFERENCES _adb_workflows(id) ON DELETE CASCADE,
103            step_index   INTEGER NOT NULL,
104            name         TEXT NOT NULL,
105            status       TEXT NOT NULL DEFAULT 'pending',
106            input        TEXT,
107            output       TEXT,
108            error        TEXT,
109            started_at   INTEGER,
110            completed_at INTEGER
111        );
112        CREATE INDEX IF NOT EXISTS idx_wf_steps ON _adb_workflow_steps(workflow_id, step_index);
113
114        -- Reasoning traces
115        CREATE TABLE IF NOT EXISTS _adb_traces (
116            id         TEXT PRIMARY KEY,
117            session_id TEXT,
118            parent_id  TEXT REFERENCES _adb_traces(id),
119            trace_type TEXT NOT NULL,
120            content    TEXT NOT NULL,
121            metadata   TEXT,
122            created_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
123        );
124        CREATE INDEX IF NOT EXISTS idx_traces_session ON _adb_traces(session_id, created_at);
125        CREATE INDEX IF NOT EXISTS idx_traces_parent  ON _adb_traces(parent_id);
126        ",
127    )?;
128
129    conn.execute(
130        "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('schema_version', ?1)",
131        rusqlite::params![SCHEMA_VERSION],
132    )?;
133    conn.execute(
134        "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('created_at', ?1)",
135        rusqlite::params![now_ms()],
136    )?;
137    Ok(())
138}
139
140pub fn check_version(conn: &Connection) -> Result<()> {
141    let version: Option<String> = conn
142        .query_row(
143            "SELECT value FROM _adb_meta WHERE key = 'schema_version'",
144            [],
145            |row| row.get(0),
146        )
147        .ok();
148    match version.as_deref() {
149        Some(v) if v == SCHEMA_VERSION => Ok(()),
150        Some(_) => Err(crate::error::AgentDbError::SchemaMigration),
151        None => Ok(()),
152    }
153}
154
155pub fn now_ms() -> i64 {
156    std::time::SystemTime::now()
157        .duration_since(std::time::UNIX_EPOCH)
158        .unwrap_or_default()
159        .as_millis() as i64
160}