Ensure unique embeddings for a given file and model (#1638)

ente-io · Jan 6, 2024 · 0da0df0 · 0da0df0
2 parents ab1b3ec + 0fad642
commit 0da0df0
Show file tree

Hide file tree

Showing 7 changed files with 327 additions and 10 deletions.
diff --git a/lib/db/embeddings_db.dart b/lib/db/embeddings_db.dart
@@ -32,14 +32,14 @@ class EmbeddingsDB {
 
   Future<void> put(Embedding embedding) {
     return _isar.writeTxn(() async {
-      await _isar.embeddings.put(embedding);
+      await _isar.embeddings.putByIndex(Embedding.index, embedding);
       Bus.instance.fire(EmbeddingUpdatedEvent());
     });
   }
 
   Future<void> putMany(List<Embedding> embeddings) {
     return _isar.writeTxn(() async {
-      await _isar.embeddings.putAll(embeddings);
+      await _isar.embeddings.putAllByIndex(Embedding.index, embeddings);
       Bus.instance.fire(EmbeddingUpdatedEvent());
     });
   }

diff --git a/lib/models/embedding.dart b/lib/models/embedding.dart
@@ -6,9 +6,12 @@ part 'embedding.g.dart';
 
 @collection
 class Embedding {
-  Id id = Isar.autoIncrement; // you can also use id = null to auto increment
+  static const index = 'unique_file_model_embedding';
+
+  Id id = Isar.autoIncrement;
   final int fileID;
   @enumerated
+  @Index(name: index, composite: [CompositeIndex('fileID')], unique: true, replace: true)
   final Model model;
   final List<double> embedding;
   int? updationTime;

diff --git a/lib/models/embedding.g.dart b/lib/models/embedding.g.dart
diff --git a/lib/services/semantic_search/embedding_store.dart b/lib/services/semantic_search/embedding_store.dart
@@ -61,6 +61,11 @@ class EmbeddingStore {
     unawaited(_pushEmbedding(file, embedding));
   }
 
+  Future<void> clearEmbeddings(Model model) async {
+    await EmbeddingsDB.instance.deleteAllForModel(model);
+    await _preferences.remove(kEmbeddingsSyncTimeKey);
+  }
+
   Future<void> _pushEmbedding(EnteFile file, Embedding embedding) async {
     final encryptionKey = getFileKey(file);
     final embeddingJSON = jsonEncode(embedding.embedding);