From d6599f178b826316e72b3cdcaa636d2f5baaa9b3 Mon Sep 17 00:00:00 2001
From: Neeraj Gupta <254676+ua741@users.noreply.github.com>
Date: Tue, 14 May 2024 16:25:44 +0530
Subject: [PATCH] [mob] Fix: use DB to identify indexed fileIDs

---
 mobile/lib/db/embeddings_db.dart                    | 13 +++++++++++++
 .../semantic_search/semantic_search_service.dart    |  3 ++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/mobile/lib/db/embeddings_db.dart b/mobile/lib/db/embeddings_db.dart
index 0eb1d3f6d7..64878a2ce4 100644
--- a/mobile/lib/db/embeddings_db.dart
+++ b/mobile/lib/db/embeddings_db.dart
@@ -63,6 +63,19 @@ class EmbeddingsDB {
     return _convertToEmbeddings(results);
   }
 
+  // Get FileIDs for a specific model
+  Future<Set<int>> getFileIDs(Model model) async {
+    final db = await _database;
+    final results = await db.getAll(
+      'SELECT $columnFileID FROM $tableName WHERE $columnModel = ?',
+      [modelToInt(model)!],
+    );
+    if (results.isEmpty) {
+      return <int>{};
+    }
+    return results.map((e) => e[columnFileID] as int).toSet();
+  }
+
   Future<void> put(Embedding embedding) async {
     final db = await _database;
     await db.execute(
diff --git a/mobile/lib/services/machine_learning/semantic_search/semantic_search_service.dart b/mobile/lib/services/machine_learning/semantic_search/semantic_search_service.dart
index 337ca913ff..18c77b43bb 100644
--- a/mobile/lib/services/machine_learning/semantic_search/semantic_search_service.dart
+++ b/mobile/lib/services/machine_learning/semantic_search/semantic_search_service.dart
@@ -225,7 +225,8 @@ class SemanticSearchService {
   Future<List<int>> _getFileIDsToBeIndexed() async {
     final uploadedFileIDs = await FilesDB.instance
         .getOwnedFileIDs(Configuration.instance.getUserID()!);
-    final embeddedFileIDs = _cachedEmbeddings.map((e) => e.fileID).toSet();
+    final embeddedFileIDs =
+        await EmbeddingsDB.instance.getFileIDs(_currentModel);
     uploadedFileIDs.removeWhere(
       (id) => embeddedFileIDs.contains(id),
     );