Skip to content

Commit

Permalink
Improve AiChatTab
Browse files Browse the repository at this point in the history 
- Use existing methods (.getDatabase().getNumberOfCitationKeyOccurrences)
- Fix method names
- Use Optional
- Enable chatting if at least one PDF file is contained
  • Loading branch information
@koppor
koppor committed Jul 4, 2024
1 parent 49dce23 commit 7760b7f
Showing 1 changed file with 25 additions and 42 deletions.
67 changes: 25 additions & 42 deletions src/main/java/org/jabref/gui/entryeditor/AiChatTab.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,18 +42,18 @@ public class AiChatTab extends EntryEditorTab {
private final FilePreferences filePreferences;
private final WorkspacePreferences workspacePreferences;
private final EntryEditorPreferences entryEditorPreferences;
private final CitationKeyPatternPreferences citationKeyPatternPreferences;
private final BibDatabaseContext bibDatabaseContext;
private final EmbeddingsGenerationTaskManager embeddingsGenerationTaskManager;
private final TaskExecutor taskExecutor;
private final CitationKeyGenerator citationKeyGenerator;

private AiChatComponent aiChatComponent = null;

private final AiService aiService;

private AiChatLogic aiChatLogic = null;

private BibEntry currentBibEntry = null;
private Optional<BibEntry> currentBibEntry = Optional.empty();

private @Nullable BibEntryChatHistory bibEntryChatHistory;

Expand All @@ -63,28 +63,22 @@ public AiChatTab(DialogService dialogService, PreferencesService preferencesServ
this.workspacePreferences = preferencesService.getWorkspacePreferences();
this.filePreferences = preferencesService.getFilePreferences();
this.entryEditorPreferences = preferencesService.getEntryEditorPreferences();
this.citationKeyPatternPreferences = preferencesService.getCitationKeyPatternPreferences();

this.aiService = aiService;

this.bibDatabaseContext = bibDatabaseContext;

this.embeddingsGenerationTaskManager = embeddingsGenerationTaskManager;

this.taskExecutor = taskExecutor;

this.citationKeyGenerator = new CitationKeyGenerator(bibDatabaseContext, preferencesService.getCitationKeyPatternPreferences());
setText(Localization.lang("AI chat"));
setTooltip(new Tooltip(Localization.lang("AI chat with full-text article")));

aiService.getEmbeddingsManager().getIngestedFilesTracker().registerListener(new FileIngestedListener());
}

private class FileIngestedListener {
@Subscribe
public void listen(FileIngestedEvent event) {
if (currentBibEntry != null) {
if (aiService.getEmbeddingsManager().getIngestedFilesTracker().haveIngestedLinkedFiles(currentBibEntry.getFiles())) {
UiTaskExecutor.runInJavaFXThread(() -> bindToEntry(currentBibEntry));
currentBibEntry.ifPresent(entry -> {
if (aiService.getEmbeddingsManager().getIngestedFilesTracker().haveIngestedLinkedFiles(entry.getFiles())) {
UiTaskExecutor.runInJavaFXThread(() -> bindToEntry(entry));
}
}
}
Expand All @@ -97,20 +91,25 @@ public boolean shouldShow(BibEntry entry) {

@Override
protected void bindToEntry(BibEntry entry) {
currentBibEntry = entry;
currentBibEntry = Optional.of(entry);

if (!aiService.getPreferences().getEnableChatWithFiles()) {
showPrivacyNotice(entry);
} else if (!checkIfCitationKeyIsAppropriate(bibDatabaseContext, entry)) {
tryToGenerateCitationKey(entry);
} else if (entry.getFiles().isEmpty()) {
showErrorNoFiles();
} else if (!entry.getFiles().stream().map(LinkedFile::getLink).map(Path::of).allMatch(FileUtil::isPDFFile)) {
} else if (!entry.getFiles().stream().map(LinkedFile::getLink).map(Path::of).anyMatch(FileUtil::isPDFFile)) {
showErrorNotPdfs();
} else if (!aiService.getEmbeddingsManager().getIngestedFilesTracker().haveIngestedLinkedFiles(currentBibEntry.getFiles())) {
} else if (!citationKeyIsValid(bibDatabaseContext, entry)) {
tryToGenerateCitationKeyThenBind(entry);
} else if (!aiService.getEmbeddingsManager().getIngestedFilesTracker().haveIngestedLinkedFiles(entry.getFiles())) {
showErrorNotIngested();
} else {
bindToCorrectEntry(entry);
// All preconditions met
embeddingsGenerationTaskManager.moveToFront(entry.getFiles());
createAiChat();
setupChatHistory();
restoreLogicalChatHistory();
buildChatUI();
}
}

Expand All @@ -126,8 +125,7 @@ private void showErrorNoFiles() {
setContent(new ErrorStateComponent(Localization.lang("Unable to chat"), Localization.lang("Please attach at least one PDF file to enable chatting with PDF files.")));
}

private void tryToGenerateCitationKey(BibEntry entry) {
CitationKeyGenerator citationKeyGenerator = new CitationKeyGenerator(bibDatabaseContext, citationKeyPatternPreferences);
private void tryToGenerateCitationKeyThenBind(BibEntry entry) {
if (citationKeyGenerator.generateAndSetKey(entry).isEmpty()) {
setContent(new ErrorStateComponent(Localization.lang("Unable to chat"), Localization.lang("Please provide a non-empty and unique citation key for this entry.")));
} else {
Expand All @@ -141,31 +139,16 @@ private void showPrivacyNotice(BibEntry entry) {
}));
}

private static boolean checkIfCitationKeyIsAppropriate(BibDatabaseContext bibDatabaseContext, BibEntry bibEntry) {
// bibEntry.getCitationKey().isPresent() is called to pleasure the linter (even if it is already checked in checkIfCitationKeyIsEmpty).
return !checkIfCitationKeyIsEmpty(bibEntry) && bibEntry.getCitationKey().isPresent() && checkIfCitationKeyIsUnique(bibDatabaseContext, bibEntry.getCitationKey().get());
}

private static boolean checkIfCitationKeyIsEmpty(BibEntry bibEntry) {
return bibEntry.getCitationKey().isEmpty() || bibEntry.getCitationKey().get().isEmpty();
private static boolean citationKeyIsValid(BibDatabaseContext bibDatabaseContext, BibEntry bibEntry) {
return !hasEmptyCitationKey(bibEntry) && bibEntry.getCitationKey().map(key -> citationKeyIsUnique(bibDatabaseContext, key)).orElse(false);
}

private static boolean checkIfCitationKeyIsUnique(BibDatabaseContext bibDatabaseContext, String citationKey) {
return bibDatabaseContext.getDatabase().getEntries().stream()
.map(BibEntry::getCitationKey)
.filter(Optional::isPresent)
.map(Optional::get)
.filter(key -> key.equals(citationKey))
.count() == 1;
private static boolean hasEmptyCitationKey(BibEntry bibEntry) {
return bibEntry.getCitationKey().map(key -> key.isEmpty()).orElse(true);
}

private void bindToCorrectEntry(BibEntry entry) {
embeddingsGenerationTaskManager.moveToFront(entry.getFiles());

createAiChat();
setupChatHistory();
restoreLogicalChatHistory();
buildChatUI(entry);
private static boolean citationKeyIsUnique(BibDatabaseContext bibDatabaseContext, String citationKey) {
return bibDatabaseContext.getDatabase().getNumberOfCitationKeyOccurrences(citationKey) == 1;
}

private void createAiChat() {
Expand All @@ -187,7 +170,7 @@ private void restoreLogicalChatHistory() {
}
}

private void buildChatUI(BibEntry entry) {
private void buildChatUI() {
aiChatComponent = new AiChatComponent(userPrompt -> {
ChatMessage userMessage = ChatMessage.user(userPrompt);
aiChatComponent.addMessage(userMessage);
Expand Down

0 comments on commit 7760b7f

Please sign in to comment.