rightkit_embed/
vector.rs

1use rusqlite::{params, Connection, OptionalExtension};
2use std::path::Path;
3use thiserror::Error;
4
5pub const DEFAULT_VECTOR_DIM: usize = 256;
6const SCHEMA: &str = "CREATE TABLE IF NOT EXISTS rightkit_vectors (
7    key TEXT PRIMARY KEY NOT NULL,
8    model_id TEXT NOT NULL,
9    dimensions INTEGER NOT NULL,
10    vector BLOB NOT NULL
11);";
12
13#[derive(Debug, Error)]
14pub enum VectorError {
15    #[error("SQLite: {0}")]
16    Sqlite(#[from] rusqlite::Error),
17    #[error("vector dimension must be {expected}, got {actual}")]
18    Dimension { expected: usize, actual: usize },
19    #[error("vector contains non-finite values")]
20    NonFinite,
21    #[error("stored vector for {key} is malformed")]
22    Malformed { key: String },
23}
24
25#[derive(Debug, Clone, PartialEq)]
26pub struct Nearest {
27    pub key: String,
28    pub score: f32,
29    pub vector: Vec<f32>,
30}
31
32/// Fold a bootstrap failure into the existing error type (keeps `VectorError` exhaustive
33/// for callers): SQLite errors pass through; filesystem/bootstrap errors surface as
34/// `SQLITE_CANTOPEN` with the message attached.
35fn db_error(error: rightkit_db::DbError) -> VectorError {
36    match error {
37        rightkit_db::DbError::Sqlite(e) => VectorError::Sqlite(e),
38        other => VectorError::Sqlite(rusqlite::Error::SqliteFailure(
39            rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_CANTOPEN),
40            Some(other.to_string()),
41        )),
42    }
43}
44
45pub struct VectorStore {
46    conn: Connection,
47    dimension: usize,
48}
49
50impl VectorStore {
51    pub fn open(path: impl AsRef<Path>, dimension: usize) -> Result<Self, VectorError> {
52        let conn = rightkit_db::DbOptions::new(path.as_ref())
53            .open_writer()
54            .map_err(db_error)?;
55        Self::from_connection(conn, dimension)
56    }
57
58    pub fn in_memory(dimension: usize) -> Result<Self, VectorError> {
59        let conn = rightkit_db::open_in_memory().map_err(db_error)?;
60        Self::from_connection(conn, dimension)
61    }
62
63    pub fn from_connection(conn: Connection, dimension: usize) -> Result<Self, VectorError> {
64        if dimension == 0 {
65            return Err(VectorError::Dimension {
66                expected: 1,
67                actual: 0,
68            });
69        }
70        conn.execute_batch(SCHEMA)?;
71        Ok(Self { conn, dimension })
72    }
73
74    pub fn dimension(&self) -> usize {
75        self.dimension
76    }
77
78    pub fn len(&self) -> Result<usize, VectorError> {
79        let count: i64 =
80            self.conn
81                .query_row("SELECT COUNT(*) FROM rightkit_vectors", [], |row| {
82                    row.get(0)
83                })?;
84        Ok(count as usize)
85    }
86
87    pub fn is_empty(&self) -> Result<bool, VectorError> {
88        Ok(self.len()? == 0)
89    }
90
91    pub fn put(&self, key: &str, model_id: &str, vector: &[f32]) -> Result<(), VectorError> {
92        validate(vector, self.dimension)?;
93        self.conn.execute(
94            "INSERT INTO rightkit_vectors(key, model_id, dimensions, vector)
95             VALUES (?1, ?2, ?3, ?4)
96             ON CONFLICT(key) DO UPDATE SET model_id=excluded.model_id,
97             dimensions=excluded.dimensions, vector=excluded.vector",
98            params![key, model_id, self.dimension as i64, encode(vector)],
99        )?;
100        Ok(())
101    }
102
103    pub fn put_batch<'a, I>(&mut self, rows: I) -> Result<(), VectorError>
104    where
105        I: IntoIterator<Item = (&'a str, &'a str, &'a [f32])>,
106    {
107        let tx = self.conn.transaction()?;
108        for (key, model_id, vector) in rows {
109            validate(vector, self.dimension)?;
110            tx.execute(
111                "INSERT INTO rightkit_vectors(key, model_id, dimensions, vector)
112                 VALUES (?1, ?2, ?3, ?4)
113                 ON CONFLICT(key) DO UPDATE SET model_id=excluded.model_id,
114                 dimensions=excluded.dimensions, vector=excluded.vector",
115                params![key, model_id, self.dimension as i64, encode(vector)],
116            )?;
117        }
118        tx.commit()?;
119        Ok(())
120    }
121
122    pub fn get(&self, key: &str) -> Result<Option<Nearest>, VectorError> {
123        let row = self
124            .conn
125            .query_row(
126                "SELECT dimensions, vector FROM rightkit_vectors WHERE key=?1",
127                params![key],
128                |row| {
129                    let dimensions: i64 = row.get(0)?;
130                    let bytes: Vec<u8> = row.get(1)?;
131                    Ok((dimensions, bytes))
132                },
133            )
134            .optional()?;
135        row.map(|(dimensions, bytes)| {
136            if dimensions != self.dimension as i64 {
137                return Err(VectorError::Malformed {
138                    key: key.to_owned(),
139                });
140            }
141            let vector = decode(&bytes, self.dimension).map_err(|_| VectorError::Malformed {
142                key: key.to_owned(),
143            })?;
144            Ok(Nearest {
145                key: key.to_owned(),
146                score: 1.0,
147                vector,
148            })
149        })
150        .transpose()
151    }
152
153    pub fn nearest(
154        &self,
155        model_id: &str,
156        query: &[f32],
157        limit: usize,
158    ) -> Result<Vec<Nearest>, VectorError> {
159        validate(query, self.dimension)?;
160        if limit == 0 {
161            return Ok(Vec::new());
162        }
163        let mut stmt = self.conn.prepare(
164            "SELECT key, dimensions, vector FROM rightkit_vectors
165             WHERE model_id=?1 ORDER BY key",
166        )?;
167        let rows = stmt.query_map(params![model_id], |row| {
168            let key: String = row.get(0)?;
169            let dimensions: i64 = row.get(1)?;
170            let bytes: Vec<u8> = row.get(2)?;
171            Ok((key, dimensions, bytes))
172        })?;
173        let mut candidates = Vec::new();
174        for row in rows {
175            let (key, dimensions, bytes) = row?;
176            if dimensions != self.dimension as i64 {
177                return Err(VectorError::Malformed { key });
178            }
179            let vector = decode(&bytes, self.dimension)
180                .map_err(|_| VectorError::Malformed { key: key.clone() })?;
181            candidates.push(Nearest {
182                score: cosine_similarity(query, &vector),
183                key,
184                vector,
185            });
186        }
187        candidates.sort_by(|left, right| {
188            right
189                .score
190                .total_cmp(&left.score)
191                .then_with(|| left.key.cmp(&right.key))
192        });
193        candidates.truncate(limit);
194        Ok(candidates)
195    }
196
197    pub fn remove(&self, key: &str) -> Result<bool, VectorError> {
198        Ok(self
199            .conn
200            .execute("DELETE FROM rightkit_vectors WHERE key=?1", params![key])?
201            != 0)
202    }
203
204    pub fn clear_model(&self, model_id: &str) -> Result<usize, VectorError> {
205        Ok(self.conn.execute(
206            "DELETE FROM rightkit_vectors WHERE model_id=?1",
207            params![model_id],
208        )?)
209    }
210}
211
212pub fn cosine_similarity(left: &[f32], right: &[f32]) -> f32 {
213    if left.len() != right.len() || left.is_empty() {
214        return 0.0;
215    }
216    let mut dot = 0.0;
217    let mut left_norm = 0.0;
218    let mut right_norm = 0.0;
219    for (&a, &b) in left.iter().zip(right) {
220        dot += a * b;
221        left_norm += a * a;
222        right_norm += b * b;
223    }
224    if left_norm == 0.0 || right_norm == 0.0 {
225        0.0
226    } else {
227        dot / (left_norm.sqrt() * right_norm.sqrt())
228    }
229}
230
231fn validate(vector: &[f32], dimension: usize) -> Result<(), VectorError> {
232    if vector.len() != dimension {
233        return Err(VectorError::Dimension {
234            expected: dimension,
235            actual: vector.len(),
236        });
237    }
238    if vector.iter().any(|value| !value.is_finite()) {
239        return Err(VectorError::NonFinite);
240    }
241    Ok(())
242}
243
244fn encode(vector: &[f32]) -> Vec<u8> {
245    vector
246        .iter()
247        .flat_map(|value| value.to_le_bytes())
248        .collect()
249}
250
251fn decode(bytes: &[u8], dimension: usize) -> Result<Vec<f32>, ()> {
252    if bytes.len() != dimension * std::mem::size_of::<f32>() {
253        return Err(());
254    }
255    let mut vector = Vec::with_capacity(dimension);
256    for chunk in bytes.as_chunks::<4>().0 {
257        vector.push(f32::from_le_bytes(*chunk));
258    }
259    if vector.iter().any(|value| !value.is_finite()) {
260        return Err(());
261    }
262    Ok(vector)
263}
rightkit_embed/vector.rs

rightkit_embed/
vector.rs
