rightkit_embed/
vector.rs

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

rightkit_embed/
vector.rs
