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_from(
51 config: &crate::library_config::LibraryConfig,
52 explicit: Option<&str>,
53) -> anyhow::Result<String> {
54 let id = explicit.unwrap_or(&config.default_model).to_string();
55 validate_model_id(&id)?;
56 Ok(id)
57}
58
59pub fn validate_model_id(id: &str) -> anyhow::Result<()> {
77 match id.split_once('/') {
78 Some((owner, name)) if !owner.is_empty() && !name.is_empty() && !name.contains('/') => {
79 Ok(())
80 }
81 _ => anyhow::bail!(
82 "invalid model id {id:?}: expected owner/name, \
83 e.g. google/siglip-base-patch16-224"
84 ),
85 }
86}
87
88#[derive(Debug, Clone)]
89pub struct PendingImage {
90 pub hash: String,
91 pub path: String,
92}
93
94pub fn ensure_embeddings_index(conn: &Connection) -> Result<()> {
102 conn.execute_batch("CREATE INDEX IF NOT EXISTS idx_file_hashes_hash ON file_hashes(hash);")
103}
104
105pub fn pending_images(conn: &Connection, model_id: &str) -> Result<Vec<PendingImage>> {
108 images_for_model(conn, model_id, false)
109}
110
111pub fn embeddable_images(conn: &Connection, model_id: &str) -> Result<Vec<PendingImage>> {
116 images_for_model(conn, model_id, true)
117}
118
119fn images_for_model(
120 conn: &Connection,
121 model_id: &str,
122 include_embedded: bool,
123) -> Result<Vec<PendingImage>> {
124 let mimes = crate::mime_probe::EMBEDDABLE_MIMES
125 .iter()
126 .map(|m| format!("'{m}'"))
127 .collect::<Vec<_>>()
128 .join(",");
129 let exts = EMBEDDABLE_EXTS
130 .iter()
131 .map(|e| format!("'{e}'"))
132 .collect::<Vec<_>>()
133 .join(",");
134 let skip_embedded = if include_embedded {
141 String::new()
142 } else {
143 "AND NOT EXISTS (SELECT 1 FROM emb.embeddings e
144 WHERE e.hash = file_hashes.hash AND e.model_id = ?1)"
145 .to_string()
146 };
147 let sql = format!(
148 "SELECT hash, MIN(path) FROM file_hashes
149 WHERE lower(COALESCE(ext, '')) != 'dng'
150 AND (mime IN ({mimes}) OR (mime IS NULL AND lower(ext) IN ({exts})))
151 {skip_embedded}
152 GROUP BY hash
153 ORDER BY hash"
154 );
155 let mut stmt = conn.prepare(&sql)?;
156 let map_row = |row: &rusqlite::Row| {
157 Ok(PendingImage {
158 hash: row.get(0)?,
159 path: row.get(1)?,
160 })
161 };
162 let rows = if include_embedded {
165 stmt.query_map([], map_row)?
166 } else {
167 stmt.query_map(params![model_id], map_row)?
168 };
169 rows.collect()
170}
171
172pub fn insert_embeddings(
174 conn: &Connection,
175 model_id: &str,
176 items: &[(String, Vec<u8>)],
177) -> Result<()> {
178 let tx = conn.unchecked_transaction()?;
179 {
180 let mut stmt = tx.prepare(
181 "INSERT OR REPLACE INTO emb.embeddings (hash, model_id, embedding, embedded_at)
182 VALUES (?1, ?2, ?3, datetime('now'))",
183 )?;
184 for (hash, blob) in items {
185 stmt.execute(params![hash, model_id, blob])?;
186 }
187 }
188 tx.commit()
189}
190
191pub fn load_embeddings(conn: &Connection, model_id: &str) -> Result<Vec<(String, Vec<u8>)>> {
195 let attached: bool = conn
196 .query_row(
197 "SELECT COUNT(*) FROM emb.sqlite_master WHERE type='table' AND name='embeddings'",
198 [],
199 |r| r.get::<_, i64>(0),
200 )
201 .unwrap_or(0)
202 > 0;
203 if !attached {
204 return Ok(Vec::new());
205 }
206 let mut stmt =
207 conn.prepare("SELECT hash, embedding FROM emb.embeddings WHERE model_id = ?1")?;
208 let rows = stmt.query_map(params![model_id], |row| Ok((row.get(0)?, row.get(1)?)))?;
209 rows.collect()
210}
211
212pub fn paths_for_hash(conn: &Connection, hash: &str) -> Result<Vec<String>> {
213 let mut stmt = conn.prepare("SELECT path FROM file_hashes WHERE hash = ?1 ORDER BY path")?;
214 let rows = stmt.query_map(params![hash], |row| row.get(0))?;
215 rows.collect()
216}
217
218#[cfg(test)]
219mod tests {
220 use super::*;
221 use rusqlite::Connection;
222
223 #[test]
224 fn library_config_model_is_used_unless_an_explicit_model_wins() {
225 let mut config = crate::library_config::LibraryConfig::default();
226 config.default_model = "owner/configured".into();
227 assert_eq!(
228 resolve_model_id_from(&config, None).unwrap(),
229 "owner/configured"
230 );
231 assert_eq!(
232 resolve_model_id_from(&config, Some("owner/explicit")).unwrap(),
233 "owner/explicit"
234 );
235 assert!(resolve_model_id_from(&config, Some("invalid")).is_err());
236 }
237
238 fn test_db_attached(tag: &str) -> Connection {
243 let ctx = crate::embeddings_db::test_context(tag);
244 let conn = Connection::open_in_memory().unwrap();
245 conn.execute_batch(
246 "CREATE TABLE file_hashes (
247 path TEXT PRIMARY KEY,
248 hash TEXT NOT NULL,
249 mime TEXT,
250 size_bytes INTEGER,
251 created_at TEXT,
252 modified_at TEXT,
253 ext TEXT,
254 phash INTEGER,
255 exif_date TEXT,
256 gps_lat REAL,
257 gps_lon REAL,
258 width INTEGER,
259 height INTEGER
260 );",
261 )
262 .unwrap();
263 ensure_embeddings_index(&conn).unwrap();
264 crate::embeddings_db::attach_in(&conn, &ctx, "owner/test-model", true).unwrap();
265 conn
266 }
267
268 fn insert_file(conn: &Connection, path: &str, hash: &str, ext: &str) {
269 conn.execute(
270 "INSERT INTO file_hashes (path, hash, ext) VALUES (?1, ?2, ?3)",
271 rusqlite::params![path, hash, ext],
272 )
273 .unwrap();
274 }
275
276 #[test]
277 fn pending_images_dedupes_by_hash_and_includes_video() {
278 let conn = test_db_attached("emb_dedupe");
279 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
280 insert_file(&conn, "/b/1-copy.jpg", "h1", "jpg"); insert_file(&conn, "/a/2.png", "h2", "png");
282 insert_file(&conn, "/a/clip.mp4", "h3", "mp4"); insert_file(&conn, "/a/other.xyz", "h4", "xyz"); let pending = pending_images(&conn, "test-model").unwrap();
286 assert_eq!(pending.len(), 3); assert!(pending.iter().any(|p| p.hash == "h1"));
288 assert!(pending.iter().any(|p| p.hash == "h2"));
289 assert!(pending.iter().any(|p| p.hash == "h3"));
290 }
291
292 #[test]
293 fn pending_images_excludes_dng_since_it_cannot_be_decoded() {
294 let conn = test_db_attached("emb_dng");
295 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
296 insert_file(&conn, "/a/raw.dng", "h2", "dng");
297
298 let pending = pending_images(&conn, "test-model").unwrap();
299 assert_eq!(pending.len(), 1);
300 assert_eq!(pending[0].hash, "h1");
301 }
302
303 #[test]
304 fn pending_images_excludes_already_embedded() {
305 let conn = test_db_attached("emb_already");
306 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
307 insert_file(&conn, "/a/2.jpg", "h2", "jpg");
308 insert_embeddings(&conn, "test-model", &[("h1".to_string(), vec![0u8; 4])]).unwrap();
309
310 let pending = pending_images(&conn, "test-model").unwrap();
311 assert_eq!(pending.len(), 1);
312 assert_eq!(pending[0].hash, "h2");
313 }
314
315 #[test]
316 fn pending_images_is_model_aware() {
317 let conn = test_db_attached("emb_modelaware");
318 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
319 insert_embeddings(&conn, "a", &[("h1".to_string(), vec![0u8; 4])]).unwrap();
320
321 assert!(pending_images(&conn, "a").unwrap().is_empty());
323
324 let pending = pending_images(&conn, "b").unwrap();
326 assert_eq!(pending.len(), 1);
327 assert_eq!(pending[0].hash, "h1");
328 }
329
330 #[test]
331 fn pending_images_uses_mime_over_a_wrong_extension() {
332 let conn = test_db_attached("emb_mime");
333 conn.execute(
334 "INSERT INTO file_hashes (path, hash, ext, mime)
335 VALUES ('/a/actually_a_jpeg.png', 'h1', 'png', 'image/jpeg')",
336 [],
337 )
338 .unwrap();
339 let pending = pending_images(&conn, "m").unwrap();
340 assert_eq!(
341 pending.len(),
342 1,
343 "a JPEG named .png must still be embeddable"
344 );
345 }
346
347 #[test]
348 fn pending_images_falls_back_to_ext_when_mime_is_null() {
349 let conn = test_db_attached("emb_nullmime");
350 conn.execute(
351 "INSERT INTO file_hashes (path, hash, ext, mime) VALUES ('/a/1.jpg', 'h1', 'jpg', NULL)",
352 [],
353 )
354 .unwrap();
355 assert_eq!(pending_images(&conn, "m").unwrap().len(), 1);
356 }
357
358 #[test]
359 fn pending_images_still_excludes_dng_even_though_its_mime_is_tiff() {
360 let conn = test_db_attached("emb_dng_mime");
363 conn.execute(
364 "INSERT INTO file_hashes (path, hash, ext, mime)
365 VALUES ('/a/raw.dng', 'h1', 'dng', 'image/tiff')",
366 [],
367 )
368 .unwrap();
369 assert!(pending_images(&conn, "m").unwrap().is_empty());
370 }
371
372 #[test]
373 fn embeddable_images_includes_already_embedded_rows() {
374 let conn = test_db_attached("emb_reprocess");
378 conn.execute(
379 "INSERT INTO file_hashes (path, hash, ext, mime)
380 VALUES ('/a/1.jpg', 'h1', 'jpg', 'image/jpeg')",
381 [],
382 )
383 .unwrap();
384 insert_embeddings(&conn, "m", &[("h1".to_string(), vec![0u8; 4])]).unwrap();
385 assert!(
386 pending_images(&conn, "m").unwrap().is_empty(),
387 "a fully embedded library has no pending work"
388 );
389 let all = embeddable_images(&conn, "m").unwrap();
390 assert_eq!(all.len(), 1, "reprocess must see the embedded hash again");
391 assert_eq!(all[0].hash, "h1");
392 }
393
394 #[test]
395 fn insert_embeddings_empty_slice_succeeds() {
396 let conn = test_db_attached("emb_empty");
397 insert_embeddings(&conn, "test-model", &[]).unwrap();
398 assert!(load_embeddings(&conn, "test-model").unwrap().is_empty());
399 }
400
401 #[test]
402 fn insert_and_load_round_trip() {
403 let conn = test_db_attached("emb_roundtrip");
404 insert_embeddings(
405 &conn,
406 "test-model",
407 &[
408 ("h1".to_string(), vec![1u8, 2, 3, 4]),
409 ("h2".to_string(), vec![5u8, 6]),
410 ],
411 )
412 .unwrap();
413
414 let rows = load_embeddings(&conn, "test-model").unwrap();
415 assert_eq!(rows.len(), 2);
416 let h1 = rows.iter().find(|(h, _)| h == "h1").unwrap();
417 assert_eq!(h1.1, vec![1u8, 2, 3, 4]);
418
419 assert!(load_embeddings(&conn, "other").unwrap().is_empty());
421 }
422
423 #[test]
424 fn load_embeddings_finds_the_table_in_the_attached_database() {
425 let conn = test_db_attached("emb_attachedprobe");
430 insert_embeddings(&conn, "test-model", &[("h1".to_string(), vec![1u8, 2])]).unwrap();
431
432 let rows = load_embeddings(&conn, "test-model").unwrap();
433 assert_eq!(rows.len(), 1, "must read through emb., not main");
434 }
435
436 #[test]
437 fn ensure_embeddings_index_creates_the_index_in_the_main_database() {
438 let conn = test_db_attached("emb_indexmain");
441 let found: i64 = conn
442 .query_row(
443 "SELECT COUNT(*) FROM main.sqlite_master
444 WHERE type='index' AND name='idx_file_hashes_hash'",
445 [],
446 |r| r.get(0),
447 )
448 .unwrap();
449 assert_eq!(found, 1);
450 }
451
452 #[test]
453 fn paths_for_hash_returns_all_duplicates() {
454 let conn = test_db_attached("emb_paths");
455 insert_file(&conn, "/a/1.jpg", "h1", "jpg");
456 insert_file(&conn, "/b/1-copy.jpg", "h1", "jpg");
457 let paths = paths_for_hash(&conn, "h1").unwrap();
458 assert_eq!(paths.len(), 2);
459 }
460
461 #[test]
462 fn default_model_id_is_a_siglip_checkpoint() {
463 assert!(
470 DEFAULT_MODEL_ID.starts_with("google/"),
471 "{DEFAULT_MODEL_ID}"
472 );
473 assert!(DEFAULT_MODEL_ID.contains("siglip"), "{DEFAULT_MODEL_ID}");
474 assert_eq!(
475 DEFAULT_MODEL_ID.matches('/').count(),
476 1,
477 "{DEFAULT_MODEL_ID}"
478 );
479 }
480
481 #[test]
482 fn is_video_ext_matches_mov_and_mp4_case_insensitively() {
483 assert!(is_video_ext("mov"));
484 assert!(is_video_ext("MP4"));
485 assert!(is_video_ext("Mov"));
486 assert!(!is_video_ext("jpg"));
487 assert!(!is_video_ext(""));
488 }
489}
490
491#[cfg(test)]
492mod model_id_tests {
493 use super::*;
494
495 #[test]
496 fn an_id_without_a_slash_is_rejected() {
497 assert!(validate_model_id("foo").is_err());
500 assert!(validate_model_id("").is_err());
501 assert!(validate_model_id("/name").is_err());
502 assert!(validate_model_id("owner/").is_err());
503 assert!(validate_model_id("a/b/c").is_err());
504 }
505
506 #[test]
507 fn a_well_formed_id_passes_even_if_never_embedded() {
508 assert!(validate_model_id("google/siglip-base-patch16-224").is_ok());
509 assert!(validate_model_id("someone/a-model-nobody-has-run").is_ok());
510 }
511
512 #[test]
513 fn the_explicit_flag_is_validated_not_just_the_config_file() {
514 let config = crate::library_config::LibraryConfig::default();
517 assert!(resolve_model_id_from(&config, Some("foo")).is_err());
518 assert_eq!(
519 resolve_model_id_from(&config, Some("owner/name")).unwrap(),
520 "owner/name"
521 );
522 }
523}