Rework embeddings cache clearing

JabRef · Aug 3, 2024 · 44e8bdf · 44e8bdf
1 parent 89d7f77
commit 44e8bdf
Show file tree

Hide file tree

Showing 5 changed files with 29 additions and 19 deletions.
diff --git a/src/main/java/org/jabref/gui/actions/StandardActions.java b/src/main/java/org/jabref/gui/actions/StandardActions.java
@@ -201,7 +201,7 @@ public enum StandardActions implements Action {
     GROUP_ENTRIES_ADD(Localization.lang("Add selected entries to this group")),
     GROUP_ENTRIES_REMOVE(Localization.lang("Remove selected entries from this group")),
 
-    REGENERATE_EMBEDDINGS_CACHE(Localization.lang("Regenerate embeddings cache"));
+    CLEAR_EMBEDDINGS_CACHE(Localization.lang("Clear embeddings cache"));
 
     private String text;
     private final String description;

diff --git a/...ef/gui/ai/RegenerateEmbeddingsAction.java → .../jabref/gui/ai/ClearEmbeddingsAction.java b/...ef/gui/ai/RegenerateEmbeddingsAction.java → .../jabref/gui/ai/ClearEmbeddingsAction.java
@@ -1,25 +1,28 @@
 package org.jabref.gui.ai;
 
+import java.util.List;
+
 import org.jabref.gui.DialogService;
 import org.jabref.gui.StateManager;
 import org.jabref.gui.actions.SimpleCommand;
 import org.jabref.gui.util.BackgroundTask;
 import org.jabref.gui.util.TaskExecutor;
 import org.jabref.logic.ai.AiService;
 import org.jabref.logic.l10n.Localization;
+import org.jabref.model.entry.LinkedFile;
 
 import static org.jabref.gui.actions.ActionHelper.needsDatabase;
 
-public class RegenerateEmbeddingsAction extends SimpleCommand {
+public class ClearEmbeddingsAction extends SimpleCommand {
     private final StateManager stateManager;
     private final DialogService dialogService;
     private final AiService aiService;
     private final TaskExecutor taskExecutor;
 
-    public RegenerateEmbeddingsAction(StateManager stateManager,
-                                      DialogService dialogService,
-                                      AiService aiService,
-                                      TaskExecutor taskExecutor) {
+    public ClearEmbeddingsAction(StateManager stateManager,
+                                 DialogService dialogService,
+                                 AiService aiService,
+                                 TaskExecutor taskExecutor) {
         this.stateManager = stateManager;
         this.dialogService = dialogService;
         this.taskExecutor = taskExecutor;
@@ -34,16 +37,25 @@ public void execute() {
         }
 
         boolean confirmed = dialogService.showConfirmationDialogAndWait(
-                Localization.lang("Regenerate embeddings cache"),
-                Localization.lang("Regenerate embeddings cache for current library?"));
+                Localization.lang("Clear embeddings cache"),
+                Localization.lang("Clear embeddings cache for current library?"));
 
         if (!confirmed) {
             return;
         }
 
-        dialogService.notify(Localization.lang("Regenerating embeddings cache..."));
+        dialogService.notify(Localization.lang("Clearing embeddings cache..."));
+
+        List<LinkedFile> linkedFile = stateManager
+                .getActiveDatabase()
+                .get()
+                .getDatabase()
+                .getEntries()
+                .stream()
+                .flatMap(entry -> entry.getFiles().stream())
+                .toList();
 
-        BackgroundTask.wrap(() -> aiService.getEmbeddingsManager().invalidate())
+        BackgroundTask.wrap(() -> aiService.getEmbeddingsManager().clearEmbeddingsFor(linkedFile))
                       .executeWith(taskExecutor);
     }
 }
diff --git a/src/main/java/org/jabref/gui/frame/MainMenu.java b/src/main/java/org/jabref/gui/frame/MainMenu.java
@@ -15,7 +15,7 @@
 import org.jabref.gui.StateManager;
 import org.jabref.gui.actions.ActionFactory;
 import org.jabref.gui.actions.StandardActions;
-import org.jabref.gui.ai.RegenerateEmbeddingsAction;
+import org.jabref.gui.ai.ClearEmbeddingsAction;
 import org.jabref.gui.auximport.NewSubLibraryAction;
 import org.jabref.gui.bibtexextractor.ExtractBibtexAction;
 import org.jabref.gui.citationkeypattern.GenerateCitationKeyAction;
@@ -314,7 +314,7 @@ private void createMenu() {
                 new SeparatorMenuItem(),
 
                 factory.createMenuItem(StandardActions.REBUILD_FULLTEXT_SEARCH_INDEX, new RebuildFulltextSearchIndexAction(stateManager, frame::getCurrentLibraryTab, dialogService, preferencesService.getFilePreferences(), taskExecutor)),
-                factory.createMenuItem(StandardActions.REGENERATE_EMBEDDINGS_CACHE, new RegenerateEmbeddingsAction(stateManager, dialogService, aiService, taskExecutor)),
+                factory.createMenuItem(StandardActions.CLEAR_EMBEDDINGS_CACHE, new ClearEmbeddingsAction(stateManager, dialogService, aiService, taskExecutor)),
 
                 new SeparatorMenuItem(),
 

diff --git a/src/main/java/org/jabref/logic/ai/FileEmbeddingsManager.java b/src/main/java/org/jabref/logic/ai/FileEmbeddingsManager.java
@@ -1,6 +1,5 @@
 package org.jabref.logic.ai;
 
-import java.util.HashSet;
 import java.util.List;
 import java.util.Optional;
 import java.util.Set;
@@ -107,8 +106,7 @@ public boolean hasIngestedLinkedFiles(List<LinkedFile> linkedFiles) {
         return hasIngestedDocuments(linkedFiles.stream().map(LinkedFile::getLink).toList());
     }
 
-    public void invalidate() {
-        Set<String> ingestedClone = new HashSet<>(getIngestedDocuments());
-        ingestedClone.forEach(this::removeDocument);
+    public void clearEmbeddingsFor(List<LinkedFile> linkedFiles) {
+        linkedFiles.stream().map(LinkedFile::getLink).forEach(this::removeDocument);
     }
 }
diff --git a/src/main/resources/l10n/JabRef_en.properties b/src/main/resources/l10n/JabRef_en.properties
@@ -2536,9 +2536,9 @@ RAG\ -\ maximum\ results\ count=RAG - maximum results count
 RAG\ -\ minimum\ score=RAG - minimum score
 RAG\ max\ results\ count\ must\ be\ greater\ than\ 0=RAG max results count must be greater than 0
 RAG\ min\ score\ must\ be\ greater\ than\ 0\ and\ less\ than\ 1=RAG min score must be greater than 0 and less than 1
-Regenerate\ embeddings\ cache=Regenerate embeddings cache
-Regenerate\ embeddings\ cache\ for\ current\ library?=Regenerate embeddings cache for current library?
-Regenerating\ embeddings\ cache...=Regenerating embeddings cache...
+Clear\ embeddings\ cache=Clear embeddings cache
+Clear\ embeddings\ cache\ for\ current\ library?=Clear embeddings cache for current library?
+Clearing\ embeddings\ cache...=Clearing embeddings cache...
 Temperature=Temperature
 Temperature\ must\ be\ between\ 0\ and\ 2=Temperature must be between 0 and 2
 The\ embeddings\ of\ the\ file\ are\ currently\ being\ generated.\ Please\ wait,\ and\ at\ the\ end\ you\ will\ be\ able\ to\ chat.=The embeddings of the file are currently being generated. Please wait, and at the end you will be able to chat.