1use rusqlite::{Connection, Result, params};
4
5pub const EMBEDDABLE_EXTS: &[&str] = &[
17 "jpg", "jpeg", "png", "gif", "webp", "bmp", "tiff", "heic", "mov", "mp4",
18];
19
20pub fn is_video_ext(ext: &str) -> bool {
24 matches!(ext.to_lowercase().as_str(), "mov" | "mp4")
25}
26
27pub const DEFAULT_MODEL_ID: &str = "google/siglip-base-patch16-224";
48
49pub fn resolve_model_id_in(
62 home: &std::path::Path,
63 explicit: Option<&str>,
64) -> anyhow::Result<String> {
65 if let Some(id) = explicit {
66 return Ok(id.to_string());
67 }
68 Ok(crate::home::load_config(home)?
69 .default_model
70 .unwrap_or_else(|| DEFAULT_MODEL_ID.to_string()))
71}
72
73pub fn resolve_model_id(explicit: Option<&str>) -> anyhow::Result<String> {
79 resolve_model_id_in(&crate::home::videre_home()?, explicit)
80}
81
82#[derive(Debug, Clone)]
83pub struct PendingImage {
84 pub hash: String,
85 pub path: String,
86}
87
88pub fn ensure_embeddings_index(conn: &Connection) -> Result<()> {
96 conn.execute_batch(
97 "CREATE INDEX IF NOT EXISTS idx_file_hashes_hash ON file_hashes(hash);",
98 )
99}
100
101pub fn pending_images(conn: &Connection, model_id: &str) -> Result<Vec<PendingImage>> {
104 let mimes = crate::mime_probe::EMBEDDABLE_MIMES
105 .iter()
106 .map(|m| format!("'{m}'"))
107 .collect::<Vec<_>>()
108 .join(",");
109 let exts = EMBEDDABLE_EXTS
110 .iter()
111 .map(|e| format!("'{e}'"))
112 .collect::<Vec<_>>()
113 .join(",");
114 let sql = format!(
121 "SELECT hash, MIN(path) FROM file_hashes
122 WHERE lower(COALESCE(ext, '')) != 'dng'
123 AND (mime IN ({mimes}) OR (mime IS NULL AND lower(ext) IN ({exts})))
124 AND NOT EXISTS (SELECT 1 FROM emb.embeddings e
125 WHERE e.hash = file_hashes.hash AND e.model_id = ?1)
126 GROUP BY hash
127 ORDER BY hash"
128 );
129 let mut stmt = conn.prepare(&sql)?;
130 let rows = stmt.query_map(params![model_id], |row| {
131 Ok(PendingImage {
132 hash: row.get(0)?,
133 path: row.get(1)?,
134 })
135 })?;
136 rows.collect()
137}
138
139pub fn insert_embeddings(
141 conn: &Connection,
142 model_id: &str,
143 items: &[(String, Vec<u8>)],
144) -> Result<()> {
145 let tx = conn.unchecked_transaction()?;
146 {
147 let mut stmt = tx.prepare(
148 "INSERT OR REPLACE INTO emb.embeddings (hash, model_id, embedding, embedded_at)
149 VALUES (?1, ?2, ?3, datetime('now'))",
150 )?;
151 for (hash, blob) in items {
152 stmt.execute(params![hash, model_id, blob])?;
153 }
154 }
155 tx.commit()
156}
157
158pub fn load_embeddings(conn: &Connection, model_id: &str) -> Result<Vec<(String, Vec<u8>)>> {
162 let attached: bool = conn
163 .query_row(
164 "SELECT COUNT(*) FROM emb.sqlite_master WHERE type='table' AND name='embeddings'",
165 [],
166 |r| r.get::<_, i64>(0),
167 )
168 .unwrap_or(0)
169 > 0;
170 if !attached {
171 return Ok(Vec::new());
172 }
173 let mut stmt =
174 conn.prepare("SELECT hash, embedding FROM emb.embeddings WHERE model_id = ?1")?;
175 let rows = stmt.query_map(params![model_id], |row| Ok((row.get(0)?, row.get(1)?)))?;
176 rows.collect()
177}
178
179pub fn paths_for_hash(conn: &Connection, hash: &str) -> Result<Vec<String>> {
180 let mut stmt =
181 conn.prepare("SELECT path FROM file_hashes WHERE hash = ?1 ORDER BY path")?;
182 let rows = stmt.query_map(params![hash], |row| row.get(0))?;
183 rows.collect()
184}
185
186#[cfg(test)]
187mod tests {
188 use super::*;
189 use rusqlite::Connection;
190
191 fn test_db_attached(tag: &str) -> Connection {
196 let lib = crate::embeddings_db::test_library(tag);
197 let conn = Connection::open_in_memory().unwrap();
198 conn.execute_batch(
199 "CREATE TABLE file_hashes (
200 path TEXT PRIMARY KEY,
201 hash TEXT NOT NULL,
202 mime TEXT,
203 size_bytes INTEGER,
204 created_at TEXT,
205 modified_at TEXT,
206 ext TEXT,
207 phash INTEGER,
208 exif_date TEXT,
209 gps_lat REAL,
210 gps_lon REAL,
211 width INTEGER,
212 height INTEGER
213 );",
214 )
215 .unwrap();
216 ensure_embeddings_index(&conn).unwrap();
217 crate::embeddings_db::attach(&conn, &lib, "test-model", true).unwrap();
218 conn
219 }
220
221 fn insert_file(conn: &Connection, path: &str, hash: &str, ext: &str) {
222 conn.execute(
223 "INSERT INTO file_hashes (path, hash, ext) VALUES (?1, ?2, ?3)",
224 rusqlite::params![path, hash, ext],
225 )
226 .unwrap();
227 }
228
229 #[test]
230 fn pending_images_dedupes_by_hash_and_includes_video() {
231 let conn = test_db_attached("emb_dedupe");
232 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
233 insert_file(&conn, "/b/1-copy.jpg", "h1", "jpg"); insert_file(&conn, "/a/2.png", "h2", "png");
235 insert_file(&conn, "/a/clip.mp4", "h3", "mp4"); insert_file(&conn, "/a/other.xyz", "h4", "xyz"); let pending = pending_images(&conn, "test-model").unwrap();
239 assert_eq!(pending.len(), 3); assert!(pending.iter().any(|p| p.hash == "h1"));
241 assert!(pending.iter().any(|p| p.hash == "h2"));
242 assert!(pending.iter().any(|p| p.hash == "h3"));
243 }
244
245 #[test]
246 fn pending_images_excludes_dng_since_it_cannot_be_decoded() {
247 let conn = test_db_attached("emb_dng");
248 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
249 insert_file(&conn, "/a/raw.dng", "h2", "dng");
250
251 let pending = pending_images(&conn, "test-model").unwrap();
252 assert_eq!(pending.len(), 1);
253 assert_eq!(pending[0].hash, "h1");
254 }
255
256 #[test]
257 fn pending_images_excludes_already_embedded() {
258 let conn = test_db_attached("emb_already");
259 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
260 insert_file(&conn, "/a/2.jpg", "h2", "jpg");
261 insert_embeddings(&conn, "test-model", &[("h1".to_string(), vec![0u8; 4])]).unwrap();
262
263 let pending = pending_images(&conn, "test-model").unwrap();
264 assert_eq!(pending.len(), 1);
265 assert_eq!(pending[0].hash, "h2");
266 }
267
268 #[test]
269 fn pending_images_is_model_aware() {
270 let conn = test_db_attached("emb_modelaware");
271 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
272 insert_embeddings(&conn, "a", &[("h1".to_string(), vec![0u8; 4])]).unwrap();
273
274 assert!(pending_images(&conn, "a").unwrap().is_empty());
276
277 let pending = pending_images(&conn, "b").unwrap();
279 assert_eq!(pending.len(), 1);
280 assert_eq!(pending[0].hash, "h1");
281 }
282
283 #[test]
284 fn pending_images_uses_mime_over_a_wrong_extension() {
285 let conn = test_db_attached("emb_mime");
286 conn.execute(
287 "INSERT INTO file_hashes (path, hash, ext, mime)
288 VALUES ('/a/actually_a_jpeg.png', 'h1', 'png', 'image/jpeg')",
289 [],
290 )
291 .unwrap();
292 let pending = pending_images(&conn, "m").unwrap();
293 assert_eq!(pending.len(), 1, "a JPEG named .png must still be embeddable");
294 }
295
296 #[test]
297 fn pending_images_falls_back_to_ext_when_mime_is_null() {
298 let conn = test_db_attached("emb_nullmime");
299 conn.execute(
300 "INSERT INTO file_hashes (path, hash, ext, mime) VALUES ('/a/1.jpg', 'h1', 'jpg', NULL)",
301 [],
302 )
303 .unwrap();
304 assert_eq!(pending_images(&conn, "m").unwrap().len(), 1);
305 }
306
307 #[test]
308 fn pending_images_still_excludes_dng_even_though_its_mime_is_tiff() {
309 let conn = test_db_attached("emb_dng_mime");
312 conn.execute(
313 "INSERT INTO file_hashes (path, hash, ext, mime)
314 VALUES ('/a/raw.dng', 'h1', 'dng', 'image/tiff')",
315 [],
316 )
317 .unwrap();
318 assert!(pending_images(&conn, "m").unwrap().is_empty());
319 }
320
321 #[test]
322 fn insert_embeddings_empty_slice_succeeds() {
323 let conn = test_db_attached("emb_empty");
324 insert_embeddings(&conn, "test-model", &[]).unwrap();
325 assert!(load_embeddings(&conn, "test-model").unwrap().is_empty());
326 }
327
328 #[test]
329 fn insert_and_load_round_trip() {
330 let conn = test_db_attached("emb_roundtrip");
331 insert_embeddings(
332 &conn,
333 "test-model",
334 &[("h1".to_string(), vec![1u8, 2, 3, 4]), ("h2".to_string(), vec![5u8, 6])],
335 )
336 .unwrap();
337
338 let rows = load_embeddings(&conn, "test-model").unwrap();
339 assert_eq!(rows.len(), 2);
340 let h1 = rows.iter().find(|(h, _)| h == "h1").unwrap();
341 assert_eq!(h1.1, vec![1u8, 2, 3, 4]);
342
343 assert!(load_embeddings(&conn, "other").unwrap().is_empty());
345 }
346
347 #[test]
348 fn load_embeddings_finds_the_table_in_the_attached_database() {
349 let conn = test_db_attached("emb_attachedprobe");
354 insert_embeddings(&conn, "test-model", &[("h1".to_string(), vec![1u8, 2])]).unwrap();
355
356 let rows = load_embeddings(&conn, "test-model").unwrap();
357 assert_eq!(rows.len(), 1, "must read through emb., not main");
358 }
359
360 #[test]
361 fn ensure_embeddings_index_creates_the_index_in_the_main_database() {
362 let conn = test_db_attached("emb_indexmain");
365 let found: i64 = conn
366 .query_row(
367 "SELECT COUNT(*) FROM main.sqlite_master
368 WHERE type='index' AND name='idx_file_hashes_hash'",
369 [],
370 |r| r.get(0),
371 )
372 .unwrap();
373 assert_eq!(found, 1);
374 }
375
376 #[test]
377 fn paths_for_hash_returns_all_duplicates() {
378 let conn = test_db_attached("emb_paths");
379 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
380 insert_file(&conn, "/b/1-copy.jpg", "h1", "jpg");
381 let paths = paths_for_hash(&conn, "h1").unwrap();
382 assert_eq!(paths.len(), 2);
383 }
384
385 #[test]
386 fn default_model_id_is_a_siglip_checkpoint() {
387 assert!(DEFAULT_MODEL_ID.starts_with("google/"), "{DEFAULT_MODEL_ID}");
394 assert!(DEFAULT_MODEL_ID.contains("siglip"), "{DEFAULT_MODEL_ID}");
395 assert_eq!(DEFAULT_MODEL_ID.matches('/').count(), 1, "{DEFAULT_MODEL_ID}");
396 }
397
398 fn cfg_home(tag: &str, toml_text: &str) -> std::path::PathBuf {
399 let dir = std::env::temp_dir().join(format!("videre_rmi_{}_{}", tag, std::process::id()));
400 let _ = std::fs::remove_dir_all(&dir);
401 std::fs::create_dir_all(&dir).unwrap();
402 if !toml_text.is_empty() {
403 std::fs::write(dir.join("config.toml"), toml_text).unwrap();
404 }
405 dir
406 }
407
408 #[test]
409 fn resolve_model_id_prefers_the_explicit_argument() {
410 let home = cfg_home("explicit", "default_model = \"owner/from-config\"\n");
411 assert_eq!(
412 resolve_model_id_in(&home, Some("owner/explicit")).unwrap(),
413 "owner/explicit"
414 );
415 let _ = std::fs::remove_dir_all(&home);
416 }
417
418 #[test]
419 fn resolve_model_id_uses_config_when_there_is_no_flag() {
420 let home = cfg_home("fromconfig", "default_model = \"owner/from-config\"\n");
421 assert_eq!(resolve_model_id_in(&home, None).unwrap(), "owner/from-config");
422 let _ = std::fs::remove_dir_all(&home);
423 }
424
425 #[test]
426 fn resolve_model_id_falls_back_to_the_builtin_default() {
427 let home = cfg_home("builtin", "");
428 assert_eq!(resolve_model_id_in(&home, None).unwrap(), DEFAULT_MODEL_ID);
429 let _ = std::fs::remove_dir_all(&home);
430 }
431
432 #[test]
433 fn videre_embed_model_env_var_has_no_effect() {
434 let home = cfg_home("noenv", "");
438 std::env::set_var("VIDERE_EMBED_MODEL", "owner/should-be-ignored");
439 let got = resolve_model_id_in(&home, None).unwrap();
440 std::env::remove_var("VIDERE_EMBED_MODEL");
441 assert_eq!(got, DEFAULT_MODEL_ID, "VIDERE_EMBED_MODEL must be ignored");
442 let _ = std::fs::remove_dir_all(&home);
443 }
444
445 #[test]
446 fn a_malformed_config_is_an_error_not_a_silent_default() {
447 let home = cfg_home("malformed", "not = = toml\n");
448 assert!(resolve_model_id_in(&home, None).is_err());
449 let _ = std::fs::remove_dir_all(&home);
450 }
451
452 #[test]
453 fn is_video_ext_matches_mov_and_mp4_case_insensitively() {
454 assert!(is_video_ext("mov"));
455 assert!(is_video_ext("MP4"));
456 assert!(is_video_ext("Mov"));
457 assert!(!is_video_ext("jpg"));
458 assert!(!is_video_ext(""));
459 }
460}