agentdb/vectors/
collection.rs

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

agentdb/vectors/
collection.rs
