Customizable AI templates

build.gradle

@@ -351,7 +351,7 @@ dependencies {
         exclude group: 'org.jetbrains.kotlin'
     }
 
-
+    implementation 'org.apache.velocity:velocity-engine-core:2.3'
     implementation platform('ai.djl:bom:0.30.0')
     implementation 'ai.djl:api'
     implementation 'ai.djl.huggingface:tokenizers'

src/main/java/module-info.java

@@ -160,6 +160,7 @@
     uses ai.djl.repository.RepositoryFactory;
     uses ai.djl.repository.zoo.ZooProvider;
     uses dev.langchain4j.spi.prompt.PromptTemplateFactory;
+    requires velocity.engine.core;
     // endregion
 
     // region: Lucene

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

@@ -162,10 +162,6 @@
             </children>
         </HBox>
 
-        <ResizableTextArea
-                fx:id="instructionTextArea"
-                wrapText="true"/>
-
         <GridPane hgap="10" vgap="10">
             <columnConstraints>
                 <ColumnConstraints hgrow="ALWAYS" percentWidth="50" />
@@ -235,5 +231,31 @@
                         glyph="REFRESH"/>
             </graphic>
         </Button>
+
+        <HBox alignment="BASELINE_CENTER">
+            <Label styleClass="sectionHeader"
+                   text="%Templates"
+                   maxWidth="Infinity"
+                   HBox.hgrow="ALWAYS"/>
+            <Button fx:id="templatesHelp"
+                    prefWidth="20.0"/>
+        </HBox>
+
+        <ComboBox
+                fx:id="currentEditingTemplateComboBox"
+                maxWidth="1.7976931348623157E308"
+                HBox.hgrow="ALWAYS"/>
+
+        <ResizableTextArea
+                fx:id="currentEditingTemplateSourceTextArea"
+                wrapText="true"/>
+
+        <Button onAction="#onResetTemplatesButtonClick"
+                text="%Reset templates to default">
+            <graphic>
+                <JabRefIconView
+                        glyph="REFRESH"/>
+            </graphic>
+        </Button>
     </children>
 </fx:root>

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

@@ -15,6 +15,7 @@
 import org.jabref.gui.preferences.AbstractPreferenceTabView;
 import org.jabref.gui.preferences.PreferencesTab;
 import org.jabref.gui.util.ViewModelListCellFactory;
+import org.jabref.logic.ai.templates.AiTemplate;
 import org.jabref.logic.help.HelpFile;
 import org.jabref.logic.l10n.Localization;
 import org.jabref.model.ai.AiProvider;
@@ -43,16 +44,19 @@ public class AiTab extends AbstractPreferenceTabView<AiTabViewModel> implements
 
     @FXML private TextField apiBaseUrlTextField;
     @FXML private SearchableComboBox<EmbeddingModel> embeddingModelComboBox;
-    @FXML private ResizableTextArea instructionTextArea;
     @FXML private TextField temperatureTextField;
     @FXML private IntegerInputField contextWindowSizeTextField;
     @FXML private IntegerInputField documentSplitterChunkSizeTextField;
     @FXML private IntegerInputField documentSplitterOverlapSizeTextField;
     @FXML private IntegerInputField ragMaxResultsCountTextField;
     @FXML private TextField ragMinScoreTextField;
 
+    @FXML private ComboBox<AiTemplate> currentEditingTemplateComboBox;
+    @FXML private ResizableTextArea currentEditingTemplateSourceTextArea;
+
     @FXML private Button generalSettingsHelp;
     @FXML private Button expertSettingsHelp;
+    @FXML private Button templatesHelp;
 
     private final ControlsFxVisualizer visualizer = new ControlsFxVisualizer();
 
