Clear embeddings for files that no longer exist locally (#1734)

ente-io · Feb 20, 2024 · 712fb85 · 712fb85
2 parents ceb4a00 + e8b0212
commit 712fb85
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 2 deletions.
diff --git a/lib/db/embeddings_db.dart b/lib/db/embeddings_db.dart
@@ -48,6 +48,19 @@ class EmbeddingsDB {
     return await _isar.embeddings.filter().updationTimeEqualTo(null).findAll();
   }
 
+  Future<void> deleteEmbeddings(List<int> fileIDs) async {
+    await _isar.writeTxn(() async {
+      final embeddings = <Embedding>[];
+      for (final fileID in fileIDs) {
+        embeddings.addAll(
+          await _isar.embeddings.filter().fileIDEqualTo(fileID).findAll(),
+        );
+      }
+      await _isar.embeddings.deleteAll(embeddings.map((e) => e.id).toList());
+      Bus.instance.fire(EmbeddingUpdatedEvent());
+    });
+  }
+
   Future<void> deleteAllForModel(Model model) async {
     await _isar.writeTxn(() async {
       final embeddings =

diff --git a/lib/services/machine_learning/semantic_search/embedding_store.dart b/lib/services/machine_learning/semantic_search/embedding_store.dart
@@ -53,13 +53,22 @@ class EmbeddingStore {
     final fileMap = await FilesDB.instance
         .getFilesFromIDs(pendingItems.map((e) => e.fileID).toList());
     _logger.info("Pushing ${pendingItems.length} embeddings");
+    final deletedEntries = <int>[];
     for (final item in pendingItems) {
       try {
-        await _pushEmbedding(fileMap[item.fileID]!, item);
+        final file = fileMap[item.fileID];
+        if (file != null) {
+          await _pushEmbedding(file, item);
+        } else {
+          deletedEntries.add(item.fileID);
+        }
       } catch (e, s) {
         _logger.severe(e, s);
       }
     }
+    if (deletedEntries.isNotEmpty) {
+      await EmbeddingsDB.instance.deleteEmbeddings(deletedEntries);
+    }
   }
 
   Future<void> storeEmbedding(EnteFile file, Embedding embedding) async {

diff --git a/lib/services/machine_learning/semantic_search/semantic_search_service.dart b/lib/services/machine_learning/semantic_search/semantic_search_service.dart
@@ -243,15 +243,23 @@ class SemanticSearchService {
 
     final ignoredCollections =
         CollectionsService.instance.getHiddenCollectionIds();
+    final deletedEntries = <int>[];
     for (final result in queryResults) {
       final file = filesMap[result.id];
       if (file != null && !ignoredCollections.contains(file.collectionID)) {
-        results.add(filesMap[result.id]!);
+        results.add(file);
+      }
+      if (file == null) {
+        deletedEntries.add(result.id);
       }
     }
 
     _logger.info(results.length.toString() + " results");
 
+    if (deletedEntries.isNotEmpty) {
+      unawaited(EmbeddingsDB.instance.deleteEmbeddings(deletedEntries));
+    }
+
     return results;
   }