Browse Source

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

Vishnu Mohandas 1 year ago
parent
commit
712fb85eec

+ 13 - 0
lib/db/embeddings_db.dart

@@ -48,6 +48,19 @@ class EmbeddingsDB {
     return await _isar.embeddings.filter().updationTimeEqualTo(null).findAll();
     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 {
   Future<void> deleteAllForModel(Model model) async {
     await _isar.writeTxn(() async {
     await _isar.writeTxn(() async {
       final embeddings =
       final embeddings =

+ 10 - 1
lib/services/machine_learning/semantic_search/embedding_store.dart

@@ -53,13 +53,22 @@ class EmbeddingStore {
     final fileMap = await FilesDB.instance
     final fileMap = await FilesDB.instance
         .getFilesFromIDs(pendingItems.map((e) => e.fileID).toList());
         .getFilesFromIDs(pendingItems.map((e) => e.fileID).toList());
     _logger.info("Pushing ${pendingItems.length} embeddings");
     _logger.info("Pushing ${pendingItems.length} embeddings");
+    final deletedEntries = <int>[];
     for (final item in pendingItems) {
     for (final item in pendingItems) {
       try {
       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) {
       } catch (e, s) {
         _logger.severe(e, s);
         _logger.severe(e, s);
       }
       }
     }
     }
+    if (deletedEntries.isNotEmpty) {
+      await EmbeddingsDB.instance.deleteEmbeddings(deletedEntries);
+    }
   }
   }
 
 
   Future<void> storeEmbedding(EnteFile file, Embedding embedding) async {
   Future<void> storeEmbedding(EnteFile file, Embedding embedding) async {

+ 9 - 1
lib/services/machine_learning/semantic_search/semantic_search_service.dart

@@ -243,15 +243,23 @@ class SemanticSearchService {
 
 
     final ignoredCollections =
     final ignoredCollections =
         CollectionsService.instance.getHiddenCollectionIds();
         CollectionsService.instance.getHiddenCollectionIds();
+    final deletedEntries = <int>[];
     for (final result in queryResults) {
     for (final result in queryResults) {
       final file = filesMap[result.id];
       final file = filesMap[result.id];
       if (file != null && !ignoredCollections.contains(file.collectionID)) {
       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");
     _logger.info(results.length.toString() + " results");
 
 
+    if (deletedEntries.isNotEmpty) {
+      unawaited(EmbeddingsDB.instance.deleteEmbeddings(deletedEntries));
+    }
+
     return results;
     return results;
   }
   }