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
68        -- Conversations / message threading
69        CREATE TABLE IF NOT EXISTS _adb_conversations (
70            id         TEXT PRIMARY KEY,
71            title      TEXT,
72            metadata   TEXT,
73            created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
74            updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
75        );
76
77        CREATE TABLE IF NOT EXISTS _adb_messages (
78            id              TEXT PRIMARY KEY,
79            conversation_id TEXT NOT NULL REFERENCES _adb_conversations(id) ON DELETE CASCADE,
80            role            TEXT NOT NULL,
81            content         TEXT NOT NULL,
82            metadata        TEXT,
83            created_at      INTEGER NOT NULL DEFAULT (strftime('%s','now'))
84        );
85        CREATE INDEX IF NOT EXISTS idx_messages_conv ON _adb_messages(conversation_id, created_at);
86
87        -- Workflow persistence
88        CREATE TABLE IF NOT EXISTS _adb_workflows (
89            id         TEXT PRIMARY KEY,
90            name       TEXT NOT NULL,
91            status     TEXT NOT NULL DEFAULT 'pending',
92            input      TEXT,
93            output     TEXT,
94            metadata   TEXT,
95            created_at INTEGER NOT NULL DEFAULT (strftime('%s','now')),
96            updated_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
97        );
98
99        CREATE TABLE IF NOT EXISTS _adb_workflow_steps (
100            id           TEXT PRIMARY KEY,
101            workflow_id  TEXT NOT NULL REFERENCES _adb_workflows(id) ON DELETE CASCADE,
102            step_index   INTEGER NOT NULL,
103            name         TEXT NOT NULL,
104            status       TEXT NOT NULL DEFAULT 'pending',
105            input        TEXT,
106            output       TEXT,
107            error        TEXT,
108            started_at   INTEGER,
109            completed_at INTEGER
110        );
111        CREATE INDEX IF NOT EXISTS idx_wf_steps ON _adb_workflow_steps(workflow_id, step_index);
112
113        -- Reasoning traces
114        CREATE TABLE IF NOT EXISTS _adb_traces (
115            id         TEXT PRIMARY KEY,
116            session_id TEXT,
117            parent_id  TEXT REFERENCES _adb_traces(id),
118            trace_type TEXT NOT NULL,
119            content    TEXT NOT NULL,
120            metadata   TEXT,
121            created_at INTEGER NOT NULL DEFAULT (strftime('%s','now'))
122        );
123        CREATE INDEX IF NOT EXISTS idx_traces_session ON _adb_traces(session_id, created_at);
124        CREATE INDEX IF NOT EXISTS idx_traces_parent  ON _adb_traces(parent_id);
125        ",
126    )?;
127
128    conn.execute(
129        "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('schema_version', ?1)",
130        rusqlite::params![SCHEMA_VERSION],
131    )?;
132    conn.execute(
133        "INSERT OR IGNORE INTO _adb_meta (key, value) VALUES ('created_at', ?1)",
134        rusqlite::params![now_ms()],
135    )?;
136    Ok(())
137}
138
139pub fn check_version(conn: &Connection) -> Result<()> {
140    let version: Option<String> = conn
141        .query_row(
142            "SELECT value FROM _adb_meta WHERE key = 'schema_version'",
143            [],
144            |row| row.get(0),
145        )
146        .ok();
147    match version.as_deref() {
148        Some(v) if v == SCHEMA_VERSION => Ok(()),
149        Some(_) => Err(crate::error::AgentDbError::SchemaMigration),
150        None => Ok(()),
151    }
152}
153
154pub fn now_ms() -> i64 {
155    std::time::SystemTime::now()
156        .duration_since(std::time::UNIX_EPOCH)
157        .unwrap_or_default()
158        .as_millis() as i64
159}