Clean saving chat history

src/main/java/module-info.java

@@ -154,4 +154,7 @@
  requires de.saxsys.mvvmfx.validation;
  requires dd.plist;
  requires mslinks;
+ requires langchain4j.embeddings.all.minilm.l6.v2;
+ requires langchain4j.core;
+ requires langchain4j.open.ai;
 }

src/main/java/org/jabref/gui/entryeditor/aichattab/ChatMessageComponent.java

@@ -68,9 +68,9 @@ public void setMessage(ChatMessage chatMessage) {
  }
 
  public void setError(String message) {
- contentPane.setStyle("-fx-background-color: -jr-red");
  contentPane.getChildren().clear();
  contentPane.setStyle("");
+ contentPane.setStyle("-fx-background-color: -jr-red");
 
  VBox paneVBox = new VBox(10);
  paneVBox.setMaxWidth(500);

src/main/java/org/jabref/gui/exporter/SaveDatabaseAction.java

@@ -8,6 +8,7 @@
 import java.nio.file.Path;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -25,11 +26,13 @@
 import org.jabref.gui.autosaveandbackup.BackupManager;
 import org.jabref.gui.importer.actions.loadchathistory.AiChatFile;
 import org.jabref.gui.importer.actions.loadchathistory.AiChatFileMessage;
-import org.jabref.gui.importer.actions.loadchathistory.LoadChatHistoryAction;
+import org.jabref.gui.importer.actions.LoadChatHistoryAction;
 import org.jabref.gui.maintable.BibEntryTableViewModel;
 import org.jabref.gui.maintable.columns.MainTableColumn;
 import org.jabref.gui.util.BackgroundTask;
 import org.jabref.gui.util.FileDialogConfiguration;
+import org.jabref.logic.ai.AiChatsFile;
+import org.jabref.logic.ai.ChatMessage;
 import org.jabref.logic.exporter.AtomicFileWriter;
 import org.jabref.logic.exporter.BibDatabaseWriter;
 import org.jabref.logic.exporter.BibWriter;
@@ -296,7 +299,7 @@ private boolean saveDatabase(Path file, boolean selectedOnly, Charset encoding,
  // TODO: What to do with selectedOnly?
 
  ObjectMapper objectMapper = new ObjectMapper();
- AiChatFile aiChatFile = makeAiChatFile(bibDatabaseContext.getDatabase());
+ AiChatsFile aiChatFile = makeAiChatsFile(bibDatabaseContext.getDatabase());
  objectMapper.writeValue(fileWriter, aiChatFile);
 
  if (fileWriter.hasEncodingProblems()) {
@@ -312,19 +315,18 @@ private boolean saveDatabase(Path file, boolean selectedOnly, Charset encoding,
  }
  }
 
- private AiChatFile makeAiChatFile(BibDatabase bibDatabase) {
- AiChatFile aiChatFile = new AiChatFile();
- aiChatFile.chatHistoryMap = new HashMap<>();
+ private AiChatsFile makeAiChatsFile(BibDatabase bibDatabase) {
+ Map<String, List<ChatMessage>> chatHistoryMap = new HashMap<>();
 
  bibDatabase.getEntries().forEach(entry ->
  entry.getCitationKey().ifPresent(citationKey -> {
- List<AiChatFileMessage> aiChatFileMessages = entry.getAiChatMessages().stream().map(AiChatFileMessage::fromLangchain).toList();
- aiChatFile.chatHistoryMap.put(citationKey, aiChatFileMessages);
+ List<ChatMessage> aiChatFileMessages = entry.getAiChatMessages();
+ chatHistoryMap.put(citationKey, aiChatFileMessages);
  }
  )
  );
 
- return aiChatFile;
+ return new AiChatsFile(chatHistoryMap);
  }
 
  private void saveWithDifferentEncoding(Path file, boolean selectedOnly, Charset encoding, Set<Character> encodingProblems, BibDatabaseWriter.SaveType saveType, SelfContainedSaveOrder saveOrder) throws SaveException {

src/main/java/org/jabref/gui/importer/actions/loadchathistory/LoadChatHistoryAction.java → src/main/java/org/jabref/gui/importer/actions/LoadChatHistoryAction.java

@@ -1,26 +1,23 @@
-package org.jabref.gui.importer.actions.loadchathistory;
+package org.jabref.gui.importer.actions;
 
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.io.InputStream;
-import java.io.InputStreamReader;
 import java.nio.file.Path;
-import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 
 import org.jabref.gui.DialogService;
-import org.jabref.gui.importer.actions.GUIPostOpenAction;
+import org.jabref.logic.ai.AiChatsFile;
+import org.jabref.logic.ai.ChatMessage;
 import org.jabref.logic.importer.ParserResult;
 import org.jabref.logic.util.io.FileUtil;
 import org.jabref.model.database.BibDatabase;
 import org.jabref.model.entry.BibEntry;
 
 import com.fasterxml.jackson.databind.ObjectMapper;
-import dev.langchain4j.data.message.ChatMessage;
-import dev.langchain4j.data.message.UserMessage;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -45,17 +42,17 @@ private void readAiChatsFile(BibDatabase bibDatabase, Path databasePath) {
  try {
  InputStream inputStream = new FileInputStream(chatFile);
  ObjectMapper objectMapper = new ObjectMapper();
- AiChatFile aiChatFile = objectMapper.readValue(inputStream, AiChatFile.class);
+ AiChatsFile aiChatFile = objectMapper.readValue(inputStream, AiChatsFile.class);
  loadAiChatHistory(bibDatabase, aiChatFile);
  } catch (FileNotFoundException e) {
- LOGGER.info("There is no .aichats file for the opened library.");
+ LOGGER.info("There is no " + AI_CHAT_HISTORY_EXTENSION + " file for the opened library.");
  } catch (IOException e) {
  LOGGER.error("An error occurred while reading " + chatPath, e);
  }
  }
 
- private void loadAiChatHistory(BibDatabase bibDatabase, AiChatFile aiChatFile) {
- aiChatFile.chatHistoryMap.forEach((citationKey, chatHistory) -> {
+ private void loadAiChatHistory(BibDatabase bibDatabase, AiChatsFile aiChatFile) {
+ aiChatFile.getChatHistoryMap().forEach((citationKey, chatHistory) -> {
  List<BibEntry> bibEntries = bibDatabase.getEntriesByCitationKey(citationKey);
 
  if (bibEntries.isEmpty()) {
@@ -64,14 +61,9 @@ private void loadAiChatHistory(BibDatabase bibDatabase, AiChatFile aiChatFile) {
  LOGGER.warn("Found a chat history for an bib entry with citation key \"" + citationKey + "\" but there are several bib entries in the database with the same key");
  } else {
  BibEntry bibEntry = bibEntries.getFirst();
- List<ChatMessage> parsedChatMessages = parseChatMessages(chatHistory);
- bibEntry.getAiChatMessages().addAll(parsedChatMessages);
+ bibEntry.getAiChatMessages().addAll(chatHistory);
  }
  });
  }
-
- private List<ChatMessage> parseChatMessages(List<AiChatFileMessage> chatHistory) {
- return chatHistory.stream().map(AiChatFileMessage::toLangchainMessage).filter(Optional::isPresent).map(Optional::get).toList();
- }
 }
 

src/main/java/org/jabref/gui/importer/actions/OpenDatabaseAction.java

@@ -20,7 +20,6 @@
 import org.jabref.gui.actions.SimpleCommand;
 import org.jabref.gui.autosaveandbackup.BackupManager;
 import org.jabref.gui.dialogs.BackupUIManager;
-import org.jabref.gui.importer.actions.loadchathistory.LoadChatHistoryAction;
 import org.jabref.gui.shared.SharedDatabaseUIManager;
 import org.jabref.gui.telemetry.Telemetry;
 import org.jabref.gui.undo.CountingUndoManager;

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

@@ -1,11 +1,11 @@
 package org.jabref.logic.ai;
 
 import java.util.List;
+import java.util.Optional;
 import java.util.UUID;
 
 import dev.langchain4j.chain.ConversationalRetrievalChain;
 import dev.langchain4j.data.message.AiMessage;
-import dev.langchain4j.data.message.ChatMessage;
 import dev.langchain4j.data.message.SystemMessage;
 import dev.langchain4j.data.message.UserMessage;
 import dev.langchain4j.data.segment.TextSegment;
@@ -70,6 +70,6 @@ public String execute(String prompt) {
  }
 
  public void restoreMessages(List<ChatMessage> messages) {
- messages.forEach(this.chatMemory::add);
+ messages.stream().map(ChatMessage::toLangchainMessage).filter(Optional::isPresent).map(Optional::get).forEach(this.chatMemory::add);
  }
 }

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

@@ -0,0 +1,20 @@
+package org.jabref.logic.ai;
+
+import java.util.List;
+import java.util.Map;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+public class AiChatsFile {
+ private final Map<String, List<ChatMessage>> chatHistoryMap;
+
+ @JsonCreator
+ public AiChatsFile(@JsonProperty("chatHistoryMap") Map<String, List<ChatMessage>> chatHistoryMap) {
+ this.chatHistoryMap = chatHistoryMap;
+ }
+
+ public Map<String, List<ChatMessage>> getChatHistoryMap() {
+ return chatHistoryMap;
+ }
+}

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

@@ -1,13 +1,9 @@
 package org.jabref.logic.ai;
 
-import java.io.FileNotFoundException;
-import java.io.IOException;
 import java.io.StringWriter;
 import java.nio.file.Path;
 import java.util.Optional;
 
-import org.jabref.logic.JabRefException;
-import org.jabref.logic.l10n.Localization;
 import org.jabref.logic.util.io.FileUtil;
 import org.jabref.logic.xmp.XmpUtilReader;
 import org.jabref.model.database.BibDatabaseContext;

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

@@ -4,8 +4,6 @@
 import dev.langchain4j.model.embedding.AllMiniLmL6V2EmbeddingModel;
 import dev.langchain4j.model.embedding.EmbeddingModel;
 import dev.langchain4j.model.openai.OpenAiChatModel;
-import dev.langchain4j.store.memory.chat.ChatMemoryStore;
-import dev.langchain4j.store.memory.chat.InMemoryChatMemoryStore;
 
 /**
  * This class maintains the connection to AI services.

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

@@ -0,0 +1,62 @@
+package org.jabref.logic.ai;
+
+import java.util.Optional;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import dev.langchain4j.data.message.AiMessage;
+import dev.langchain4j.data.message.UserMessage;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class ChatMessage {
+ private static final Logger LOGGER = LoggerFactory.getLogger(ChatMessage.class);
+
+ private final ChatMessageType type;
+ private final String content;
+
+ @JsonCreator
+ public ChatMessage(@JsonProperty("type") ChatMessageType type,
+ @JsonProperty("content") String content) {
+ this.type = type;
+ this.content = content;
+ }
+
+ public ChatMessageType getType() {
+ return type;
+ }
+
+ public String getContent() {
+ return content;
+ }
+
+ public static Optional<ChatMessage> fromLangchain(dev.langchain4j.data.message.ChatMessage chatMessage) {
+ switch (chatMessage) {
+ case UserMessage userMessage -> {
+ return Optional.of(new ChatMessage(ChatMessageType.USER, userMessage.singleText()));
+ }
+ case AiMessage aiMessage -> {
+ return Optional.of(new ChatMessage(ChatMessageType.ASSISTANT, aiMessage.text()));
+ }
+ default -> {
+ LOGGER.error("Unable to convert langchain4j chat message to JabRef chat message, the type is {}", chatMessage.getClass());
+ return Optional.empty();
+ }
+ }
+ }
+
+ public Optional<dev.langchain4j.data.message.ChatMessage> toLangchainMessage() {
+ switch (type) {
+ case ChatMessageType.USER -> {
+ return Optional.of(new UserMessage(content));
+ }
+ case ChatMessageType.ASSISTANT -> {
+ return Optional.of(new AiMessage(content));
+ }
+ default -> {
+ LOGGER.error("Unable to convert JabRef chat message to langchain4j chat message, the type is {}", type);
+ return Optional.empty();
+ }
+ }
+ }
+}

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

@@ -0,0 +1,6 @@
+package org.jabref.logic.ai;
+
+public enum ChatMessageType {
+ USER,
+ ASSISTANT,
+}

src/main/java/org/jabref/model/entry/BibEntry.java

@@ -18,17 +18,14 @@
 import java.util.function.BiFunction;
 import java.util.stream.Collectors;
 
-import dev.langchain4j.data.message.ChatMessage;
-import dev.langchain4j.data.segment.TextSegment;
-import dev.langchain4j.store.embedding.EmbeddingStore;
-import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
 import javafx.beans.Observable;
 import javafx.beans.property.ObjectProperty;
 import javafx.beans.property.SimpleObjectProperty;
 import javafx.collections.FXCollections;
 import javafx.collections.ObservableMap;
 
 import org.jabref.architecture.AllowedToUseLogic;
+import org.jabref.logic.ai.ChatMessage;
 import org.jabref.logic.bibtex.FileFieldWriter;
 import org.jabref.logic.importer.util.FileFieldParser;
 import org.jabref.model.FieldChange;