Add options to auto generate embeddings and summaries

src/main/java/org/jabref/gui/JabRefGUI.java

@@ -166,13 +166,15 @@ public void initialize() {
  Injector.setModelOrService(ClipBoardManager.class, clipBoardManager);
 
  JabRefGUI.aiService = new AiService(
+ stateManager,
  preferences.getAiPreferences(),
  preferences.getFilePreferences(),
  dialogService,
  taskExecutor);
  Injector.setModelOrService(AiService.class, aiService);
 
  JabRefGUI.chatHistoryService = new ChatHistoryService(
+ stateManager,
  preferences.getCitationKeyPatternPreferences(),
  dialogService);
  Injector.setModelOrService(ChatHistoryService.class, chatHistoryService);

src/main/java/org/jabref/gui/ai/chatting/chathistory/ChatHistoryService.java

@@ -25,7 +25,6 @@
 import org.jabref.model.groups.AbstractGroup;
 import org.jabref.model.groups.GroupTreeNode;
 
-import com.airhacks.afterburner.injection.Injector;
 import com.google.common.eventbus.Subscribe;
 import dev.langchain4j.data.message.ChatMessage;
 import org.slf4j.Logger;
@@ -58,7 +57,7 @@ public class ChatHistoryService implements AutoCloseable {
 
  private static final String CHAT_HISTORY_FILE_NAME = "chat-histories.mv";
 
- private final StateManager stateManager = Injector.instantiateModelOrService(StateManager.class);
+ private final StateManager stateManager;
 
  private final CitationKeyPatternPreferences citationKeyPatternPreferences;
 
@@ -79,31 +78,25 @@ private record ChatHistoryManagementRecord(Optional<BibDatabaseContext> bibDatab
  return o1 == o2 ? 0 : o1.getGroup().getName().compareTo(o2.getGroup().getName());
  });
 
- public ChatHistoryService(CitationKeyPatternPreferences citationKeyPatternPreferences, NotificationService notificationService) {
+ public ChatHistoryService(StateManager stateManager, CitationKeyPatternPreferences citationKeyPatternPreferences, NotificationService notificationService) {
+ this.stateManager = stateManager;
  this.citationKeyPatternPreferences = citationKeyPatternPreferences;
  this.implementation = new MVStoreChatHistoryStorage(Directories.getAiFilesDirectory().resolve(CHAT_HISTORY_FILE_NAME), notificationService);
- configureHistoryTransfer();
- }
-
- public ChatHistoryService(CitationKeyPatternPreferences citationKeyPatternPreferences,
- ChatHistoryStorage implementation) {
- this.citationKeyPatternPreferences = citationKeyPatternPreferences;
- this.implementation = implementation;
 
- configureHistoryTransfer();
+ configureDatabaseListeners();
  }
 
- private void configureHistoryTransfer() {
+ private void configureDatabaseListeners() {
  stateManager.getOpenDatabases().addListener((ListChangeListener<BibDatabaseContext>) change -> {
  while (change.next()) {
  if (change.wasAdded()) {
- change.getAddedSubList().forEach(this::configureHistoryTransfer);
+ change.getAddedSubList().forEach(this::configureDatabaseListeners);
  }
  }
  });
  }
 
- private void configureHistoryTransfer(BibDatabaseContext bibDatabaseContext) {
+ private void configureDatabaseListeners(BibDatabaseContext bibDatabaseContext) {
  bibDatabaseContext.getMetaData().getGroups().ifPresent(rootGroupTreeNode -> {
  rootGroupTreeNode.iterateOverTree().forEach(groupNode -> {
  groupNode.getGroup().nameProperty().addListener((observable, oldValue, newValue) -> {
@@ -264,6 +257,7 @@ public void close() {
  new HashSet<>(groupsChatHistory.keySet()).forEach(this::closeChatHistoryForGroup);
 
  implementation.commit();
+ implementation.close();
  }
 
  private void transferGroupHistory(BibDatabaseContext bibDatabaseContext, GroupTreeNode groupTreeNode, String oldName, String newName) {

src/main/java/org/jabref/gui/preferences/ai/AiTab.fxml

@@ -43,6 +43,21 @@
  </children>
  </HBox>
 
+ <CheckBox fx:id="autoGenerateEmbeddings"
+ mnemonicParsing="false"
+ text="%Automatically generate embeddings for new entries"
+ HBox.hgrow="ALWAYS"
+ maxWidth="Infinity"/>
+
+ <CheckBox fx:id="autoGenerateSummaries"
+ mnemonicParsing="false"
+ text="%Automatically generate summaries for new entries"
+ HBox.hgrow="ALWAYS"
+ maxWidth="Infinity"/>
+
+ <Label styleClass="sectionHeader"
+ text="%Connection"/>
+
  <HBox alignment="CENTER_LEFT"
  layoutX="10.0"
  layoutY="306.0"

src/main/java/org/jabref/gui/preferences/ai/AiTab.java

@@ -29,6 +29,8 @@ public class AiTab extends AbstractPreferenceTabView<AiTabViewModel> implements
  private static final String HUGGING_FACE_CHAT_MODEL_PROMPT = "TinyLlama/TinyLlama_v1.1 (or any other model name)";
 
  @FXML private CheckBox enableAi;
+ @FXML private CheckBox autoGenerateEmbeddings;
+ @FXML private CheckBox autoGenerateSummaries;
 
  @FXML private ComboBox<AiProvider> aiProviderComboBox;
  @FXML private ComboBox<String> chatModelComboBox;
@@ -72,6 +74,10 @@ public void initialize() {
  this.viewModel = new AiTabViewModel(preferences);
 
  enableAi.selectedProperty().bindBidirectional(viewModel.enableAi());
+ autoGenerateSummaries.selectedProperty().bindBidirectional(viewModel.autoGenerateSummaries());
+ autoGenerateSummaries.disableProperty().bind(viewModel.disableAutoGenerateSummaries());
+ autoGenerateEmbeddings.selectedProperty().bindBidirectional(viewModel.autoGenerateEmbeddings());
+ autoGenerateEmbeddings.disableProperty().bind(viewModel.disableAutoGenerateEmbeddings());
 
  new ViewModelListCellFactory<AiProvider>()
  .withText(AiProvider::toString)

src/main/java/org/jabref/gui/preferences/ai/AiTabViewModel.java

@@ -37,6 +37,10 @@ public class AiTabViewModel implements PreferenceTabViewModel {
  private final Locale oldLocale;
 
  private final BooleanProperty enableAi = new SimpleBooleanProperty();
+ private final BooleanProperty autoGenerateEmbeddings = new SimpleBooleanProperty();
+ private final BooleanProperty disableAutoGenerateEmbeddings = new SimpleBooleanProperty();
+ private final BooleanProperty autoGenerateSummaries = new SimpleBooleanProperty();
+ private final BooleanProperty disableAutoGenerateSummaries = new SimpleBooleanProperty();
 
  private final ListProperty<AiProvider> aiProvidersList =
  new SimpleListProperty<>(FXCollections.observableArrayList(AiProvider.values()));
@@ -295,6 +299,8 @@ public void setValues() {
  huggingFaceChatModel.setValue(aiPreferences.getHuggingFaceChatModel());
 
  enableAi.setValue(aiPreferences.getEnableAi());
+ autoGenerateSummaries.setValue(aiPreferences.getAutoGenerateSummaries());
+ autoGenerateEmbeddings.setValue(aiPreferences.getAutoGenerateEmbeddings());
 
  selectedAiProvider.setValue(aiPreferences.getAiProvider());
 
@@ -313,6 +319,8 @@ public void setValues() {
  @Override
  public void storeSettings() {
  aiPreferences.setEnableAi(enableAi.get());
+ aiPreferences.setAutoGenerateEmbeddings(autoGenerateEmbeddings.get());
+ aiPreferences.setAutoGenerateSummaries(autoGenerateSummaries.get());
 
  aiPreferences.setAiProvider(selectedAiProvider.get());
 
@@ -407,6 +415,22 @@ public BooleanProperty enableAi() {
  return enableAi;
  }
 
+ public BooleanProperty autoGenerateEmbeddings() {
+ return autoGenerateEmbeddings;
+ }
+
+ public BooleanProperty disableAutoGenerateEmbeddings() {
+ return disableAutoGenerateEmbeddings;
+ }
+
+ public BooleanProperty autoGenerateSummaries() {
+ return autoGenerateSummaries;
+ }
+
+ public BooleanProperty disableAutoGenerateSummaries() {
+ return disableAutoGenerateSummaries;
+ }
+
  public ReadOnlyListProperty<AiProvider> aiProvidersProperty() {
  return aiProvidersList;
  }

src/main/java/org/jabref/logic/ai/AiDefaultPreferences.java

@@ -52,6 +52,8 @@ public class AiDefaultPreferences {
  );
 
  public static final boolean ENABLE_CHAT = false;
+ public static final boolean AUTO_GENERATE_EMBEDDINGS = false;
+ public static final boolean AUTO_GENERATE_SUMMARIES = false;
 
  public static final AiProvider PROVIDER = AiProvider.OPEN_AI;
 

src/main/java/org/jabref/logic/ai/AiPreferences.java

@@ -31,6 +31,8 @@ public class AiPreferences {
  private static final String KEYRING_AI_SERVICE_ACCOUNT = "apiKey";
 
  private final BooleanProperty enableAi;
+ private final BooleanProperty autoGenerateEmbeddings;
+ private final BooleanProperty autoGenerateSummaries;
 
  private final ObjectProperty<AiProvider> aiProvider;
 
@@ -58,6 +60,8 @@ public class AiPreferences {
  private Runnable apiKeyChangeListener;
 
  public AiPreferences(boolean enableAi,
+ boolean autoGenerateEmbeddings,
+ boolean autoGenerateSummaries,
  AiProvider aiProvider,
  String openAiChatModel,
  String mistralAiChatModel,
@@ -78,6 +82,8 @@ public AiPreferences(boolean enableAi,
  double ragMinScore
  ) {
  this.enableAi = new SimpleBooleanProperty(enableAi);
+ this.autoGenerateEmbeddings = new SimpleBooleanProperty(autoGenerateEmbeddings);
+ this.autoGenerateSummaries = new SimpleBooleanProperty(autoGenerateSummaries);
 
  this.aiProvider = new SimpleObjectProperty<>(aiProvider);
 
@@ -143,6 +149,30 @@ public void setEnableAi(boolean enableAi) {
  this.enableAi.set(enableAi);
  }
 
+ public BooleanProperty autoGenerateEmbeddingsProperty() {
+ return autoGenerateEmbeddings;
+ }
+
+ public boolean getAutoGenerateEmbeddings() {
+ return autoGenerateEmbeddings.get();
+ }
+
+ public void setAutoGenerateEmbeddings(boolean autoGenerateEmbeddings) {
+ this.autoGenerateEmbeddings.set(autoGenerateEmbeddings);
+ }
+
+ public BooleanProperty autoGenerateSummariesProperty() {
+ return autoGenerateSummaries;
+ }
+
+ public boolean getAutoGenerateSummaries() {
+ return autoGenerateSummaries.get();
+ }
+
+ public void setAutoGenerateSummaries(boolean autoGenerateSummaries) {
+ this.autoGenerateSummaries.set(autoGenerateSummaries);
+ }
+
  public ObjectProperty<AiProvider> aiProviderProperty() {
  return aiProvider;
  }

src/main/java/org/jabref/logic/ai/AiService.java

@@ -6,6 +6,7 @@
 import javafx.beans.property.BooleanProperty;
 import javafx.beans.property.SimpleBooleanProperty;
 
+import org.jabref.gui.StateManager;
 import org.jabref.logic.FilePreferences;
 import org.jabref.logic.ai.chatting.AiChatService;
 import org.jabref.logic.ai.chatting.model.JabRefChatLanguageModel;
@@ -52,7 +53,8 @@ public class AiService implements AutoCloseable {
  private final IngestionService ingestionService;
  private final SummariesService summariesService;
 
- public AiService(AiPreferences aiPreferences,
+ public AiService(StateManager stateManager,
+ AiPreferences aiPreferences,
  FilePreferences filePreferences,
  NotificationService notificationService,
  TaskExecutor taskExecutor
@@ -64,8 +66,11 @@ public AiService(AiPreferences aiPreferences,
  this.mvStoreSummariesStorage = new MVStoreSummariesStorage(Directories.getAiFilesDirectory().resolve(SUMMARIES_FILE_NAME), notificationService);
 
  this.jabRefEmbeddingModel = new JabRefEmbeddingModel(aiPreferences, notificationService, taskExecutor);
+
  this.aiChatService = new AiChatService(aiPreferences, jabRefChatLanguageModel, jabRefEmbeddingModel, mvStoreEmbeddingStore, cachedThreadPool);
+
  this.ingestionService = new IngestionService(
+ stateManager,
  aiPreferences,
  shutdownSignal,
  jabRefEmbeddingModel,
@@ -74,7 +79,16 @@ public AiService(AiPreferences aiPreferences,
  filePreferences,
  taskExecutor
  );
- this.summariesService = new SummariesService(aiPreferences, mvStoreSummariesStorage, jabRefChatLanguageModel, shutdownSignal, filePreferences, taskExecutor);
+
+ this.summariesService = new SummariesService(
+ stateManager,
+ aiPreferences,
+ mvStoreSummariesStorage,
+ jabRefChatLanguageModel,
+ shutdownSignal,
+ filePreferences,
+ taskExecutor
+ );
  }
 
  public JabRefChatLanguageModel getChatLanguageModel() {

src/main/java/org/jabref/logic/ai/ingestion/IngestionService.java

@@ -1,21 +1,26 @@
 package org.jabref.logic.ai.ingestion;
 
 import java.util.ArrayList;
-import java.util.HashMap;
 import java.util.List;
-import java.util.Map;
+import java.util.TreeMap;
 
 import javafx.beans.property.ReadOnlyBooleanProperty;
 import javafx.beans.property.StringProperty;
+import javafx.collections.ListChangeListener;
 
+import org.jabref.gui.StateManager;
 import org.jabref.logic.FilePreferences;
 import org.jabref.logic.ai.AiPreferences;
 import org.jabref.logic.ai.processingstatus.ProcessingInfo;
 import org.jabref.logic.ai.processingstatus.ProcessingState;
 import org.jabref.logic.util.TaskExecutor;
 import org.jabref.model.database.BibDatabaseContext;
+import org.jabref.model.database.event.EntriesAddedEvent;
 import org.jabref.model.entry.LinkedFile;
+import org.jabref.model.entry.event.FieldChangedEvent;
+import org.jabref.model.entry.field.StandardField;
 
+import com.google.common.eventbus.Subscribe;
 import dev.langchain4j.data.segment.TextSegment;
 import dev.langchain4j.model.embedding.EmbeddingModel;
 import dev.langchain4j.store.embedding.EmbeddingStore;
@@ -25,25 +30,29 @@
  * Use this class in the logic and UI.
  */
 public class IngestionService {
- private final Map<LinkedFile, ProcessingInfo<LinkedFile, Void>> ingestionStatusMap = new HashMap<>();
+ // We use a {@link TreeMap} here for the same reasons we use it in {@link ChatHistoryService}.
+ private final TreeMap<LinkedFile, ProcessingInfo<LinkedFile, Void>> ingestionStatusMap = new TreeMap<>((o1, o2) -> o1 == o2 ? 0 : o1.getLink().compareTo(o2.getLink()));
 
  private final List<List<LinkedFile>> listsUnderIngestion = new ArrayList<>();
 
+ private final AiPreferences aiPreferences;
  private final FilePreferences filePreferences;
  private final TaskExecutor taskExecutor;
 
  private final FileEmbeddingsManager fileEmbeddingsManager;
 
  private final ReadOnlyBooleanProperty shutdownSignal;
 
- public IngestionService(AiPreferences aiPreferences,
+ public IngestionService(StateManager stateManager,
+ AiPreferences aiPreferences,
  ReadOnlyBooleanProperty shutdownSignal,
  EmbeddingModel embeddingModel,
  EmbeddingStore<TextSegment> embeddingStore,
  FullyIngestedDocumentsTracker fullyIngestedDocumentsTracker,
  FilePreferences filePreferences,
  TaskExecutor taskExecutor
  ) {
+ this.aiPreferences = aiPreferences;
  this.filePreferences = filePreferences;
  this.taskExecutor = taskExecutor;
 
@@ -56,6 +65,49 @@ public IngestionService(AiPreferences aiPreferences,
  );
 
  this.shutdownSignal = shutdownSignal;
+
+ configureDatabaseListeners(stateManager);
+ }
+
+ private void configureDatabaseListeners(StateManager stateManager) {
+ stateManager.getOpenDatabases().addListener((ListChangeListener<BibDatabaseContext>) change -> {
+ while (change.next()) {
+ if (change.wasAdded()) {
+ change.getAddedSubList().forEach(this::configureDatabaseListeners);
+ }
+ }
+ });
+ }
+
+ private void configureDatabaseListeners(BibDatabaseContext bibDatabaseContext) {
+ // GC was eating the listeners, so we have to fall back to the event bus.
+ bibDatabaseContext.getDatabase().registerListener(new EntriesChangedListener(bibDatabaseContext));
+ }
+
+ private class EntriesChangedListener {
+ private final BibDatabaseContext bibDatabaseContext;
+
+ public EntriesChangedListener(BibDatabaseContext bibDatabaseContext) {
+ this.bibDatabaseContext = bibDatabaseContext;
+ }
+
+ @Subscribe
+ public void listen(EntriesAddedEvent e) {
+ e.getBibEntries().forEach(entry -> {
+ if (aiPreferences.getAutoGenerateEmbeddings()) {
+ entry.getFiles().forEach(linkedFile -> ingest(linkedFile, bibDatabaseContext));
+ }
+
+ entry.registerListener(this);
+ });
+ }
+
+ @Subscribe
+ public void listen(FieldChangedEvent e) {
+ if (e.getField() == StandardField.FILE && aiPreferences.getAutoGenerateEmbeddings()) {
+ e.getBibEntry().getFiles().forEach(linkedFile -> ingest(linkedFile, bibDatabaseContext));
+ }
+ }
  }
 
  /**