feat(discovery): AI tagging, semantic search, similarity, duplicate finder, folder management

Adds a full discovery and organisation layer to the gallery. ## AI Tagger - WD tagger (ONNX via ort) with DirectML/CPU acceleration - Per-image and per-folder tag queuing; configurable batch size and confidence threshold; model downloaded on first use via ureq/zip - Tags stored with source ('ai' | 'user'); user tags are never overwritten by AI; AI tags protected from accidental promotion - Tagger pause/resume per folder; in-flight batches discarded cleanly on pause or cancellation without leaving jobs stuck in processing ## Semantic & Tag Search - CLIP text-query embedding via candle (HuggingFace hub model) - Progressive candidate doubling with filter post-processing so folder/rating/media-kind filters do not silently under-return results - Tag search with DB-backed pagination (offset + COUNT total) - Filename search unchanged; prefix syntax: s:, t:, f: in search bar ## Similarity Search - HNSW in-memory index (hnsw_rs) with monotonic embedding_revision counter for cache invalidation; build_index retries if a concurrent write advances the revision during construction - Folder-scoped similar search uses brute-force cosine scan (no k limit) - Region-based similarity: crop an arbitrary rectangle in the lightbox, embed it with CLIP, find the nearest images globally or within folder - Pagination for both similar and region results; scope toggle (all media / current folder) re-runs the query without reopening image ## Duplicate Finder - Three-phase detection: stat (size) → sample hash (4×16 KB windows) → full-file hash (xxh3, large files only) to eliminate false positives - Filesystem-first deletion: only removes DB rows for files successfully deleted from disk; failed files remain visible for retry - Persisted scan cache (SQLite) with invalidation on reindex and deletion; both global and per-folder scopes invalidated after any deletion ## Tag Cloud & Explore - Visual cluster explore: k-means over CLIP embeddings, configurable k, representative thumbnail per cluster; cache keyed on xxh3 of embedding set - Tag explore: ranked tag list with image counts and representative thumbnail per tag; invalidated when AI tagging completes or folder removed - Tag autocomplete for search bar ## Folder Management - Inline rename (display name only, not OS rename) - Relocate: file-picker to update folder path; all image paths rewritten using SUBSTR prefix replacement to avoid corrupting paths that contain the folder name as a substring - Missing-folder recovery banner: Locate or Remove, never silent deletion - Right-click context menu on sidebar items: Reindex, Rename, Locate, Remove; hover buttons preserved alongside context menu ## Background Tasks - Per-folder progress panel with embedding, tagging, caption, and thumbnail stage breakdown - Retry button per task for failed embeddings and tagging jobs - Completed-task notifications (system tray via tauri-plugin-notification) - Scan errors shown inline; WalkDir permission errors protect existing records from deletion rather than cascading data loss ## Backend correctness - sqlite-vec DML performed outside transactions throughout (virtual-table limitation); embedding revision incremented on delete as well as insert - Tagging job discard checks status = 'processing' so paused jobs reset to 'pending' are also skipped, not just cancelled ones - Stale async responses in Lightbox guarded with cancellation flags and currentImageIdRef for both tag-load and tag-add callbacks - Duplicate cache cleared on reindex; folder-removal clears vector rows outside transaction before deleting the folder row
2026-06-07 22:43:16 +00:00
parent 8905baf4a5
commit 0ca4d142d8
32 changed files with 9539 additions and 760 deletions
@@ -1,5 +1,5 @@
 use anyhow::{anyhow, Result};
-use rusqlite::{ffi::sqlite3_auto_extension, Connection};
+use rusqlite::{ffi::sqlite3_auto_extension, Connection, Error as SqliteError};
 use sqlite_vec::sqlite3_vec_init;
 use std::sync::Once;

@@ -19,8 +19,13 @@ pub fn migrate(conn: &Connection) -> Result<()> {
        "CREATE VIRTUAL TABLE IF NOT EXISTS image_vec USING vec0(
            image_id INTEGER PRIMARY KEY,
            embedding FLOAT[{}] distance_metric=cosine
+        );
+
+        CREATE VIRTUAL TABLE IF NOT EXISTS caption_vec USING vec0(
+            image_id INTEGER PRIMARY KEY,
+            embedding FLOAT[{}] distance_metric=cosine
        );",
-        CLIP_VECTOR_DIM
+        CLIP_VECTOR_DIM, CLIP_VECTOR_DIM
    ))?;
    Ok(())
 }
