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