agentdb/vectors/
collection.rs

1use crate::error::{AgentDbError, Result};
2use crate::filter;
3use crate::schema::now_ms;
4use crate::vectors::hnsw::{DistanceMetric, HnswIndex};
5use rusqlite::{params, Connection};
6use serde_json::Value;
7use std::sync::{Arc, Mutex};
8use uuid::Uuid;
9
10/// A single vector entry
11#[derive(Debug, Clone)]
12pub struct VectorEntry {
13    pub id: String,
14    pub vector: Vec<f32>,
15    pub metadata: Option<Value>,
16}
17
18/// A single vector search result
19#[derive(Debug, Clone)]
20pub struct SearchResult {
21    pub id: String,
22    pub score: f32,
23    pub metadata: Option<Value>,
24}
25
26/// Options controlling a vector search
27#[derive(Debug, Clone)]
28pub struct SearchOptions {
29    pub top_k: usize,
30    pub metric: DistanceMetric,
31    pub filter: Option<Value>,
32}
33
34impl Default for SearchOptions {
35    fn default() -> Self {
36        Self {
37            top_k: 10,
38            metric: DistanceMetric::Cosine,
39            filter: None,
40        }
41    }
42}
43
44/// An entry for batch upsert
45#[derive(Debug, Clone)]
46pub struct BatchEntry {
47    pub id: String,
48    pub vector: Vec<f32>,
49    pub metadata: Option<Value>,
50}
51
52/// A named vector collection
53pub struct Collection {
54    pub id: String,
55    pub name: String,
56    pub dim: usize,
57    pub(crate) conn: Arc<Mutex<Connection>>,
58    index: Mutex<Option<HnswIndex>>,
59    metric: DistanceMetric,
60}
61
62impl Collection {
63    pub(crate) fn new(
64        id: String,
65        name: String,
66        dim: usize,
67        metric: DistanceMetric,
68        conn: Arc<Mutex<Connection>>,
69    ) -> Self {
70        Self {
71            id,
72            name,
73            dim,
74            conn,
75            index: Mutex::new(None),
76            metric,
77        }
78    }
79
80    /// Insert or update a single vector
81    pub fn upsert(&self, entry: VectorEntry) -> Result<()> {
82        if entry.vector.len() != self.dim {
83            return Err(AgentDbError::DimensionMismatch {
84                expected: self.dim,
85                got: entry.vector.len(),
86            });
87        }
88        let blob: Vec<u8> = entry.vector.iter().flat_map(|f| f.to_le_bytes()).collect();
89        let meta = entry.metadata.as_ref().map(|m| m.to_string());
90        let conn = self.conn.lock().unwrap();
91        // INSERT OR IGNORE returns changes()=1 for a new row, 0 for a duplicate.
92        let inserted = conn.execute(
93            "INSERT OR IGNORE INTO _adb_vectors (id, collection_id, vector, metadata, created_at)
94             VALUES (?1, ?2, ?3, ?4, ?5)",
95            params![entry.id, self.id, blob, meta, now_ms()],
96        )?;
97        if inserted == 0 {
98            conn.execute(
99                "UPDATE _adb_vectors SET vector = ?1, metadata = ?2
100                 WHERE id = ?3 AND collection_id = ?4",
101                params![blob, meta, entry.id, self.id],
102            )?;
103        }
104        conn.execute(
105            "INSERT INTO _adb_hnsw_index (collection_id, index_blob, built_at, is_dirty)
106             VALUES (?1, X'', ?2, 1)
107             ON CONFLICT(collection_id) DO UPDATE SET is_dirty = 1",
108            params![self.id, now_ms()],
109        )?;
110        if inserted > 0 {
111            conn.execute(
112                "UPDATE _adb_collections SET count = count + 1 WHERE id = ?1",
113                params![self.id],
114            )?;
115        }
116        *self.index.lock().unwrap() = None;
117        Ok(())
118    }
119
120    /// Insert or update multiple vectors in a single transaction
121    pub fn upsert_batch(&self, entries: Vec<BatchEntry>) -> Result<usize> {
122        if entries.is_empty() {
123            return Ok(0);
124        }
125        for e in &entries {
126            if e.vector.len() != self.dim {
127                return Err(AgentDbError::DimensionMismatch {
128                    expected: self.dim,
129                    got: e.vector.len(),
130                });
131            }
132        }
133        let conn = self.conn.lock().unwrap();
134        conn.execute_batch("BEGIN")?;
135        let result: Result<usize> = (|| {
136            let mut new_rows: usize = 0;
137            for e in &entries {
138                let blob: Vec<u8> = e.vector.iter().flat_map(|f| f.to_le_bytes()).collect();
139                let meta = e.metadata.as_ref().map(|m| m.to_string());
140                // INSERT OR IGNORE: changes()=1 for new, 0 for existing
141                let inserted = conn.execute(
142                    "INSERT OR IGNORE INTO _adb_vectors
143                         (id, collection_id, vector, metadata, created_at)
144                     VALUES (?1, ?2, ?3, ?4, ?5)",
145                    params![e.id, self.id, blob, meta, now_ms()],
146                )?;
147                if inserted == 0 {
148                    conn.execute(
149                        "UPDATE _adb_vectors SET vector = ?1, metadata = ?2
150                         WHERE id = ?3 AND collection_id = ?4",
151                        params![blob, meta, e.id, self.id],
152                    )?;
153                } else {
154                    new_rows += 1;
155                }
156            }
157            conn.execute(
158                "INSERT INTO _adb_hnsw_index (collection_id, index_blob, built_at, is_dirty)
159                 VALUES (?1, X'', ?2, 1)
160                 ON CONFLICT(collection_id) DO UPDATE SET is_dirty = 1",
161                params![self.id, now_ms()],
162            )?;
163            if new_rows > 0 {
164                conn.execute(
165                    "UPDATE _adb_collections SET count = count + ?1 WHERE id = ?2",
166                    params![new_rows as i64, self.id],
167                )?;
168            }
169            Ok(new_rows)
170        })();
171        match result {
172            Ok(new_rows) => {
173                conn.execute_batch("COMMIT")?;
174                *self.index.lock().unwrap() = None;
175                Ok(new_rows)
176            }
177            Err(e) => {
178                let _ = conn.execute_batch("ROLLBACK");
179                Err(e)
180            }
181        }
182    }
183
184    /// ANN search with optional advanced metadata filtering
185    pub fn search(&self, query: &[f32], opts: SearchOptions) -> Result<Vec<SearchResult>> {
186        if query.len() != self.dim {
187            return Err(AgentDbError::DimensionMismatch {
188                expected: self.dim,
189                got: query.len(),
190            });
191        }
192        self.ensure_index()?;
193        let guard = self.index.lock().unwrap();
194        let index = guard.as_ref().unwrap();
195        let fetch_k = if opts.filter.is_some() {
196            (opts.top_k * 10).max(50)
197        } else {
198            opts.top_k
199        };
200        let raw = index.search(query, fetch_k);
201        let conn = self.conn.lock().unwrap();
202        let mut out = Vec::new();
203        for (id, score) in raw {
204            let meta_str: Option<String> = conn
205                .query_row(
206                    "SELECT metadata FROM _adb_vectors
207                     WHERE id = ?1 AND collection_id = ?2",
208                    params![id, self.id],
209                    |r| r.get(0),
210                )
211                .ok()
212                .flatten();
213            let meta: Option<Value> = meta_str
214                .as_deref()
215                .and_then(|s| serde_json::from_str(s).ok());
216            if let Some(ref f) = opts.filter {
217                match &meta {
218                    Some(m) if filter::matches(m, f) => {}
219                    _ => continue,
220                }
221            }
222            out.push(SearchResult {
223                id,
224                score,
225                metadata: meta,
226            });
227            if out.len() >= opts.top_k {
228                break;
229            }
230        }
231        Ok(out)
232    }
233
234    /// Delete a vector by ID
235    pub fn delete(&self, id: &str) -> Result<()> {
236        let conn = self.conn.lock().unwrap();
237        conn.execute(
238            "DELETE FROM _adb_vectors WHERE id = ?1 AND collection_id = ?2",
239            params![id, self.id],
240        )?;
241        conn.execute(
242            "UPDATE _adb_hnsw_index SET is_dirty = 1 WHERE collection_id = ?1",
243            params![self.id],
244        )?;
245        *self.index.lock().unwrap() = None;
246        Ok(())
247    }
248
249    /// Rebuild the HNSW index from stored vectors
250    pub fn reindex(&self) -> Result<()> {
251        let mut index = HnswIndex::new(16, 200, self.metric.clone());
252        // Scope the borrow of conn so stmt and rows are dropped before the second lock.
253        {
254            let conn = self.conn.lock().unwrap();
255            let mut stmt =
256                conn.prepare("SELECT id, vector FROM _adb_vectors WHERE collection_id = ?1")?;
257            let rows = stmt.query_map(params![self.id], |row| {
258                Ok((row.get::<_, String>(0)?, row.get::<_, Vec<u8>>(1)?))
259            })?;
260            for row in rows {
261                let (id, blob) = row?;
262                let vec: Vec<f32> = blob
263                    .chunks_exact(4)
264                    .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
265                    .collect();
266                index.insert(&id, vec);
267            }
268        }
269        let serialized = index.serialize()?;
270        {
271            let conn = self.conn.lock().unwrap();
272            conn.execute(
273                "INSERT INTO _adb_hnsw_index (collection_id, index_blob, built_at, is_dirty)
274                 VALUES (?1, ?2, ?3, 0)
275                 ON CONFLICT(collection_id) DO UPDATE SET
276                   index_blob = excluded.index_blob,
277                   built_at   = excluded.built_at,
278                   is_dirty   = 0",
279                params![self.id, serialized, now_ms()],
280            )?;
281        }
282        *self.index.lock().unwrap() = Some(index);
283        Ok(())
284    }
285
286    fn ensure_index(&self) -> Result<()> {
287        if self.index.lock().unwrap().is_some() {
288            return Ok(());
289        }
290        self.reindex()
291    }
292
293    /// Number of vectors in this collection
294    pub fn count(&self) -> Result<i64> {
295        let conn = self.conn.lock().unwrap();
296        Ok(conn.query_row(
297            "SELECT count FROM _adb_collections WHERE id = ?1",
298            params![self.id],
299            |r| r.get(0),
300        )?)
301    }
302}
303
304/// Manages all vector collections
305pub struct VectorStore {
306    conn: Arc<Mutex<Connection>>,
307}
308
309impl VectorStore {
310    pub(crate) fn new(conn: Arc<Mutex<Connection>>) -> Self {
311        Self { conn }
312    }
313
314    pub fn collection(&self, name: &str, dim: usize) -> Result<Collection> {
315        self.collection_with_metric(name, dim, DistanceMetric::Cosine)
316    }
317
318    pub fn collection_with_metric(
319        &self,
320        name: &str,
321        dim: usize,
322        metric: DistanceMetric,
323    ) -> Result<Collection> {
324        let conn = self.conn.lock().unwrap();
325        let existing: Option<(String, usize, String)> = conn
326            .query_row(
327                "SELECT id, dim, metric FROM _adb_collections WHERE name = ?1",
328                params![name],
329                |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)),
330            )
331            .ok();
332        if let Some((id, edim, mstr)) = existing {
333            if edim != dim {
334                return Err(AgentDbError::DimensionMismatch {
335                    expected: edim,
336                    got: dim,
337                });
338            }
339            let m = match mstr.as_str() {
340                "euclidean" => DistanceMetric::Euclidean,
341                "dot" => DistanceMetric::DotProduct,
342                _ => DistanceMetric::Cosine,
343            };
344            return Ok(Collection::new(
345                id,
346                name.to_string(),
347                dim,
348                m,
349                Arc::clone(&self.conn),
350            ));
351        }
352        let id = Uuid::new_v4().to_string();
353        let mstr = match &metric {
354            DistanceMetric::Cosine => "cosine",
355            DistanceMetric::Euclidean => "euclidean",
356            DistanceMetric::DotProduct => "dot",
357        };
358        conn.execute(
359            "INSERT INTO _adb_collections (id, name, dim, metric, count, created_at)
360             VALUES (?1, ?2, ?3, ?4, 0, ?5)",
361            params![id, name, dim, mstr, now_ms()],
362        )?;
363        Ok(Collection::new(
364            id,
365            name.to_string(),
366            dim,
367            metric,
368            Arc::clone(&self.conn),
369        ))
370    }
371
372    pub fn list_collections(&self) -> Result<Vec<(String, usize, i64)>> {
373        let conn = self.conn.lock().unwrap();
374        let mut stmt =
375            conn.prepare("SELECT name, dim, count FROM _adb_collections ORDER BY name")?;
376        let rows = stmt.query_map([], |row| {
377            Ok((
378                row.get::<_, String>(0)?,
379                row.get::<_, usize>(1)?,
380                row.get::<_, i64>(2)?,
381            ))
382        })?;
383        rows.map(|r| r.map_err(AgentDbError::Sqlite)).collect()
384    }
385
386    pub fn drop_collection(&self, name: &str) -> Result<()> {
387        let conn = self.conn.lock().unwrap();
388        conn.execute(
389            "DELETE FROM _adb_collections WHERE name = ?1",
390            params![name],
391        )?;
392        Ok(())
393    }
394}
agentdb/vectors/collection.rs

agentdb/vectors/
collection.rs