@@ -28,6 +33,18 @@ pub fn migrate(conn: &Connection) -> Result<()> {
 #[allow(dead_code)]
 pub fn delete_embedding(conn: &Connection, image_id: i64) -> Result<()> {
    conn.execute("DELETE FROM image_vec WHERE image_id = ?1", [image_id])?;
+    // Advance the revision so any cached HNSW index is invalidated after deletions.
+    conn.execute(
+        "INSERT INTO app_kv (key, value) VALUES ('embedding_revision', 1)
+         ON CONFLICT(key) DO UPDATE SET value = value + 1",
+        [],
+    )?;
+    Ok(())
+}
+
+#[allow(dead_code)]
+pub fn delete_caption_embedding(conn: &Connection, image_id: i64) -> Result<()> {
+    conn.execute("DELETE FROM caption_vec WHERE image_id = ?1", [image_id])?;
    Ok(())
 }

@@ -50,26 +67,74 @@ pub fn upsert_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) ->
    Ok(())
 }

-pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) -> Result<Vec<i64>> {
-    let embedding: Vec<u8> = conn.query_row(
+#[allow(dead_code)]
+pub fn upsert_caption_embedding(conn: &Connection, image_id: i64, embedding: &[f32]) -> Result<()> {
+    if embedding.len() != CLIP_VECTOR_DIM {
+        return Err(anyhow!(
+            "expected {}-dimensional embedding, got {}",
+            CLIP_VECTOR_DIM,
+            embedding.len()
+        ));
+    }
+
+    let packed = pack_f32(embedding);
+    conn.execute("DELETE FROM caption_vec WHERE image_id = ?1", [image_id])?;
+    conn.execute(
+        "INSERT INTO caption_vec (image_id, embedding) VALUES (?1, ?2)",
+        (&image_id, &packed),
+    )?;
+    Ok(())
+}
+
+pub fn find_similar_image_ids(
+    conn: &Connection,
+    image_id: i64,
+    limit: usize,
+    folder_id: Option<i64>,
+) -> Result<Vec<i64>> {
+    let embedding: Vec<u8> = match conn.query_row(
        "SELECT embedding FROM image_vec WHERE image_id = ?1",
        [image_id],
        |row| row.get(0),
-    )?;
+    ) {
+        Ok(embedding) => embedding,
+        Err(SqliteError::QueryReturnedNoRows) => return Ok(Vec::new()),
+        Err(error) => return Err(error.into()),
+    };

+    if let Some(folder_id) = folder_id {
+        // Brute-force cosine scan scoped to the folder — avoids the KNN k=4096 limit
+        // and returns exact nearest neighbours within the folder.
+        let mut stmt = conn.prepare(
+            "SELECT v.image_id
+             FROM image_vec v
+             JOIN images i ON i.id = v.image_id
+             WHERE i.folder_id = ?2
+               AND v.image_id != ?3
+             ORDER BY vec_distance_cosine(v.embedding, vec_f32(?1)) ASC
+             LIMIT ?4",
+        )?;
+        let rows = stmt.query_map((&embedding, folder_id, image_id, limit as i64), |row| {
+            row.get::<_, i64>(0)
+        })?;
+        return Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?);
+    }
+
+    // Global KNN search (no folder filter) — use the ANN index.
    let mut stmt = conn.prepare(
        "SELECT image_id
         FROM image_vec
         WHERE embedding MATCH vec_f32(?1)
           AND k = ?2",
    )?;
-    let rows = stmt.query_map((&embedding, (limit as i64) + 1), |row| row.get::<_, i64>(0))?;
+    let rows = stmt
+        .query_map((&embedding, (limit + 1) as i64), |row| row.get::<_, i64>(0))?
+        .collect::<rusqlite::Result<Vec<_>>>()?;

    let mut ids = Vec::new();
    for row in rows {
-        let candidate_id = row?;
-        if candidate_id != image_id {
-            ids.push(candidate_id);
+        if row != image_id {
+            ids.push(row);
        }
        if ids.len() >= limit {
            break;
@@ -78,6 +143,106 @@ pub fn find_similar_image_ids(conn: &Connection, image_id: i64, limit: usize) ->
    Ok(ids)
 }

+// pub fn find_similar_image_matches(
+//     conn: &Connection,
+//     image_id: i64,
+//     folder_id: Option<i64>,
+//     threshold: f32,
+//     offset: usize,
+//     limit: usize,
+// ) -> Result<Vec<(i64, f32)>> {
+//     let embedding: Vec<u8> = match conn.query_row(
+//         "SELECT embedding FROM image_vec WHERE image_id = ?1",
+//         [image_id],
+//         |row| row.get(0),
+//     ) {
+//         Ok(embedding) => embedding,
+//         Err(SqliteError::QueryReturnedNoRows) => return Ok(Vec::new()),
+//         Err(error) => return Err(error.into()),
+//     };
+
+//     let query = match folder_id {
+//         Some(_) => {
+//             "SELECT v.image_id, vec_distance_cosine(v.embedding, vec_f32(?1)) AS distance
+//              FROM image_vec v
+//              JOIN images i ON i.id = v.image_id
+//              WHERE i.folder_id = ?2
+//                AND v.image_id != ?3
+//                AND vec_distance_cosine(v.embedding, vec_f32(?1)) <= ?4
+//              ORDER BY distance ASC
+//              LIMIT ?5 OFFSET ?6"
+//         }
+//         None => {
+//             "SELECT v.image_id, vec_distance_cosine(v.embedding, vec_f32(?1)) AS distance
+//              FROM image_vec v
+//              WHERE v.image_id != ?2
+//                AND vec_distance_cosine(v.embedding, vec_f32(?1)) <= ?3
+//              ORDER BY distance ASC
+//              LIMIT ?4 OFFSET ?5"
+//         }
+//     };
+
+//     let mut stmt = conn.prepare(query)?;
+//     match folder_id {
+//         Some(folder_id) => Ok(stmt
+//             .query_map(
+//                 (
+//                     &embedding,
+//                     folder_id,
+//                     image_id,
+//                     threshold,
+//                     limit as i64,
+//                     offset as i64,
+//                 ),
+//                 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, f32>(1)?)),
+//             )?
+//             .collect::<rusqlite::Result<Vec<_>>>()?),
+//         None => Ok(stmt
+//             .query_map(
+//                 (&embedding, image_id, threshold, limit as i64, offset as i64),
+//                 |row| Ok((row.get::<_, i64>(0)?, row.get::<_, f32>(1)?)),
+//             )?
+//             .collect::<rusqlite::Result<Vec<_>>>()?),
+//     }
+// }
+
+pub fn get_image_embedding(conn: &Connection, image_id: i64) -> Result<Option<Vec<f32>>> {
+    let embedding: Result<Vec<u8>, rusqlite::Error> = conn.query_row(
+        "SELECT embedding FROM image_vec WHERE image_id = ?1",
+        [image_id],
+        |row| row.get(0),
+    );
+
+    match embedding {
+        Ok(bytes) => Ok(Some(unpack_f32(&bytes))),
+        Err(SqliteError::QueryReturnedNoRows) => Ok(None),
+        Err(error) => Err(error.into()),
+    }
+}
+
+pub fn get_embedding_revision(conn: &Connection) -> Result<String> {
+    // Use the monotonically incremented app_kv counter so that two embeddings
+    // saved within the same clock second still advance the revision, preventing
+    // the HNSW cache from serving stale vectors.
+    let revision: i64 = conn
+        .query_row(
+            "SELECT COALESCE((SELECT value FROM app_kv WHERE key = 'embedding_revision'), 0)",
+            [],
+            |row| row.get(0),
+        )
+        .unwrap_or(0);
+    Ok(revision.to_string())
+}
+
+// fn image_ids_for_folder(
+//     conn: &Connection,
+//     folder_id: i64,
+// ) -> Result<std::collections::HashSet<i64>> {
+//     let mut stmt = conn.prepare("SELECT id FROM images WHERE folder_id = ?1")?;
+//     let rows = stmt.query_map([folder_id], |row| row.get::<_, i64>(0))?;
+//     Ok(rows.collect::<rusqlite::Result<std::collections::HashSet<_>>>()?)
+// }
+
 /// Returns all stored image embeddings with their image IDs, optionally filtered to one folder.
 /// Each entry is `(image_id, normalized_f32_embedding)`.
 pub fn get_all_image_embeddings_with_ids(
@@ -155,6 +320,149 @@ pub fn search_image_ids_by_embedding(
    Ok(ids)
 }

+/// Brute-force cosine search scoped to a single folder, ordered by ascending distance.
+/// Used for region-based similarity search where we want folder-scoped results.
+pub fn search_image_ids_by_embedding_in_folder(
+    conn: &Connection,
+    embedding: &[f32],
+    folder_id: i64,
+    exclude_image_id: Option<i64>,
+    limit: usize,
+) -> Result<Vec<i64>> {
+    if embedding.len() != CLIP_VECTOR_DIM {
+        return Err(anyhow!(
+            "expected {}-dimensional embedding, got {}",
+            CLIP_VECTOR_DIM,
+            embedding.len()
+        ));
+    }
+
+    let packed = pack_f32(embedding);
+    let exclude_id = exclude_image_id.unwrap_or(-1);
+    let mut stmt = conn.prepare(
+        "SELECT v.image_id
+         FROM image_vec v
+         JOIN images i ON i.id = v.image_id
+         WHERE i.folder_id = ?2
+           AND v.image_id != ?3
+         ORDER BY vec_distance_cosine(v.embedding, vec_f32(?1)) ASC
+         LIMIT ?4",
+    )?;
+    let rows = stmt.query_map((&packed, folder_id, exclude_id, limit as i64), |row| {
+        row.get::<_, i64>(0)
+    })?;
+    Ok(rows.collect::<rusqlite::Result<Vec<_>>>()?)
+}
+
+#[allow(dead_code)]
+pub fn search_caption_ids_by_embedding(
+    conn: &Connection,
+    embedding: &[f32],
+    limit: usize,
+) -> Result<Vec<i64>> {
+    if embedding.len() != CLIP_VECTOR_DIM {
+        return Err(anyhow!(
+            "expected {}-dimensional embedding, got {}",
+            CLIP_VECTOR_DIM,
+            embedding.len()
+        ));
+    }
+
+    let packed = pack_f32(embedding);
+    let mut stmt = conn.prepare(
+        "SELECT image_id
+         FROM caption_vec
+         WHERE embedding MATCH vec_f32(?1)
+           AND k = ?2",
+    )?;
+    let rows = stmt.query_map((&packed, limit as i64), |row| row.get::<_, i64>(0))?;
+
+    let mut ids = Vec::new();
+    for row in rows {
+        ids.push(row?);
+        if ids.len() >= limit {
+            break;
+        }
+    }
+    Ok(ids)
+}
+
+pub fn count_image_vectors(conn: &Connection) -> Result<i64> {
+    conn.query_row("SELECT COUNT(*) FROM image_vec", [], |row| row.get(0))
+        .map_err(Into::into)
+}
+
+#[allow(dead_code)]
+pub fn count_caption_vectors(conn: &Connection) -> Result<i64> {
+    conn.query_row("SELECT COUNT(*) FROM caption_vec", [], |row| row.get(0))
+        .map_err(Into::into)
+}
+
+pub fn delete_orphaned_embeddings(conn: &Connection) -> Result<usize> {
+    let image_ids = {
+        let mut stmt = conn.prepare("SELECT id FROM images")?;
+        let rows = stmt
+            .query_map([], |row| row.get::<_, i64>(0))?
+            .collect::<rusqlite::Result<std::collections::HashSet<_>>>()?;
+        rows
+    };
+    let vector_ids = {
+        let mut stmt = conn.prepare("SELECT image_id FROM image_vec")?;
+        let rows = stmt
+            .query_map([], |row| row.get::<_, i64>(0))?
+            .collect::<rusqlite::Result<Vec<_>>>()?;
+        rows
+    };
+    let orphaned_ids = vector_ids
+        .into_iter()
+        .filter(|image_id| !image_ids.contains(image_id))
+        .collect::<Vec<_>>();
+
+    for image_id in &orphaned_ids {
+        delete_embedding(conn, *image_id)?;
+    }
+
+    Ok(orphaned_ids.len())
+}
+
+#[allow(dead_code)]
+pub fn delete_orphaned_caption_embeddings(conn: &Connection) -> Result<usize> {
+    let image_ids = {
+        let mut stmt = conn.prepare("SELECT id FROM images")?;
+        let rows = stmt
+            .query_map([], |row| row.get::<_, i64>(0))?
+            .collect::<rusqlite::Result<std::collections::HashSet<_>>>()?;
+        rows
+    };
+    let vector_ids = {
+        let mut stmt = conn.prepare("SELECT image_id FROM caption_vec")?;
+        let rows = stmt
+            .query_map([], |row| row.get::<_, i64>(0))?
+            .collect::<rusqlite::Result<Vec<_>>>()?;
+        rows
+    };
+    let orphaned_ids = vector_ids
+        .into_iter()
+        .filter(|image_id| !image_ids.contains(image_id))
+        .collect::<Vec<_>>();
+
+    for image_id in &orphaned_ids {
+        delete_caption_embedding(conn, *image_id)?;
+    }
+
+    Ok(orphaned_ids.len())
+}
+
+pub fn has_image_vector(conn: &Connection, image_id: i64) -> Result<bool> {
+    conn.query_row(
+        "SELECT EXISTS(SELECT 1 FROM image_vec WHERE image_id = ?1)",
+        [image_id],
+        |row| row.get::<_, i64>(0),
+    )
+    .map(|value| value != 0)
+    .map_err(Into::into)
+}
+
 #[allow(dead_code)]
 fn pack_f32(values: &[f32]) -> Vec<u8> {
    let mut out = Vec::with_capacity(values.len() * std::mem::size_of::<f32>());