1use anyhow::{Context, Result};
12use rusqlite::Connection;
13use std::path::{Path, PathBuf};
14
15pub const ATTACH_ALIAS: &str = "emb";
17
18const DB_EXT: &str = "db";
20
21pub fn model_slug(model_id: &str) -> String {
27 model_id.replace('/', "--")
28}
29
30pub fn model_from_slug(slug: &str) -> String {
33 slug.replacen("--", "/", 1)
34}
35
36fn validate_context_storage(ctx: &crate::library::LibraryContext) -> Result<()> {
37 ctx.ensure_root_identity()?;
38 crate::library_locks::reject_dir_redirect(&ctx.paths.embeddings, "the embeddings directory")?;
39 Ok(())
40}
41
42pub fn db_path_in(ctx: &crate::library::LibraryContext, model_id: &str) -> Result<PathBuf> {
47 crate::embeddings::validate_model_id(model_id)?;
48 validate_context_storage(ctx)?;
49 let path = ctx
50 .paths
51 .embeddings
52 .join(format!("{}.{DB_EXT}", model_slug(model_id)));
53 crate::library_locks::reject_redirect(&path, "the model database")?;
54 Ok(path)
55}
56
57pub const PAGE_SIZE: i64 = 16384;
73
74fn init_model_db(path: &Path) -> Result<()> {
81 if let Some(parent) = path.parent() {
82 std::fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
83 }
84 let conn = Connection::open(path).with_context(|| format!("create {}", path.display()))?;
85 conn.pragma_update(None, "page_size", PAGE_SIZE)
86 .context("set page_size")?;
87 conn.pragma_update(None, "journal_mode", "WAL")
88 .context("set journal_mode")?;
89 conn.execute_batch(
90 "CREATE TABLE IF NOT EXISTS embeddings (
91 hash TEXT PRIMARY KEY NOT NULL,
92 model_id TEXT NOT NULL,
93 embedding BLOB NOT NULL,
94 embedded_at TEXT NOT NULL
95 );",
96 )
97 .with_context(|| format!("create embeddings table in {}", path.display()))?;
98 Ok(())
99}
100
101pub fn attach_in(
103 conn: &Connection,
104 ctx: &crate::library::LibraryContext,
105 model_id: &str,
106 create: bool,
107) -> Result<()> {
108 crate::library_locks::verify_state(ctx)?;
109 let path = db_path_in(ctx, model_id)?;
110 if !path.exists() {
111 if !create {
112 let available = list_models_in(ctx).unwrap_or_default();
113 let available = if available.is_empty() {
114 "(none)".to_string()
115 } else {
116 available.join(", ")
117 };
118 anyhow::bail!(
119 "no embeddings for {model_id} in this library\n \
120 expected: {}\n available: {available}\n \
121 run: videre embed --model {model_id}",
122 path.display()
123 );
124 }
125 std::fs::create_dir_all(&ctx.paths.embeddings)
126 .with_context(|| format!("create {}", ctx.paths.embeddings.display()))?;
127 validate_context_storage(ctx)?;
128 init_model_db(&path)?;
129 }
130 crate::library_locks::reject_redirect(&path, "the model database")?;
131 let meta =
132 std::fs::symlink_metadata(&path).with_context(|| format!("inspect {}", path.display()))?;
133 anyhow::ensure!(meta.is_file(), "{} is not a regular file", path.display());
134 conn.execute(
135 &format!("ATTACH DATABASE ?1 AS {ATTACH_ALIAS}"),
136 [path.to_string_lossy().as_ref()],
137 )
138 .with_context(|| format!("attach {}", path.display()))?;
139 if create {
140 conn.execute_batch(
141 "CREATE TABLE IF NOT EXISTS emb.embeddings (
142 hash TEXT PRIMARY KEY NOT NULL,
143 model_id TEXT NOT NULL,
144 embedding BLOB NOT NULL,
145 embedded_at TEXT NOT NULL
146 );",
147 )
148 .with_context(|| format!("ensure schema in {}", path.display()))?;
149 }
150 Ok(())
151}
152
153pub fn list_models_in(ctx: &crate::library::LibraryContext) -> Result<Vec<String>> {
155 validate_context_storage(ctx)?;
156 let entries = match std::fs::read_dir(&ctx.paths.embeddings) {
157 Ok(entries) => entries,
158 Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
159 Err(error) => {
160 return Err(error).with_context(|| format!("read {}", ctx.paths.embeddings.display()))
161 }
162 };
163 let mut models = Vec::new();
164 for entry in entries {
165 let entry = entry.with_context(|| format!("read {}", ctx.paths.embeddings.display()))?;
166 let path = entry.path();
167 if path.extension().is_none_or(|extension| extension != DB_EXT) {
168 continue;
169 }
170 crate::library_locks::reject_redirect(&path, "the model database")?;
171 let meta = std::fs::symlink_metadata(&path)
172 .with_context(|| format!("inspect {}", path.display()))?;
173 anyhow::ensure!(meta.is_file(), "{} is not a regular file", path.display());
174 if let Some(stem) = path.file_stem() {
175 models.push(model_from_slug(&stem.to_string_lossy()));
176 }
177 }
178 models.sort();
179 Ok(models)
180}
181
182#[derive(Debug, Clone, PartialEq, serde::Serialize)]
184pub struct ModelEmbeddingCount {
185 pub model_id: String,
186 pub count: i64,
187 pub dims: i64,
191 pub size_bytes: i64,
192}
193
194pub fn counts_by_model_in(
196 ctx: &crate::library::LibraryContext,
197) -> Result<Vec<ModelEmbeddingCount>> {
198 let mut out = Vec::new();
199 for model_id in list_models_in(ctx)? {
200 let path = db_path_in(ctx, &model_id)?;
201 let size_bytes = std::fs::metadata(&path)
202 .with_context(|| format!("inspect {}", path.display()))?
203 .len() as i64;
204 let conn = Connection::open(&path).with_context(|| format!("open {}", path.display()))?;
205 let count: i64 = conn
206 .query_row("SELECT COUNT(*) FROM embeddings", [], |row| row.get(0))
207 .unwrap_or(0);
208 let dims: i64 = conn
209 .query_row(
210 "SELECT LENGTH(embedding) / 2 FROM embeddings LIMIT 1",
211 [],
212 |row| row.get(0),
213 )
214 .unwrap_or(0);
215 out.push(ModelEmbeddingCount {
216 model_id,
217 count,
218 dims,
219 size_bytes,
220 });
221 }
222 Ok(out)
223}
224
225pub fn attach_for_read_in(
232 conn: &Connection,
233 ctx: &crate::library::LibraryContext,
234 model_id: &str,
235) -> Result<()> {
236 attach_in(conn, ctx, model_id, false)
237}
238
239pub fn attach_for_read_or_placeholder_in(
247 conn: &Connection,
248 ctx: &crate::library::LibraryContext,
249 model_id: &str,
250) -> Result<()> {
251 if attach_for_read_in(conn, ctx, model_id).is_ok() {
252 return Ok(());
253 }
254 conn.execute("ATTACH DATABASE ':memory:' AS emb", [])
255 .context("attach placeholder embeddings database")?;
256 conn.execute_batch(
257 "CREATE TABLE emb.embeddings (
258 hash TEXT PRIMARY KEY NOT NULL,
259 model_id TEXT NOT NULL,
260 embedding BLOB NOT NULL,
261 embedded_at TEXT NOT NULL
262 );",
263 )
264 .context("create placeholder embeddings table")?;
265 Ok(())
266}
267
268pub fn detach(conn: &Connection) -> Result<()> {
271 conn.execute(&format!("DETACH DATABASE {ATTACH_ALIAS}"), [])
272 .context("detach embeddings database")?;
273 Ok(())
274}
275
276#[cfg(test)]
285pub(crate) fn test_context(tag: &str) -> crate::library::LibraryContext {
286 let base = std::env::temp_dir().join(format!("videre-embdb-{}-{}", tag, std::process::id()));
287 let _ = std::fs::remove_dir_all(&base);
288 let root = base.join("lib");
289 std::fs::create_dir_all(&root).unwrap();
290 let ctx = crate::library::LibraryContext::new(&root, &base.join("cache")).unwrap();
291 std::fs::create_dir_all(&ctx.paths.state).unwrap();
294 ctx
295}
296
297#[cfg(test)]
298mod tests {
299 use super::*;
300
301 fn explicit_context(parent: &Path, name: &str) -> crate::library::LibraryContext {
302 let root = parent.join(name);
303 std::fs::create_dir(&root).unwrap();
304 crate::library::LibraryContext::new(&root, &parent.join("cache")).unwrap()
305 }
306
307 #[test]
308 fn explicit_model_stores_are_isolated_and_keep_f16_dimensions() {
309 let temp = tempfile::tempdir().unwrap();
310 let a = explicit_context(temp.path(), "a");
311 let b = explicit_context(temp.path(), "b");
312 let a_conn = crate::library_db::initialize(&a).unwrap();
313 let b_conn = crate::library_db::initialize(&b).unwrap();
314 for model in ["owner/model-a", "owner/model-b"] {
315 attach_in(&a_conn, &a, model, true).unwrap();
316 a_conn
317 .execute(
318 "INSERT INTO emb.embeddings (hash, model_id, embedding, embedded_at)
319 VALUES ('shared', ?1, zeroblob(1536), 'now')",
320 [model],
321 )
322 .unwrap();
323 detach(&a_conn).unwrap();
324 }
325 attach_in(&b_conn, &b, "owner/model-a", true).unwrap();
326 b_conn
327 .execute(
328 "INSERT INTO emb.embeddings (hash, model_id, embedding, embedded_at)
329 VALUES ('shared', ?1, zeroblob(8), 'now')",
330 ["owner/model-a"],
331 )
332 .unwrap();
333 detach(&b_conn).unwrap();
334
335 assert_eq!(
336 list_models_in(&a).unwrap(),
337 vec!["owner/model-a", "owner/model-b"]
338 );
339 let a_counts = counts_by_model_in(&a).unwrap();
340 assert_eq!(a_counts.len(), 2);
341 assert!(a_counts
342 .iter()
343 .all(|count| count.count == 1 && count.dims == 768));
344 let b_counts = counts_by_model_in(&b).unwrap();
345 assert_eq!(b_counts[0].count, 1);
346 assert_eq!(b_counts[0].dims, 4);
347 assert_ne!(
348 db_path_in(&a, "owner/model-a").unwrap(),
349 db_path_in(&b, "owner/model-a").unwrap()
350 );
351 }
352
353 #[test]
354 fn explicit_read_attachment_does_not_create_a_missing_store() {
355 let temp = tempfile::tempdir().unwrap();
356 let ctx = explicit_context(temp.path(), "library");
357 let conn = crate::library_db::initialize(&ctx).unwrap();
358 let expected = db_path_in(&ctx, "owner/missing").unwrap();
359 let error = attach_for_read_in(&conn, &ctx, "owner/missing").unwrap_err();
360 assert!(format!("{error:#}").contains("no embeddings for owner/missing"));
361 assert!(!expected.exists());
362 assert!(!ctx.paths.embeddings.exists());
363 }
364
365 #[test]
368 fn model_slug_replaces_the_owner_separator() {
369 assert_eq!(
370 model_slug("google/siglip2-base-patch16-384"),
371 "google--siglip2-base-patch16-384"
372 );
373 }
374
375 #[test]
376 fn model_slug_round_trips_through_model_from_slug() {
377 for id in [
378 "google/siglip2-base-patch16-384",
379 "google/siglip-so400m-patch14-384",
380 "google/siglip-base-patch16-224",
381 ] {
382 assert_eq!(model_from_slug(&model_slug(id)), id);
383 }
384 }
385
386 #[test]
387 fn model_slug_contains_no_path_separator() {
388 assert!(!model_slug("google/siglip2-base-patch16-384").contains('/'));
391 }
392
393 #[test]
394 fn db_path_in_joins_the_embeddings_dir_and_model_slug() {
395 let ctx = test_context("dbpath");
396 let p = db_path_in(&ctx, "google/siglip2-base-patch16-384").unwrap();
397 assert_eq!(
398 p.file_name().unwrap(),
399 "google--siglip2-base-patch16-384.db"
400 );
401 assert_eq!(p.parent().unwrap(), ctx.paths.embeddings);
402 }
403
404 #[test]
405 fn attach_with_create_makes_a_database_with_the_chosen_page_size() {
406 let ctx = test_context("create");
407 let conn = Connection::open_in_memory().unwrap();
408 attach_in(&conn, &ctx, "google/siglip2-base-patch16-384", true).unwrap();
409
410 let ps: i64 = conn
413 .query_row("PRAGMA emb.page_size", [], |r| r.get(0))
414 .unwrap();
415 assert_eq!(ps, PAGE_SIZE);
416 }
417
418 #[test]
419 fn attach_with_create_is_idempotent_and_preserves_rows() {
420 let ctx = test_context("idem");
421 let model = "google/siglip2-base-patch16-384";
422
423 let c1 = Connection::open_in_memory().unwrap();
424 attach_in(&c1, &ctx, model, true).unwrap();
425 c1.execute(
426 "INSERT INTO emb.embeddings (hash, model_id, embedding, embedded_at)
427 VALUES ('h1', ?1, X'0102', '2026-08-05T00:00:00')",
428 [model],
429 )
430 .unwrap();
431 detach(&c1).unwrap();
432 drop(c1);
433
434 let c2 = Connection::open_in_memory().unwrap();
435 attach_in(&c2, &ctx, model, true).unwrap();
436 let n: i64 = c2
437 .query_row("SELECT COUNT(*) FROM emb.embeddings", [], |r| r.get(0))
438 .unwrap();
439 assert_eq!(n, 1, "re-attaching must not clobber existing rows");
440 }
441
442 #[test]
443 fn attach_with_create_repairs_a_file_that_exists_without_the_table() {
444 let ctx = test_context("repair");
449 let model = "google/siglip2-base-patch16-384";
450 let path = db_path_in(&ctx, model).unwrap();
451 std::fs::create_dir_all(path.parent().unwrap()).unwrap();
452 std::fs::write(&path, b"").unwrap(); let conn = Connection::open_in_memory().unwrap();
455 attach_in(&conn, &ctx, model, true).unwrap();
456 let n: i64 = conn
457 .query_row("SELECT COUNT(*) FROM emb.embeddings", [], |r| r.get(0))
458 .expect("the table must exist after attach_in(create: true)");
459 assert_eq!(n, 0);
460 }
461
462 #[test]
463 fn attach_without_create_errors_and_names_available_models() {
464 let ctx = test_context("missing");
465 let conn = Connection::open_in_memory().unwrap();
466 attach_in(&conn, &ctx, "google/siglip2-base-patch16-384", true).unwrap();
467 detach(&conn).unwrap();
468
469 let err = attach_in(&conn, &ctx, "google/siglip-base-patch16-224", false).unwrap_err();
470 let msg = format!("{err:#}");
471 assert!(
472 msg.contains("no embeddings for google/siglip-base-patch16-224"),
473 "{msg}"
474 );
475 assert!(
476 msg.contains("google/siglip2-base-patch16-384"),
477 "error must list what IS available: {msg}"
478 );
479 assert!(msg.contains("videre embed --model"), "{msg}");
480 }
481
482 #[test]
483 fn two_models_do_not_see_each_others_rows() {
484 let ctx = test_context("isolate");
485 let a = "google/siglip2-base-patch16-384";
486 let b = "google/siglip-base-patch16-224";
487
488 let conn = Connection::open_in_memory().unwrap();
489 attach_in(&conn, &ctx, a, true).unwrap();
490 conn.execute(
491 "INSERT INTO emb.embeddings (hash, model_id, embedding, embedded_at)
492 VALUES ('h1', ?1, X'0102', 'now')",
493 [a],
494 )
495 .unwrap();
496 detach(&conn).unwrap();
497
498 attach_in(&conn, &ctx, b, true).unwrap();
499 let n: i64 = conn
500 .query_row("SELECT COUNT(*) FROM emb.embeddings", [], |r| r.get(0))
501 .unwrap();
502 assert_eq!(n, 0, "model b must not see model a's rows");
503 }
504
505 #[test]
506 fn attached_table_is_visible_through_emb_sqlite_master() {
507 let ctx = test_context("master");
512 let conn = Connection::open_in_memory().unwrap();
513 attach_in(&conn, &ctx, "google/siglip2-base-patch16-384", true).unwrap();
514
515 let found: i64 = conn
516 .query_row(
517 "SELECT COUNT(*) FROM emb.sqlite_master
518 WHERE type='table' AND name='embeddings'",
519 [],
520 |r| r.get(0),
521 )
522 .unwrap();
523 assert_eq!(found, 1);
524
525 let unqualified: i64 = conn
526 .query_row(
527 "SELECT COUNT(*) FROM sqlite_master
528 WHERE type='table' AND name='embeddings'",
529 [],
530 |r| r.get(0),
531 )
532 .unwrap();
533 assert_eq!(unqualified, 0, "documents exactly why emb. is required");
534 }
535
536 #[test]
537 fn detach_allows_attaching_a_different_model_on_the_same_connection() {
538 let ctx = test_context("reattach");
539 let conn = Connection::open_in_memory().unwrap();
540 attach_in(&conn, &ctx, "google/siglip2-base-patch16-384", true).unwrap();
541 detach(&conn).unwrap();
542 attach_in(&conn, &ctx, "google/siglip-base-patch16-224", true).unwrap();
543 detach(&conn).unwrap();
544 }
545
546 #[test]
547 fn list_models_returns_sorted_ids_and_ignores_unrelated_files() {
548 let ctx = test_context("list");
549 let conn = Connection::open_in_memory().unwrap();
550 for m in [
551 "google/siglip2-base-patch16-384",
552 "google/siglip-base-patch16-224",
553 ] {
554 attach_in(&conn, &ctx, m, true).unwrap();
555 detach(&conn).unwrap();
556 }
557 std::fs::write(ctx.paths.embeddings.join("notes.txt"), b"x").unwrap();
559
560 let models = list_models_in(&ctx).unwrap();
561 assert_eq!(
562 models,
563 vec![
564 "google/siglip-base-patch16-224".to_string(),
565 "google/siglip2-base-patch16-384".to_string(),
566 ]
567 );
568 }
569
570 #[test]
571 fn list_models_on_a_library_with_no_embeddings_is_empty_not_an_error() {
572 let ctx = test_context("listempty");
573 assert!(list_models_in(&ctx).unwrap().is_empty());
574 }
575
576 #[test]
577 fn counts_by_model_reports_rows_dims_and_size() {
578 let ctx = test_context("counts");
579 let model = "google/siglip2-base-patch16-384";
580 let conn = Connection::open_in_memory().unwrap();
581 attach_in(&conn, &ctx, model, true).unwrap();
582 conn.execute(
584 "INSERT INTO emb.embeddings (hash, model_id, embedding, embedded_at)
585 VALUES ('h1', ?1, zeroblob(1536), 'now')",
586 [model],
587 )
588 .unwrap();
589 detach(&conn).unwrap();
590
591 let counts = counts_by_model_in(&ctx).unwrap();
592 assert_eq!(counts.len(), 1);
593 assert_eq!(counts[0].model_id, model);
594 assert_eq!(counts[0].count, 1);
595 assert_eq!(
596 counts[0].dims, 768,
597 "dims derive from blob length, not a table"
598 );
599 assert!(counts[0].size_bytes > 0);
600 }
601
602 #[test]
603 fn counts_by_model_reports_zero_dims_for_an_empty_model_database() {
604 let ctx = test_context("countsempty");
605 let conn = Connection::open_in_memory().unwrap();
606 attach_in(&conn, &ctx, "google/siglip2-base-patch16-384", true).unwrap();
607 detach(&conn).unwrap();
608
609 let counts = counts_by_model_in(&ctx).unwrap();
610 assert_eq!(counts.len(), 1);
611 assert_eq!(counts[0].count, 0);
612 assert_eq!(counts[0].dims, 0);
613 }
614
615 #[test]
616 fn path_computation_creates_nothing() {
617 let ctx = test_context("nocreate");
620 let _ = db_path_in(&ctx, "google/siglip2-base-patch16-384").unwrap();
621 assert!(
622 !ctx.paths.embeddings.exists(),
623 "path computation must not create {:?}",
624 ctx.paths.embeddings
625 );
626 }
627}