Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add options to auto generate embeddings and summaries #11833

Merged
merged 7 commits into from
Oct 6, 2024
Merged

Add options to auto generate embeddings and summaries #11833

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
b4a9c94
Add options to auto generate embeddings and summaries. Unfortunately,…
InAnYan Sep 25, 2024
57717fa
GC ate my listeners
InAnYan Sep 26, 2024
2f1990e
Fix checkers
InAnYan Sep 26, 2024
51a693e
Merge branch 'upstream-main' into generate-ai-for-new
InAnYan Oct 3, 2024
519a8fd
Fix architecture issues
InAnYan Oct 3, 2024
386c5ac
Move ChatHistoryService to logic layer
InAnYan Oct 6, 2024
8466e83
Move ChatHistoryService to logic layer
InAnYan Oct 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions src/main/java/org/jabref/gui/LibraryTab.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import javafx.util.Duration;

import org.jabref.gui.actions.StandardActions;
import org.jabref.gui.ai.chatting.chathistory.ChatHistoryService;
import org.jabref.gui.autocompleter.AutoCompletePreferences;
import org.jabref.gui.autocompleter.PersonNameSuggestionProvider;
import org.jabref.gui.autocompleter.SuggestionProvider;
Expand All @@ -62,8 +63,10 @@
import org.jabref.gui.undo.UndoableFieldChange;
import org.jabref.gui.undo.UndoableInsertEntries;
import org.jabref.gui.undo.UndoableRemoveEntries;
import org.jabref.gui.util.AiUiConnector;
import org.jabref.gui.util.OptionalObjectProperty;
import org.jabref.gui.util.UiTaskExecutor;
import org.jabref.logic.ai.AiService;
import org.jabref.logic.citationstyle.CitationStyleCache;
import org.jabref.logic.importer.FetcherClientException;
import org.jabref.logic.importer.FetcherException;
Expand Down Expand Up @@ -169,6 +172,9 @@ private enum PanelMode { MAIN_TABLE, MAIN_TABLE_AND_ENTRY_EDITOR }
private ImportHandler importHandler;
private LuceneManager luceneManager;

private final AiService aiService;
private final ChatHistoryService chatHistoryService;

