Skip to main content

vein_database/vector/
mod.rs

1pub mod lancedb;
2pub mod table_names;
3
4use anyhow::Result;
5
6pub use lancedb::Entry;
7pub use table_names::list_model_tables;
8
9enum LanceBackend {
10    Lance(lancedb::LanceDb),
11}
12
13pub struct VectorDb {
14    backend: LanceBackend,
15}
16
17impl VectorDb {
18    pub async fn new(table_name: &str, vector_dim: usize) -> Result<Self> {
19        let db = lancedb::LanceDb::new(table_name, vector_dim).await?;
20        Ok(Self {
21            backend: LanceBackend::Lance(db),
22        })
23    }
24
25    pub async fn post(&self, id: &str, content: &str, vector: Vec<f32>) -> Result<String> {
26        match &self.backend {
27            LanceBackend::Lance(db) => db.post(id, content, vector).await,
28        }
29    }
30
31    pub async fn get(&self, query_vector: &[f32], limit: usize) -> Result<Vec<Entry>> {
32        match &self.backend {
33            LanceBackend::Lance(db) => db.get(query_vector, limit).await,
34        }
35    }
36
37    pub async fn patch(&self, id: &str, new_content: &str, new_vector: Vec<f32>) -> Result<()> {
38        match &self.backend {
39            LanceBackend::Lance(db) => db.patch(id, new_content, new_vector).await,
40        }
41    }
42
43    pub async fn delete(&self, id: &str) -> Result<()> {
44        match &self.backend {
45            LanceBackend::Lance(db) => db.delete(id).await,
46        }
47    }
48
49    pub async fn rebuild_index(&self) -> Result<()> {
50        match &self.backend {
51            LanceBackend::Lance(db) => db.rebuild_index().await,
52        }
53    }
54}