@@ -74,14 +78,14 @@ public void initialize() {
         new ViewModelListCellFactory<AiProvider>()
                 .withText(AiProvider::toString)
                 .install(aiProviderComboBox);
-        aiProviderComboBox.setItems(viewModel.aiProvidersProperty());
+        aiProviderComboBox.itemsProperty().bind(viewModel.aiProvidersProperty());
         aiProviderComboBox.valueProperty().bindBidirectional(viewModel.selectedAiProviderProperty());
         aiProviderComboBox.disableProperty().bind(viewModel.disableBasicSettingsProperty());
 
         new ViewModelListCellFactory<String>()
                 .withText(text -> text)
                 .install(chatModelComboBox);
-        chatModelComboBox.setItems(viewModel.chatModelsProperty());
+        chatModelComboBox.itemsProperty().bind(viewModel.chatModelsProperty());
         chatModelComboBox.valueProperty().bindBidirectional(viewModel.selectedChatModelProperty());
         chatModelComboBox.disableProperty().bind(viewModel.disableBasicSettingsProperty());
 
@@ -123,9 +127,6 @@ public void initialize() {
             apiBaseUrlTextField.setDisable(newValue || viewModel.disableExpertSettingsProperty().get())
         );
 
-        instructionTextArea.textProperty().bindBidirectional(viewModel.instructionProperty());
-        instructionTextArea.disableProperty().bind(viewModel.disableExpertSettingsProperty());
-
         // bindBidirectional doesn't work well with number input fields ({@link IntegerInputField}, {@link DoubleInputField}),
         // so they are expanded into `addListener` calls.
 
@@ -180,7 +181,6 @@ public void initialize() {
             visualizer.initVisualization(viewModel.getChatModelValidationStatus(), chatModelComboBox);
             visualizer.initVisualization(viewModel.getApiBaseUrlValidationStatus(), apiBaseUrlTextField);
             visualizer.initVisualization(viewModel.getEmbeddingModelValidationStatus(), embeddingModelComboBox);
-            visualizer.initVisualization(viewModel.getSystemMessageValidationStatus(), instructionTextArea);
             visualizer.initVisualization(viewModel.getTemperatureTypeValidationStatus(), temperatureTextField);
             visualizer.initVisualization(viewModel.getTemperatureRangeValidationStatus(), temperatureTextField);
             visualizer.initVisualization(viewModel.getMessageWindowSizeValidationStatus(), contextWindowSizeTextField);
@@ -191,9 +191,18 @@ public void initialize() {
             visualizer.initVisualization(viewModel.getRagMinScoreRangeValidationStatus(), ragMinScoreTextField);
         });
 
+        new ViewModelListCellFactory<AiTemplate>()
+                .withText(AiTemplate::getLocalizedName)
+                .install(currentEditingTemplateComboBox);
+        currentEditingTemplateComboBox.itemsProperty().bind(viewModel.templatesProperty());
+        currentEditingTemplateComboBox.valueProperty().bindBidirectional(viewModel.currentEditingTemplate());
+
+        currentEditingTemplateSourceTextArea.textProperty().bindBidirectional(viewModel.currentEditingTemplateSource());
+
         ActionFactory actionFactory = new ActionFactory();
         actionFactory.configureIconButton(StandardActions.HELP, new HelpAction(HelpFile.AI_GENERAL_SETTINGS, dialogService, preferences.getExternalApplicationsPreferences()), generalSettingsHelp);
         actionFactory.configureIconButton(StandardActions.HELP, new HelpAction(HelpFile.AI_EXPERT_SETTINGS, dialogService, preferences.getExternalApplicationsPreferences()), expertSettingsHelp);
+        actionFactory.configureIconButton(StandardActions.HELP, new HelpAction(HelpFile.AI_TEMPLATES, dialogService, preferences.getExternalApplicationsPreferences()), templatesHelp);
     }
 
     @Override
@@ -206,6 +215,11 @@ private void onResetExpertSettingsButtonClick() {
         viewModel.resetExpertSettings();
     }
 
+    @FXML
+    private void onResetTemplatesButtonClick() {
+        viewModel.resetTemplates();
+    }
+
     public ReadOnlyBooleanProperty aiEnabledProperty() {
         return enableAi.selectedProperty();
     }

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

@@ -1,7 +1,9 @@
 package org.jabref.gui.preferences.ai;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Objects;
 
 import javafx.beans.property.BooleanProperty;
@@ -20,6 +22,7 @@
 import org.jabref.gui.preferences.PreferenceTabViewModel;
 import org.jabref.logic.ai.AiDefaultPreferences;
 import org.jabref.logic.ai.AiPreferences;
+import org.jabref.logic.ai.templates.AiTemplate;
 import org.jabref.logic.l10n.Localization;
 import org.jabref.logic.preferences.CliPreferences;
 import org.jabref.logic.util.LocalizedNumbers;