/**
* @param isDummyContext Indicates whether the database context is a dummy. A dummy context is used to display a progress indicator while parsing the database.
* If the context is a dummy, the Lucene index should not be created, as both the dummy context and the actual context share the same index path {@link BibDatabaseContext#getFulltextIndexPath()}.
Expand All @@ -178,6 +184,8 @@ private enum PanelMode { MAIN_TABLE, MAIN_TABLE_AND_ENTRY_EDITOR }
private LibraryTab(BibDatabaseContext bibDatabaseContext,
LibraryTabContainer tabContainer,
DialogService dialogService,
AiService aiService,
ChatHistoryService chatHistoryService,
GuiPreferences preferences,
StateManager stateManager,
FileUpdateMonitor fileUpdateMonitor,
Expand All @@ -197,6 +205,8 @@ private LibraryTab(BibDatabaseContext bibDatabaseContext,
this.clipBoardManager = clipBoardManager;
this.taskExecutor = taskExecutor;
this.directoryMonitorManager = new DirectoryMonitorManager(Injector.instantiateModelOrService(DirectoryMonitor.class));
this.aiService = aiService;
this.chatHistoryService = chatHistoryService;

initializeComponentsAndListeners(isDummyContext);

Expand Down Expand Up @@ -249,6 +259,8 @@ private void initializeComponentsAndListeners(boolean isDummyContext) {

this.entryEditor = createEntryEditor();

AiUiConnector.setupDatabase(aiService, chatHistoryService, bibDatabaseContext);

Platform.runLater(() -> {
EasyBind.subscribe(changedProperty, this::updateTabTitle);
stateManager.getOpenDatabases().addListener((ListChangeListener<BibDatabaseContext>) c ->
Expand Down Expand Up @@ -1052,6 +1064,8 @@ public void resetChangedProperties() {
public static LibraryTab createLibraryTab(BackgroundTask<ParserResult> dataLoadingTask,
Path file,
DialogService dialogService,
AiService aiService,
ChatHistoryService chatHistoryService,
GuiPreferences preferences,
StateManager stateManager,
LibraryTabContainer tabContainer,
Expand All @@ -1067,6 +1081,8 @@ public static LibraryTab createLibraryTab(BackgroundTask<ParserResult> dataLoadi
context,
tabContainer,
dialogService,
aiService,
chatHistoryService,
preferences,
stateManager,
fileUpdateMonitor,
Expand All @@ -1088,6 +1104,8 @@ public static LibraryTab createLibraryTab(BackgroundTask<ParserResult> dataLoadi
public static LibraryTab createLibraryTab(BibDatabaseContext databaseContext,
LibraryTabContainer tabContainer,
DialogService dialogService,
AiService aiService,
ChatHistoryService chatHistoryService,
GuiPreferences preferences,
StateManager stateManager,
FileUpdateMonitor fileUpdateMonitor,
Expand All @@ -1101,6 +1119,8 @@ public static LibraryTab createLibraryTab(BibDatabaseContext databaseContext,
databaseContext,
tabContainer,
dialogService,
aiService,
chatHistoryService,
preferences,
stateManager,
fileUpdateMonitor,
Expand Down
69 changes: 13 additions & 56 deletions src/main/java/org/jabref/gui/ai/chatting/chathistory/ChatHistoryService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
import java.util.TreeMap;

import javafx.collections.FXCollections;
import javafx.collections.ListChangeListener;
import javafx.collections.ObservableList;

import org.jabref.gui.StateManager;
Expand All @@ -25,7 +24,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;
Expand Down Expand Up @@ -58,12 +56,13 @@ public class ChatHistoryService implements AutoCloseable {

private static final String CHAT_HISTORY_FILE_NAME = "chat-histories.mv";

private final StateManager stateManager = Injector.instantiateModelOrService(StateManager.class);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think, you can move the class to logic. No gui dependency as far as I can see.

Also the JavaDoc can be updated --> no more use of StateManager.


private final CitationKeyPatternPreferences citationKeyPatternPreferences;

private final ChatHistoryStorage implementation;

// Note about `Optional<BibDatabaseContext>`: it was necessary in previous version, but currently we never save an `Optional.empty()`.
// However, we decided to left it here: to reduce migrations and to make possible to chat with a {@link BibEntry} without {@link BibDatabaseContext}
// ({@link BibDatabaseContext} is required only for load/store of the chat).
private record ChatHistoryManagementRecord(Optional<BibDatabaseContext> bibDatabaseContext, ObservableList<ChatMessage> chatHistory) { }

// We use a {@link TreeMap} here to store {@link BibEntry} chat histories by their id.
Expand All @@ -82,28 +81,9 @@ private record ChatHistoryManagementRecord(Optional<BibDatabaseContext> bibDatab
public ChatHistoryService(CitationKeyPatternPreferences citationKeyPatternPreferences, NotificationService notificationService) {
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();
}

private void configureHistoryTransfer() {
stateManager.getOpenDatabases().addListener((ListChangeListener<BibDatabaseContext>) change -> {
while (change.next()) {
if (change.wasAdded()) {
change.getAddedSubList().forEach(this::configureHistoryTransfer);
}
}
});
}

private void configureHistoryTransfer(BibDatabaseContext bibDatabaseContext) {
public void setupDatabase(BibDatabaseContext bibDatabaseContext) {
bibDatabaseContext.getMetaData().getGroups().ifPresent(rootGroupTreeNode -> {
rootGroupTreeNode.iterateOverTree().forEach(groupNode -> {
groupNode.getGroup().nameProperty().addListener((observable, oldValue, newValue) -> {
Expand All @@ -125,20 +105,18 @@ private void configureHistoryTransfer(BibDatabaseContext bibDatabaseContext) {
});
}

public ObservableList<ChatMessage> getChatHistoryForEntry(BibEntry entry) {
public ObservableList<ChatMessage> getChatHistoryForEntry(BibDatabaseContext bibDatabaseContext, BibEntry entry) {
return bibEntriesChatHistory.computeIfAbsent(entry, entryArg -> {
Optional<BibDatabaseContext> bibDatabaseContext = findBibDatabaseForEntry(entry);

ObservableList<ChatMessage> chatHistory;

if (bibDatabaseContext.isEmpty() || entry.getCitationKey().isEmpty() || !correctCitationKey(bibDatabaseContext.get(), entry) || bibDatabaseContext.get().getDatabasePath().isEmpty()) {
if (entry.getCitationKey().isEmpty() || !correctCitationKey(bibDatabaseContext, entry) || bibDatabaseContext.getDatabasePath().isEmpty()) {
chatHistory = FXCollections.observableArrayList();
} else {
List<ChatMessage> chatMessagesList = implementation.loadMessagesForEntry(bibDatabaseContext.get().getDatabasePath().get(), entry.getCitationKey().get());
List<ChatMessage> chatMessagesList = implementation.loadMessagesForEntry(bibDatabaseContext.getDatabasePath().get(), entry.getCitationKey().get());
chatHistory = FXCollections.observableArrayList(chatMessagesList);
}

return new ChatHistoryManagementRecord(bibDatabaseContext, chatHistory);
return new ChatHistoryManagementRecord(Optional.of(bibDatabaseContext), chatHistory);
}).chatHistory;
}

Expand Down Expand Up @@ -173,24 +151,22 @@ public void closeChatHistoryForEntry(BibEntry entry) {
bibEntriesChatHistory.remove(entry);
}

public ObservableList<ChatMessage> getChatHistoryForGroup(GroupTreeNode group) {
public ObservableList<ChatMessage> getChatHistoryForGroup(BibDatabaseContext bibDatabaseContext, GroupTreeNode group) {
return groupsChatHistory.computeIfAbsent(group, groupArg -> {
Optional<BibDatabaseContext> bibDatabaseContext = findBibDatabaseForGroup(group);

ObservableList<ChatMessage> chatHistory;

if (bibDatabaseContext.isEmpty() || bibDatabaseContext.get().getDatabasePath().isEmpty()) {
if (bibDatabaseContext.getDatabasePath().isEmpty()) {
chatHistory = FXCollections.observableArrayList();
} else {
List<ChatMessage> chatMessagesList = implementation.loadMessagesForGroup(
bibDatabaseContext.get().getDatabasePath().get(),
bibDatabaseContext.getDatabasePath().get(),
group.getGroup().getName()
);

chatHistory = FXCollections.observableArrayList(chatMessagesList);
}

return new ChatHistoryManagementRecord(bibDatabaseContext, chatHistory);
return new ChatHistoryManagementRecord(Optional.of(bibDatabaseContext), chatHistory);
}).chatHistory;
}

Expand Down Expand Up @@ -235,26 +211,6 @@ private void tryToGenerateCitationKey(BibDatabaseContext bibDatabaseContext, Bib
new CitationKeyGenerator(bibDatabaseContext, citationKeyPatternPreferences).generateAndSetKey(bibEntry);
}

private Optional<BibDatabaseContext> findBibDatabaseForEntry(BibEntry entry) {
return stateManager
.getOpenDatabases()
.stream()
.filter(dbContext -> dbContext.getDatabase().getEntries().contains(entry))
.findFirst();
}

private Optional<BibDatabaseContext> findBibDatabaseForGroup(GroupTreeNode group) {
return stateManager
.getOpenDatabases()
.stream()
.filter(dbContext ->
dbContext.getMetaData().groupsBinding().get().map(groupTreeNode ->
groupTreeNode.containsGroup(group.getGroup())
).orElse(false)
)
.findFirst();
}

@Override
public void close() {
// We need to clone `bibEntriesChatHistory.keySet()` because closeChatHistoryForEntry() modifies the `bibEntriesChatHistory` map.
Expand All @@ -264,6 +220,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) {
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/org/jabref/gui/entryeditor/AiChatTab.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,7 +135,7 @@ private void bindToCorrectEntry(BibEntry entry) {

setContent(new AiChatGuardedComponent(
chatName,
chatHistoryService.getChatHistoryForEntry(entry),
chatHistoryService.getChatHistoryForEntry(bibDatabaseContext, entry),
bibDatabaseContext,
FXCollections.observableArrayList(new ArrayList<>(List.of(entry))),
aiService,
Expand Down
6 changes: 6 additions & 0 deletions src/main/java/org/jabref/gui/frame/JabRefFrame.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,7 @@ public JabRefFrame(Stage mainStage,
this.viewModel = new JabRefFrameViewModel(
preferences,
aiService,
chatHistoryService,
stateManager,
dialogService,
this,
Expand Down Expand Up @@ -157,6 +158,7 @@ public JabRefFrame(Stage mainStage,
this.sidePane = new SidePane(
this,
this.preferences,
aiService,
chatHistoryService,
Injector.instantiateModelOrService(JournalAbbreviationRepository.class),
taskExecutor,
Expand Down Expand Up @@ -200,6 +202,7 @@ private void initLayout() {
stateManager,
preferences,
aiService,
chatHistoryService,
fileUpdateMonitor,
taskExecutor,
entryTypesManager,
Expand Down Expand Up @@ -437,6 +440,8 @@ public void addTab(@NonNull BibDatabaseContext databaseContext, boolean raisePan
databaseContext,
this,
dialogService,
aiService,
chatHistoryService,
preferences,
stateManager,
fileUpdateMonitor,
Expand Down Expand Up @@ -514,6 +519,7 @@ private OpenDatabaseAction getOpenDatabaseAction() {
this,
preferences,
aiService,
chatHistoryService,
dialogService,
stateManager,
fileUpdateMonitor,
Expand Down
5 changes: 5 additions & 0 deletions src/main/java/org/jabref/gui/frame/JabRefFrameViewModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
import org.jabref.gui.LibraryTab;
import org.jabref.gui.LibraryTabContainer;
import org.jabref.gui.StateManager;
import org.jabref.gui.ai.chatting.chathistory.ChatHistoryService;
import org.jabref.gui.externalfiles.AutoLinkFilesAction;
import org.jabref.gui.importer.ImportEntriesDialog;
import org.jabref.gui.importer.ParserResultWarningDialog;
Expand Down Expand Up @@ -50,6 +51,7 @@ public class JabRefFrameViewModel implements UiMessageHandler {

private final GuiPreferences preferences;
private final AiService aiService;
private final ChatHistoryService chatHistoryService;
private final StateManager stateManager;
private final DialogService dialogService;
private final LibraryTabContainer tabContainer;
Expand All @@ -61,6 +63,7 @@ public class JabRefFrameViewModel implements UiMessageHandler {

public JabRefFrameViewModel(GuiPreferences preferences,
AiService aiService,
ChatHistoryService chatHistoryService,
StateManager stateManager,
DialogService dialogService,
LibraryTabContainer tabContainer,
Expand All @@ -71,6 +74,7 @@ public JabRefFrameViewModel(GuiPreferences preferences,
TaskExecutor taskExecutor) {
this.preferences = preferences;
this.aiService = aiService;
this.chatHistoryService = chatHistoryService;
this.stateManager = stateManager;
this.dialogService = dialogService;
this.tabContainer = tabContainer;
Expand Down Expand Up @@ -214,6 +218,7 @@ private void openDatabases(List<ParserResult> parserResults) {
dialogService,
preferences,
aiService,
chatHistoryService,
stateManager,
entryTypesManager,
fileUpdateMonitor,
Expand Down
6 changes: 5 additions & 1 deletion src/main/java/org/jabref/gui/frame/MainToolBar.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
import org.jabref.gui.actions.ActionFactory;
import org.jabref.gui.actions.ActionHelper;
import org.jabref.gui.actions.StandardActions;
import org.jabref.gui.ai.chatting.chathistory.ChatHistoryService;
import org.jabref.gui.bibtexextractor.ExtractBibtexActionOnline;
import org.jabref.gui.citationkeypattern.GenerateCitationKeyAction;
import org.jabref.gui.cleanup.CleanupAction;
Expand Down Expand Up @@ -58,6 +59,7 @@ public class MainToolBar extends ToolBar {
private final StateManager stateManager;
private final GuiPreferences preferences;
private final AiService aiService;
private final ChatHistoryService chatHistoryService;
private final FileUpdateMonitor fileUpdateMonitor;
private final TaskExecutor taskExecutor;
private final BibEntryTypesManager entryTypesManager;
Expand All @@ -75,6 +77,7 @@ public MainToolBar(LibraryTabContainer tabContainer,
StateManager stateManager,
GuiPreferences preferences,
AiService aiService,
ChatHistoryService chatHistoryService,
FileUpdateMonitor fileUpdateMonitor,
TaskExecutor taskExecutor,
BibEntryTypesManager entryTypesManager,
Expand All @@ -87,6 +90,7 @@ public MainToolBar(LibraryTabContainer tabContainer,
this.stateManager = stateManager;
this.preferences = preferences;
this.aiService = aiService;
this.chatHistoryService = chatHistoryService;
this.fileUpdateMonitor = fileUpdateMonitor;
this.taskExecutor = taskExecutor;
this.entryTypesManager = entryTypesManager;
Expand Down Expand Up @@ -114,7 +118,7 @@ private void createToolBar() {
getItems().addAll(
new HBox(
factory.createIconButton(StandardActions.NEW_LIBRARY, new NewDatabaseAction(frame, preferences)),
factory.createIconButton(StandardActions.OPEN_LIBRARY, new OpenDatabaseAction(frame, preferences, aiService, dialogService, stateManager, fileUpdateMonitor, entryTypesManager, undoManager, clipBoardManager, taskExecutor)),
factory.createIconButton(StandardActions.OPEN_LIBRARY, new OpenDatabaseAction(frame, preferences, aiService, chatHistoryService, dialogService, stateManager, fileUpdateMonitor, entryTypesManager, undoManager, clipBoardManager, taskExecutor)),
factory.createIconButton(StandardActions.SAVE_LIBRARY, new SaveAction(SaveAction.SaveMethod.SAVE, frame::getCurrentLibraryTab, dialogService, preferences, stateManager))),

leftSpacer,
Expand Down
2 changes: 1 addition & 1 deletion src/main/java/org/jabref/gui/groups/GroupTreeViewModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -399,7 +399,7 @@ public void chatWithGroup(GroupNodeViewModel group) {
StringProperty nameProperty = new SimpleStringProperty(Localization.lang("Group %0", groupNameProperty.get()));
groupNameProperty.addListener((obs, oldValue, newValue) -> nameProperty.setValue(Localization.lang("Group %0", groupNameProperty.get())));

ObservableList<ChatMessage> chatHistory = chatHistoryService.getChatHistoryForGroup(group.getGroupNode());
ObservableList<ChatMessage> chatHistory = chatHistoryService.getChatHistoryForGroup(currentDatabase.get(), group.getGroupNode());
ObservableList<BibEntry> bibEntries = FXCollections.observableArrayList(group.getGroupNode().findMatches(currentDatabase.get().getDatabase()));

openAiChat(nameProperty, chatHistory, currentDatabase.get(), bibEntries);
Expand Down
Loading
Loading