1use rusqlite::{params, Connection, Result};
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 let id = match explicit {
66 Some(id) => id.to_string(),
67 None => crate::home::load_config(home)?
68 .default_model
69 .unwrap_or_else(|| DEFAULT_MODEL_ID.to_string()),
70 };
71 validate_model_id(&id)?;
72 Ok(id)
73}
74
75pub fn validate_model_id(id: &str) -> anyhow::Result<()> {
93 match id.split_once('/') {
94 Some((owner, name)) if !owner.is_empty() && !name.is_empty() && !name.contains('/') => {
95 Ok(())
96 }
97 _ => anyhow::bail!(
98 "invalid model id {id:?}: expected owner/name, \
99 e.g. google/siglip-base-patch16-224"
100 ),
101 }
102}
103
104pub fn resolve_model_id(explicit: Option<&str>) -> anyhow::Result<String> {
110 resolve_model_id_in(&crate::home::videre_home()?, explicit)
111}
112
113#[derive(Debug, Clone)]
114pub struct PendingImage {
115 pub hash: String,
116 pub path: String,
117}
118
119pub fn ensure_embeddings_index(conn: &Connection) -> Result<()> {
127 conn.execute_batch("CREATE INDEX IF NOT EXISTS idx_file_hashes_hash ON file_hashes(hash);")
128}
129
130pub fn pending_images(conn: &Connection, model_id: &str) -> Result<Vec<PendingImage>> {
133 let mimes = crate::mime_probe::EMBEDDABLE_MIMES
134 .iter()
135 .map(|m| format!("'{m}'"))
136 .collect::<Vec<_>>()
137 .join(",");
138 let exts = EMBEDDABLE_EXTS
139 .iter()
140 .map(|e| format!("'{e}'"))
141 .collect::<Vec<_>>()
142 .join(",");
143 let sql = format!(
150 "SELECT hash, MIN(path) FROM file_hashes
151 WHERE lower(COALESCE(ext, '')) != 'dng'
152 AND (mime IN ({mimes}) OR (mime IS NULL AND lower(ext) IN ({exts})))
153 AND NOT EXISTS (SELECT 1 FROM emb.embeddings e
154 WHERE e.hash = file_hashes.hash AND e.model_id = ?1)
155 GROUP BY hash
156 ORDER BY hash"
157 );
158 let mut stmt = conn.prepare(&sql)?;
159 let rows = stmt.query_map(params![model_id], |row| {
160 Ok(PendingImage {
161 hash: row.get(0)?,
162 path: row.get(1)?,
163 })
164 })?;
165 rows.collect()
166}
167
168pub fn insert_embeddings(
170 conn: &Connection,
171 model_id: &str,
172 items: &[(String, Vec<u8>)],
173) -> Result<()> {
174 let tx = conn.unchecked_transaction()?;
175 {
176 let mut stmt = tx.prepare(
177 "INSERT OR REPLACE INTO emb.embeddings (hash, model_id, embedding, embedded_at)
178 VALUES (?1, ?2, ?3, datetime('now'))",
179 )?;
180 for (hash, blob) in items {
181 stmt.execute(params![hash, model_id, blob])?;
182 }
183 }
184 tx.commit()
185}
186
187pub fn load_embeddings(conn: &Connection, model_id: &str) -> Result<Vec<(String, Vec<u8>)>> {
191 let attached: bool = conn
192 .query_row(
193 "SELECT COUNT(*) FROM emb.sqlite_master WHERE type='table' AND name='embeddings'",
194 [],
195 |r| r.get::<_, i64>(0),
196 )
197 .unwrap_or(0)
198 > 0;
199 if !attached {
200 return Ok(Vec::new());
201 }
202 let mut stmt =
203 conn.prepare("SELECT hash, embedding FROM emb.embeddings WHERE model_id = ?1")?;
204 let rows = stmt.query_map(params![model_id], |row| Ok((row.get(0)?, row.get(1)?)))?;
205 rows.collect()
206}
207
208pub fn paths_for_hash(conn: &Connection, hash: &str) -> Result<Vec<String>> {
209 let mut stmt = conn.prepare("SELECT path FROM file_hashes WHERE hash = ?1 ORDER BY path")?;
210 let rows = stmt.query_map(params![hash], |row| row.get(0))?;
211 rows.collect()
212}
213
214#[cfg(test)]
215mod tests {
216 use super::*;
217 use rusqlite::Connection;
218
219 fn test_db_attached(tag: &str) -> Connection {
224 let lib = crate::embeddings_db::test_library(tag);
225 let conn = Connection::open_in_memory().unwrap();
226 conn.execute_batch(
227 "CREATE TABLE file_hashes (
228 path TEXT PRIMARY KEY,
229 hash TEXT NOT NULL,
230 mime TEXT,
231 size_bytes INTEGER,
232 created_at TEXT,
233 modified_at TEXT,
234 ext TEXT,
235 phash INTEGER,
236 exif_date TEXT,
237 gps_lat REAL,
238 gps_lon REAL,
239 width INTEGER,
240 height INTEGER
241 );",
242 )
243 .unwrap();
244 ensure_embeddings_index(&conn).unwrap();
245 crate::embeddings_db::attach(&conn, &lib, "test-model", true).unwrap();
246 conn
247 }
248
249 fn insert_file(conn: &Connection, path: &str, hash: &str, ext: &str) {
250 conn.execute(
251 "INSERT INTO file_hashes (path, hash, ext) VALUES (?1, ?2, ?3)",
252 rusqlite::params![path, hash, ext],
253 )
254 .unwrap();
255 }
256
257 #[test]
258 fn pending_images_dedupes_by_hash_and_includes_video() {
259 let conn = test_db_attached("emb_dedupe");
260 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
261 insert_file(&conn, "/b/1-copy.jpg", "h1", "jpg"); insert_file(&conn, "/a/2.png", "h2", "png");
263 insert_file(&conn, "/a/clip.mp4", "h3", "mp4"); insert_file(&conn, "/a/other.xyz", "h4", "xyz"); let pending = pending_images(&conn, "test-model").unwrap();
267 assert_eq!(pending.len(), 3); assert!(pending.iter().any(|p| p.hash == "h1"));
269 assert!(pending.iter().any(|p| p.hash == "h2"));
270 assert!(pending.iter().any(|p| p.hash == "h3"));
271 }
272
273 #[test]
274 fn pending_images_excludes_dng_since_it_cannot_be_decoded() {
275 let conn = test_db_attached("emb_dng");
276 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
277 insert_file(&conn, "/a/raw.dng", "h2", "dng");
278
279 let pending = pending_images(&conn, "test-model").unwrap();
280 assert_eq!(pending.len(), 1);
281 assert_eq!(pending[0].hash, "h1");
282 }
283
284 #[test]
285 fn pending_images_excludes_already_embedded() {
286 let conn = test_db_attached("emb_already");
287 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
288 insert_file(&conn, "/a/2.jpg", "h2", "jpg");
289 insert_embeddings(&conn, "test-model", &[("h1".to_string(), vec![0u8; 4])]).unwrap();
290
291 let pending = pending_images(&conn, "test-model").unwrap();
292 assert_eq!(pending.len(), 1);
293 assert_eq!(pending[0].hash, "h2");
294 }
295
296 #[test]
297 fn pending_images_is_model_aware() {
298 let conn = test_db_attached("emb_modelaware");
299 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
300 insert_embeddings(&conn, "a", &[("h1".to_string(), vec![0u8; 4])]).unwrap();
301
302 assert!(pending_images(&conn, "a").unwrap().is_empty());
304
305 let pending = pending_images(&conn, "b").unwrap();
307 assert_eq!(pending.len(), 1);
308 assert_eq!(pending[0].hash, "h1");
309 }
310
311 #[test]
312 fn pending_images_uses_mime_over_a_wrong_extension() {
313 let conn = test_db_attached("emb_mime");
314 conn.execute(
315 "INSERT INTO file_hashes (path, hash, ext, mime)
316 VALUES ('/a/actually_a_jpeg.png', 'h1', 'png', 'image/jpeg')",
317 [],
318 )
319 .unwrap();
320 let pending = pending_images(&conn, "m").unwrap();
321 assert_eq!(
322 pending.len(),
323 1,
324 "a JPEG named .png must still be embeddable"
325 );
326 }
327
328 #[test]
329 fn pending_images_falls_back_to_ext_when_mime_is_null() {
330 let conn = test_db_attached("emb_nullmime");
331 conn.execute(
332 "INSERT INTO file_hashes (path, hash, ext, mime) VALUES ('/a/1.jpg', 'h1', 'jpg', NULL)",
333 [],
334 )
335 .unwrap();
336 assert_eq!(pending_images(&conn, "m").unwrap().len(), 1);
337 }
338
339 #[test]
340 fn pending_images_still_excludes_dng_even_though_its_mime_is_tiff() {
341 let conn = test_db_attached("emb_dng_mime");
344 conn.execute(
345 "INSERT INTO file_hashes (path, hash, ext, mime)
346 VALUES ('/a/raw.dng', 'h1', 'dng', 'image/tiff')",
347 [],
348 )
349 .unwrap();
350 assert!(pending_images(&conn, "m").unwrap().is_empty());
351 }
352
353 #[test]
354 fn insert_embeddings_empty_slice_succeeds() {
355 let conn = test_db_attached("emb_empty");
356 insert_embeddings(&conn, "test-model", &[]).unwrap();
357 assert!(load_embeddings(&conn, "test-model").unwrap().is_empty());
358 }
359
360 #[test]
361 fn insert_and_load_round_trip() {
362 let conn = test_db_attached("emb_roundtrip");
363 insert_embeddings(
364 &conn,
365 "test-model",
366 &[
367 ("h1".to_string(), vec![1u8, 2, 3, 4]),
368 ("h2".to_string(), vec![5u8, 6]),
369 ],
370 )
371 .unwrap();
372
373 let rows = load_embeddings(&conn, "test-model").unwrap();
374 assert_eq!(rows.len(), 2);
375 let h1 = rows.iter().find(|(h, _)| h == "h1").unwrap();
376 assert_eq!(h1.1, vec![1u8, 2, 3, 4]);
377
378 assert!(load_embeddings(&conn, "other").unwrap().is_empty());
380 }
381
382 #[test]
383 fn load_embeddings_finds_the_table_in_the_attached_database() {
384 let conn = test_db_attached("emb_attachedprobe");
389 insert_embeddings(&conn, "test-model", &[("h1".to_string(), vec![1u8, 2])]).unwrap();
390
391 let rows = load_embeddings(&conn, "test-model").unwrap();
392 assert_eq!(rows.len(), 1, "must read through emb., not main");
393 }
394
395 #[test]
396 fn ensure_embeddings_index_creates_the_index_in_the_main_database() {
397 let conn = test_db_attached("emb_indexmain");
400 let found: i64 = conn
401 .query_row(
402 "SELECT COUNT(*) FROM main.sqlite_master
403 WHERE type='index' AND name='idx_file_hashes_hash'",
404 [],
405 |r| r.get(0),
406 )
407 .unwrap();
408 assert_eq!(found, 1);
409 }
410
411 #[test]
412 fn paths_for_hash_returns_all_duplicates() {
413 let conn = test_db_attached("emb_paths");
414 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
415 insert_file(&conn, "/b/1-copy.jpg", "h1", "jpg");
416 let paths = paths_for_hash(&conn, "h1").unwrap();
417 assert_eq!(paths.len(), 2);
418 }
419
420 #[test]
421 fn default_model_id_is_a_siglip_checkpoint() {
422 assert!(
429 DEFAULT_MODEL_ID.starts_with("google/"),
430 "{DEFAULT_MODEL_ID}"
431 );
432 assert!(DEFAULT_MODEL_ID.contains("siglip"), "{DEFAULT_MODEL_ID}");
433 assert_eq!(
434 DEFAULT_MODEL_ID.matches('/').count(),
435 1,
436 "{DEFAULT_MODEL_ID}"
437 );
438 }
439
440 fn cfg_home(tag: &str, toml_text: &str) -> std::path::PathBuf {
441 let dir = std::env::temp_dir().join(format!("videre_rmi_{}_{}", tag, std::process::id()));
442 let _ = std::fs::remove_dir_all(&dir);
443 std::fs::create_dir_all(&dir).unwrap();
444 if !toml_text.is_empty() {
445 std::fs::write(dir.join("config.toml"), toml_text).unwrap();
446 }
447 dir
448 }
449
450 #[test]
451 fn resolve_model_id_prefers_the_explicit_argument() {
452 let home = cfg_home("explicit", "default_model = \"owner/from-config\"\n");
453 assert_eq!(
454 resolve_model_id_in(&home, Some("owner/explicit")).unwrap(),
455 "owner/explicit"
456 );
457 let _ = std::fs::remove_dir_all(&home);
458 }
459
460 #[test]
461 fn resolve_model_id_uses_config_when_there_is_no_flag() {
462 let home = cfg_home("fromconfig", "default_model = \"owner/from-config\"\n");
463 assert_eq!(
464 resolve_model_id_in(&home, None).unwrap(),
465 "owner/from-config"
466 );
467 let _ = std::fs::remove_dir_all(&home);
468 }
469
470 #[test]
471 fn resolve_model_id_falls_back_to_the_builtin_default() {
472 let home = cfg_home("builtin", "");
473 assert_eq!(resolve_model_id_in(&home, None).unwrap(), DEFAULT_MODEL_ID);
474 let _ = std::fs::remove_dir_all(&home);
475 }
476
477 #[test]
478 fn videre_embed_model_env_var_has_no_effect() {
479 let home = cfg_home("noenv", "");
483 std::env::set_var("VIDERE_EMBED_MODEL", "owner/should-be-ignored");
484 let got = resolve_model_id_in(&home, None).unwrap();
485 std::env::remove_var("VIDERE_EMBED_MODEL");
486 assert_eq!(got, DEFAULT_MODEL_ID, "VIDERE_EMBED_MODEL must be ignored");
487 let _ = std::fs::remove_dir_all(&home);
488 }
489
490 #[test]
491 fn a_malformed_config_is_an_error_not_a_silent_default() {
492 let home = cfg_home("malformed", "not = = toml\n");
493 assert!(resolve_model_id_in(&home, None).is_err());
494 let _ = std::fs::remove_dir_all(&home);
495 }
496
497 #[test]
498 fn is_video_ext_matches_mov_and_mp4_case_insensitively() {
499 assert!(is_video_ext("mov"));
500 assert!(is_video_ext("MP4"));
501 assert!(is_video_ext("Mov"));
502 assert!(!is_video_ext("jpg"));
503 assert!(!is_video_ext(""));
504 }
505}
506
507#[cfg(test)]
508mod model_id_tests {
509 use super::*;
510
511 #[test]
512 fn an_id_without_a_slash_is_rejected() {
513 assert!(validate_model_id("foo").is_err());
516 assert!(validate_model_id("").is_err());
517 assert!(validate_model_id("/name").is_err());
518 assert!(validate_model_id("owner/").is_err());
519 assert!(validate_model_id("a/b/c").is_err());
520 }
521
522 #[test]
523 fn a_well_formed_id_passes_even_if_never_embedded() {
524 assert!(validate_model_id("google/siglip-base-patch16-224").is_ok());
525 assert!(validate_model_id("someone/a-model-nobody-has-run").is_ok());
526 }
527
528 #[test]
529 fn the_explicit_flag_is_validated_not_just_the_config_file() {
530 let dir = tempfile::tempdir().unwrap();
533 assert!(resolve_model_id_in(dir.path(), Some("foo")).is_err());
534 assert_eq!(
535 resolve_model_id_in(dir.path(), Some("owner/name")).unwrap(),
536 "owner/name"
537 );
538 }
539}