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
32pub struct VectorStore {
33 conn: Connection,
34 dimension: usize,
35}
36
37impl VectorStore {
38 pub fn open(path: impl AsRef<Path>, dimension: usize) -> Result<Self, VectorError> {
39 Self::from_connection(Connection::open(path)?, dimension)
40 }
41
42 pub fn in_memory(dimension: usize) -> Result<Self, VectorError> {
43 Self::from_connection(Connection::open_in_memory()?, dimension)
44 }
45
46 pub fn from_connection(conn: Connection, dimension: usize) -> Result<Self, VectorError> {
47 if dimension == 0 {
48 return Err(VectorError::Dimension {
49 expected: 1,
50 actual: 0,
51 });
52 }
53 conn.execute_batch(SCHEMA)?;
54 Ok(Self { conn, dimension })
55 }
56
57 pub fn dimension(&self) -> usize {
58 self.dimension
59 }
60
61 pub fn len(&self) -> Result<usize, VectorError> {
62 let count: i64 =
63 self.conn
64 .query_row("SELECT COUNT(*) FROM rightkit_vectors", [], |row| {
65 row.get(0)
66 })?;
67 Ok(count as usize)
68 }
69
70 pub fn is_empty(&self) -> Result<bool, VectorError> {
71 Ok(self.len()? == 0)
72 }
73
74 pub fn put(&self, key: &str, model_id: &str, vector: &[f32]) -> Result<(), VectorError> {
75 validate(vector, self.dimension)?;
76 self.conn.execute(
77 "INSERT INTO rightkit_vectors(key, model_id, dimensions, vector)
78 VALUES (?1, ?2, ?3, ?4)
79 ON CONFLICT(key) DO UPDATE SET model_id=excluded.model_id,
80 dimensions=excluded.dimensions, vector=excluded.vector",
81 params![key, model_id, self.dimension as i64, encode(vector)],
82 )?;
83 Ok(())
84 }
85
86 pub fn put_batch<'a, I>(&mut self, rows: I) -> Result<(), VectorError>
87 where
88 I: IntoIterator<Item = (&'a str, &'a str, &'a [f32])>,
89 {
90 let tx = self.conn.transaction()?;
91 for (key, model_id, vector) in rows {
92 validate(vector, self.dimension)?;
93 tx.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 }
101 tx.commit()?;
102 Ok(())
103 }
104
105 pub fn get(&self, key: &str) -> Result<Option<Nearest>, VectorError> {
106 let row = self
107 .conn
108 .query_row(
109 "SELECT dimensions, vector FROM rightkit_vectors WHERE key=?1",
110 params![key],
111 |row| {
112 let dimensions: i64 = row.get(0)?;
113 let bytes: Vec<u8> = row.get(1)?;
114 Ok((dimensions, bytes))
115 },
116 )
117 .optional()?;
118 row.map(|(dimensions, bytes)| {
119 if dimensions != self.dimension as i64 {
120 return Err(VectorError::Malformed {
121 key: key.to_owned(),
122 });
123 }
124 let vector = decode(&bytes, self.dimension).map_err(|_| VectorError::Malformed {
125 key: key.to_owned(),
126 })?;
127 Ok(Nearest {
128 key: key.to_owned(),
129 score: 1.0,
130 vector,
131 })
132 })
133 .transpose()
134 }
135
136 pub fn nearest(
137 &self,
138 model_id: &str,
139 query: &[f32],
140 limit: usize,
141 ) -> Result<Vec<Nearest>, VectorError> {
142 validate(query, self.dimension)?;
143 if limit == 0 {
144 return Ok(Vec::new());
145 }
146 let mut stmt = self.conn.prepare(
147 "SELECT key, dimensions, vector FROM rightkit_vectors
148 WHERE model_id=?1 ORDER BY key",
149 )?;
150 let rows = stmt.query_map(params![model_id], |row| {
151 let key: String = row.get(0)?;
152 let dimensions: i64 = row.get(1)?;
153 let bytes: Vec<u8> = row.get(2)?;
154 Ok((key, dimensions, bytes))
155 })?;
156 let mut candidates = Vec::new();
157 for row in rows {
158 let (key, dimensions, bytes) = row?;
159 if dimensions != self.dimension as i64 {
160 return Err(VectorError::Malformed { key });
161 }
162 let vector = decode(&bytes, self.dimension)
163 .map_err(|_| VectorError::Malformed { key: key.clone() })?;
164 candidates.push(Nearest {
165 score: cosine_similarity(query, &vector),
166 key,
167 vector,
168 });
169 }
170 candidates.sort_by(|left, right| {
171 right
172 .score
173 .total_cmp(&left.score)
174 .then_with(|| left.key.cmp(&right.key))
175 });
176 candidates.truncate(limit);
177 Ok(candidates)
178 }
179
180 pub fn remove(&self, key: &str) -> Result<bool, VectorError> {
181 Ok(self
182 .conn
183 .execute("DELETE FROM rightkit_vectors WHERE key=?1", params![key])?
184 != 0)
185 }
186
187 pub fn clear_model(&self, model_id: &str) -> Result<usize, VectorError> {
188 Ok(self.conn.execute(
189 "DELETE FROM rightkit_vectors WHERE model_id=?1",
190 params![model_id],
191 )?)
192 }
193}
194
195pub fn cosine_similarity(left: &[f32], right: &[f32]) -> f32 {
196 if left.len() != right.len() || left.is_empty() {
197 return 0.0;
198 }
199 let mut dot = 0.0;
200 let mut left_norm = 0.0;
201 let mut right_norm = 0.0;
202 for (&a, &b) in left.iter().zip(right) {
203 dot += a * b;
204 left_norm += a * a;
205 right_norm += b * b;
206 }
207 if left_norm == 0.0 || right_norm == 0.0 {
208 0.0
209 } else {
210 dot / (left_norm.sqrt() * right_norm.sqrt())
211 }
212}
213
214fn validate(vector: &[f32], dimension: usize) -> Result<(), VectorError> {
215 if vector.len() != dimension {
216 return Err(VectorError::Dimension {
217 expected: dimension,
218 actual: vector.len(),
219 });
220 }
221 if vector.iter().any(|value| !value.is_finite()) {
222 return Err(VectorError::NonFinite);
223 }
224 Ok(())
225}
226
227fn encode(vector: &[f32]) -> Vec<u8> {
228 vector
229 .iter()
230 .flat_map(|value| value.to_le_bytes())
231 .collect()
232}
233
234fn decode(bytes: &[u8], dimension: usize) -> Result<Vec<f32>, ()> {
235 if bytes.len() != dimension * std::mem::size_of::<f32>() {
236 return Err(());
237 }
238 let mut vector = Vec::with_capacity(dimension);
239 for chunk in bytes.as_chunks::<4>().0 {
240 vector.push(f32::from_le_bytes(*chunk));
241 }
242 if vector.iter().any(|value| !value.is_finite()) {
243 return Err(());
244 }
245 Ok(vector)
246}