@@ -79,7 +82,18 @@ public class AiTabViewModel implements PreferenceTabViewModel {
     private final StringProperty huggingFaceApiBaseUrl = new SimpleStringProperty();
     private final StringProperty gpt4AllApiBaseUrl = new SimpleStringProperty();
 
-    private final StringProperty instruction = new SimpleStringProperty();
+    private final ListProperty<AiTemplate> templatesList =
+            new SimpleListProperty<>(FXCollections.observableArrayList(AiTemplate.values()));
+    private final ObjectProperty<AiTemplate> currentEditingTemplate = new SimpleObjectProperty<>(AiTemplate.CHATTING_SYSTEM_MESSAGE);
+
+    private final Map<AiTemplate, StringProperty> templateSources = Map.of(
+            AiTemplate.CHATTING_SYSTEM_MESSAGE, new SimpleStringProperty(),
+            AiTemplate.CHATTING_USER_MESSAGE, new SimpleStringProperty(),
+            AiTemplate.SUMMARIZATION_CHUNK, new SimpleStringProperty(),
+            AiTemplate.SUMMARIZATION_COMBINE, new SimpleStringProperty()
+    );
+    private final StringProperty currentEditingTemplateSource = new SimpleStringProperty();
+
     private final StringProperty temperature = new SimpleStringProperty();
     private final IntegerProperty contextWindowSize = new SimpleIntegerProperty();
     private final IntegerProperty documentSplitterChunkSize = new SimpleIntegerProperty();
@@ -96,7 +110,6 @@ public class AiTabViewModel implements PreferenceTabViewModel {
     private final Validator chatModelValidator;
     private final Validator apiBaseUrlValidator;
     private final Validator embeddingModelValidator;
-    private final Validator instructionValidator;
     private final Validator temperatureTypeValidator;
     private final Validator temperatureRangeValidator;
     private final Validator contextWindowSizeValidator;
@@ -222,6 +235,14 @@ public AiTabViewModel(CliPreferences preferences) {
             }
         });
 
+        this.currentEditingTemplateSource.addListener((observable, oldValue, newValue) -> {
+            templateSources.get(currentEditingTemplate.get()).set(newValue);
+        });
+
+        this.currentEditingTemplate.addListener((observable, oldValue, newValue) -> {
+            currentEditingTemplateSource.set(templateSources.get(newValue).get());
+        });
+
         this.apiKeyValidator = new FunctionBasedValidator<>(
                 currentApiKey,
                 token -> !StringUtil.isBlank(token),
@@ -242,11 +263,6 @@ public AiTabViewModel(CliPreferences preferences) {
                 Objects::nonNull,
                 ValidationMessage.error(Localization.lang("Embedding model has to be provided")));
 
-        this.instructionValidator = new FunctionBasedValidator<>(
-                instruction,
-                message -> !StringUtil.isBlank(message),
-                ValidationMessage.error(Localization.lang("The instruction has to be provided")));
-
         this.temperatureTypeValidator = new FunctionBasedValidator<>(
                 temperature,
                 temp -> LocalizedNumbers.stringToDouble(temp).isPresent(),
@@ -318,7 +334,12 @@ public void setValues() {
         customizeExpertSettings.setValue(aiPreferences.getCustomizeExpertSettings());
 
         selectedEmbeddingModel.setValue(aiPreferences.getEmbeddingModel());
-        instruction.setValue(aiPreferences.getInstruction());
+
+        Arrays.stream(AiTemplate.values()).forEach(template ->
+                templateSources.get(template).set(aiPreferences.getTemplate(template)));
+
+        currentEditingTemplateSource.set(templateSources.get(currentEditingTemplate.get()).get());
+
         temperature.setValue(LocalizedNumbers.doubleToString(aiPreferences.getTemperature()));
         contextWindowSize.setValue(aiPreferences.getContextWindowSize());
         documentSplitterChunkSize.setValue(aiPreferences.getDocumentSplitterChunkSize());
@@ -359,7 +380,9 @@ public void storeSettings() {
         aiPreferences.setHuggingFaceApiBaseUrl(huggingFaceApiBaseUrl.get() == null ? "" : huggingFaceApiBaseUrl.get());
         aiPreferences.setGpt4AllApiBaseUrl(gpt4AllApiBaseUrl.get() == null ? "" : gpt4AllApiBaseUrl.get());
 
-        aiPreferences.setInstruction(instruction.get());
+        Arrays.stream(AiTemplate.values()).forEach(template ->
+                aiPreferences.setTemplate(template, templateSources.get(template).get()));
+
         // We already check the correctness of temperature and RAG minimum score in validators, so we don't need to check it here.
         aiPreferences.setTemperature(LocalizedNumbers.stringToDouble(oldLocale, temperature.get()).get());
         aiPreferences.setContextWindowSize(contextWindowSize.get());
@@ -373,8 +396,6 @@ public void resetExpertSettings() {
         String resetApiBaseUrl = selectedAiProvider.get().getApiUrl();
         currentApiBaseUrl.set(resetApiBaseUrl);
 
-        instruction.set(AiDefaultPreferences.SYSTEM_MESSAGE);
-
         contextWindowSize.set(AiDefaultPreferences.getContextWindowSize(selectedAiProvider.get(), currentChatModel.get()));
 
         temperature.set(LocalizedNumbers.doubleToString(AiDefaultPreferences.TEMPERATURE));
@@ -384,6 +405,13 @@ public void resetExpertSettings() {
         ragMinScore.set(LocalizedNumbers.doubleToString(AiDefaultPreferences.RAG_MIN_SCORE));
     }
 
+    public void resetTemplates() {
+        Arrays.stream(AiTemplate.values()).forEach(template ->
+                templateSources.get(template).set(AiDefaultPreferences.TEMPLATES.get(template)));
+
+        currentEditingTemplateSource.set(templateSources.get(currentEditingTemplate.get()).get());
+    }
+
     @Override
     public boolean validateSettings() {
         if (enableAi.get()) {
@@ -410,7 +438,6 @@ public boolean validateExpertSettings() {
         List<Validator> validators = List.of(
                 apiBaseUrlValidator,
                 embeddingModelValidator,
-                instructionValidator,
                 temperatureTypeValidator,
                 temperatureRangeValidator,
                 contextWindowSizeValidator,
@@ -484,8 +511,16 @@ public BooleanProperty disableApiBaseUrlProperty() {
         return disableApiBaseUrl;
     }
 
-    public StringProperty instructionProperty() {
-        return instruction;
+    public ListProperty<AiTemplate> templatesProperty() {
+        return templatesList;
+    }
+
+    public ObjectProperty<AiTemplate> currentEditingTemplate() {
+        return currentEditingTemplate;
+    }
+
+    public StringProperty currentEditingTemplateSource() {
+        return currentEditingTemplateSource;
     }
 
     public StringProperty temperatureProperty() {
@@ -536,10 +571,6 @@ public ValidationStatus getEmbeddingModelValidationStatus() {
         return embeddingModelValidator.getValidationStatus();
     }
 
-    public ValidationStatus getSystemMessageValidationStatus() {
-        return instructionValidator.getValidationStatus();
-    }
-
     public ValidationStatus getTemperatureTypeValidationStatus() {
         return temperatureTypeValidator.getValidationStatus();
     }

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

@@ -4,6 +4,7 @@
 import java.util.List;
 import java.util.Map;
 
+import org.jabref.logic.ai.templates.AiTemplate;
 import org.jabref.model.ai.AiProvider;
 import org.jabref.model.ai.EmbeddingModel;
 
@@ -80,6 +81,46 @@ public String toString() {
 
     public static final int FALLBACK_CONTEXT_WINDOW_SIZE = 8196;
 
+    public static final Map<AiTemplate, String> TEMPLATES = Map.of(
+            AiTemplate.CHATTING_SYSTEM_MESSAGE, """
+                    You are an AI assistant that analyses research papers. You answer questions about papers.
+                    You will be supplied with the necessary information. The supplied information will contain mentions of papers in form '@citationKey'.
+                    Whenever you refer to a paper, use its citation key in the same form with @ symbol. Whenever you find relevant information, always use the citation key.
+
+                    Here are the papers you are analyzing:
+                    #foreach( $entry in $entries )
+                    ${CanonicalBibEntry.getCanonicalRepresentation($entry)}
+                    #end""",
+
+            AiTemplate.CHATTING_USER_MESSAGE, """
+                    $message
+
+                    Here is some relevant information for you:
+                    #foreach( $excerpt in $excerpts )
+                    ${excerpt.citationKey()}:
+                    ${excerpt.text()}
+                    #end""",
+
+            AiTemplate.SUMMARIZATION_CHUNK, """
+                Please provide an overview of the following text. It is a part of a scientific paper.
+                The summary should include the main objectives, methodologies used, key findings, and conclusions.
+                Mention any significant experiments, data, or discussions presented in the paper.
+
+                DOCUMENT:
+                $document
+
+                OVERVIEW:""",
+
+            AiTemplate.SUMMARIZATION_COMBINE, """
+                You have written an overview of a scientific paper. You have been collecting notes from various parts
+                of the paper. Now your task is to combine all of the notes in one structured message.
+
+                SUMMARIES:
+                $summaries
+
+                FINAL OVERVIEW:"""
+    );
+
     public static List<String> getAvailableModels(AiProvider aiProvider) {
         return Arrays.stream(AiDefaultPreferences.PredefinedChatModel.values())
                      .filter(model -> model.getAiProvider() == aiProvider)