Support method name lookup generation (#280)

src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java

@@ -7,6 +7,7 @@ import com.intellij.openapi.project.ProjectManager;
 import com.intellij.openapi.project.ProjectManagerListener;
 import com.intellij.openapi.startup.StartupActivity;
 import ee.carlrobert.codegpt.actions.editor.EditorActionsUtil;
+import ee.carlrobert.codegpt.completions.CompletionLookupService;
 import ee.carlrobert.codegpt.completions.you.YouUserManager;
 import ee.carlrobert.codegpt.completions.you.auth.AuthenticationHandler;
 import ee.carlrobert.codegpt.completions.you.auth.SessionVerificationJob;
@@ -31,6 +32,7 @@ public class PluginStartupActivity implements StartupActivity {
 
   @Override
   public void runActivity(@NotNull Project project) {
+    project.getService(CompletionLookupService.class).subscribeToLookupTopic();
     EditorActionsUtil.refreshActions();
 
     var authenticationResponse = YouUserManager.getInstance().getAuthenticationResponse();

src/main/java/ee/carlrobert/codegpt/completions/CompletionLookupService.java

@@ -0,0 +1,99 @@
+package ee.carlrobert.codegpt.completions;
+
+import com.intellij.codeInsight.completion.PrefixMatcher;
+import com.intellij.codeInsight.lookup.LookupElementBuilder;
+import com.intellij.codeInsight.lookup.LookupManagerListener;
+import com.intellij.codeInsight.lookup.impl.LookupImpl;
+import com.intellij.openapi.application.Application;
+import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.components.Service;
+import com.intellij.openapi.project.Project;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiUtilCore;
+import ee.carlrobert.codegpt.Icons;
+import ee.carlrobert.codegpt.credentials.OpenAICredentialsManager;
+import ee.carlrobert.codegpt.settings.configuration.ConfigurationState;
+import ee.carlrobert.codegpt.settings.state.SettingsState;
+import java.util.Optional;
+import org.jetbrains.annotations.Nullable;
+
+@Service(Service.Level.PROJECT)
+public final class CompletionLookupService {
+
+  private final Project project;
+
+  private CompletionLookupService(Project project) {
+    this.project = project;
+  }
+
+  public void subscribeToLookupTopic() {
+    project.getMessageBus()
+        .connect()
+        .subscribe(LookupManagerListener.TOPIC, getLookupManagerListener());
+  }
+
+  private @Nullable String getCompletionResponse(String prompt) {
+    var selectedService = SettingsState.getInstance().getSelectedService();
+    switch (selectedService) {
+      case OPENAI:
+      case AZURE:
+        return Optional.ofNullable(CompletionClientProvider.getOpenAIClient()
+                .getChatCompletion(
+                    CompletionRequestProvider.buildOpenAILookupCompletionRequest(prompt))
+                .getChoices())
+            .map(choices -> choices.get(0).getMessage().getContent())
+            .orElse(null);
+      // TODO
+      /*case LLAMA_CPP:
+        var request = CompletionRequestProvider.buildLlamaLookupCompletionRequest(prompt);
+        return CompletionClientProvider.getLlamaClient()
+            .getChatCompletion(request)
+            .getContent();*/
+      default:
+        return null;
+    }
+  }
+
+  private void addCompletionLookupValues(
+      LookupImpl lookup,
+      Application application,
+      String prompt) {
+    Optional.ofNullable(getCompletionResponse(prompt))
+        .ifPresent(response -> {
+          for (var value : response.split(",")) {
+            application.runReadAction(() -> {
+              lookup.addItem(
+                  LookupElementBuilder.create(value.trim()).withIcon(Icons.SparkleIcon),
+                  PrefixMatcher.ALWAYS_TRUE);
+            });
+            application.invokeLater(() -> lookup.refreshUi(true, true));
+          }
+        });
+  }
+
+  private LookupManagerListener getLookupManagerListener() {
+    var application = ApplicationManager.getApplication();
+    var configuration = ConfigurationState.getInstance();
+    var credentialsManager = OpenAICredentialsManager.getInstance();
+    return (oldLookup, newLookup) -> {
+      if (!configuration.isMethodNameGenerationEnabled()
+          || !credentialsManager.isApiKeySet()
+          || !(newLookup instanceof LookupImpl)) {
+        return;
+      }
+
+      var lookup = (LookupImpl) newLookup;
+      Optional.ofNullable(lookup.getPsiElement())
+          .map(PsiElement::getContext)
+          .ifPresent(context ->
+              application.runReadAction(() -> {
+                var type = PsiUtilCore.getElementType(context);
+                if ("METHOD".equals(type.toString())) {
+                  var selection = context.getText();
+                  application.executeOnPooledThread(
+                      () -> addCompletionLookupValues(lookup, application, selection));
+                }
+              }));
+    };
+  }
+}

src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java

@@ -1,5 +1,6 @@
 package ee.carlrobert.codegpt.completions;
 
+import static ee.carlrobert.codegpt.util.file.FileUtils.getResourceContent;
 import static java.util.stream.Collectors.toList;
 
 import com.intellij.openapi.application.ApplicationManager;
@@ -7,18 +8,21 @@ import com.intellij.openapi.diagnostic.Logger;
 import ee.carlrobert.codegpt.CodeGPTPlugin;
 import ee.carlrobert.codegpt.EncodingManager;
 import ee.carlrobert.codegpt.completions.llama.LlamaModel;
+import ee.carlrobert.codegpt.completions.llama.PromptTemplate;
 import ee.carlrobert.codegpt.conversations.Conversation;
 import ee.carlrobert.codegpt.conversations.ConversationsState;
 import ee.carlrobert.codegpt.conversations.message.Message;
 import ee.carlrobert.codegpt.settings.configuration.ConfigurationState;
 import ee.carlrobert.codegpt.settings.service.ServiceType;
 import ee.carlrobert.codegpt.settings.state.LlamaSettingsState;
+import ee.carlrobert.codegpt.settings.state.OpenAISettingsState;
 import ee.carlrobert.codegpt.settings.state.SettingsState;
 import ee.carlrobert.codegpt.settings.state.YouSettingsState;
 import ee.carlrobert.codegpt.telemetry.core.configuration.TelemetryConfiguration;
 import ee.carlrobert.codegpt.telemetry.core.service.UserId;
 import ee.carlrobert.embedding.EmbeddingsService;
 import ee.carlrobert.llm.client.llama.completion.LlamaCompletionRequest;
+import ee.carlrobert.llm.client.openai.completion.OpenAICompletionRequest;
 import ee.carlrobert.llm.client.openai.completion.chat.OpenAIChatCompletionModel;
 import ee.carlrobert.llm.client.openai.completion.chat.request.OpenAIChatCompletionMessage;
 import ee.carlrobert.llm.client.openai.completion.chat.request.OpenAIChatCompletionRequest;
@@ -68,6 +72,25 @@ public class CompletionRequestProvider {
     this.conversation = conversation;
   }
 
+  public static OpenAICompletionRequest buildOpenAILookupCompletionRequest(
+      String context) {
+    return new OpenAIChatCompletionRequest.Builder(
+        List.of(
+            new OpenAIChatCompletionMessage("system",
+                getResourceContent("/prompts/method-name-generator.txt")),
+            new OpenAIChatCompletionMessage("user", context)))
+        .setModel(OpenAISettingsState.getInstance().getModel())
+        .setStream(false)
+        .build();
+  }
+
+  public static LlamaCompletionRequest buildLlamaLookupCompletionRequest(String context) {
+    return new LlamaCompletionRequest.Builder(PromptTemplate.LLAMA
+        .buildPrompt(getResourceContent("/prompts/method-name-generator.txt"), context, List.of()))
+        .setStream(false)
+        .build();
+  }
+
   public LlamaCompletionRequest buildLlamaCompletionRequest(Message message) {
     var settings = LlamaSettingsState.getInstance();
     var promptTemplate = settings.isUseCustomModel()

src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java

@@ -10,7 +10,6 @@ import com.intellij.openapi.actionSystem.AnActionEvent;
 import com.intellij.openapi.keymap.impl.ui.EditKeymapsDialog;
 import com.intellij.openapi.ui.ComponentValidator;
 import com.intellij.openapi.ui.ValidationInfo;
-import com.intellij.openapi.util.Disposer;
 import com.intellij.ui.AnActionButton;
 import com.intellij.ui.TitledSeparator;
 import com.intellij.ui.ToolbarDecorator;
@@ -26,8 +25,6 @@ import ee.carlrobert.codegpt.CodeGPTBundle;
 import ee.carlrobert.codegpt.actions.editor.EditorActionsUtil;
 import ee.carlrobert.codegpt.util.SwingUtils;
 import java.awt.Dimension;
-import java.awt.event.KeyEvent;
-import java.awt.event.KeyListener;
 import java.util.Arrays;
 import java.util.LinkedHashMap;
 import java.util.Map;
@@ -45,6 +42,7 @@ public class ConfigurationComponent {
   private final JPanel mainPanel;
   private final JBTable table;
   private final JBCheckBox openNewTabCheckBox;
+  private final JBCheckBox methodNameGenerationCheckBox;
   private final JTextArea systemPromptTextArea;
   private final IntegerField maxTokensField;
   private final JBTextField temperatureField;
@@ -52,7 +50,7 @@ public class ConfigurationComponent {
   public ConfigurationComponent(Disposable parentDisposable, ConfigurationState configuration) {
     table = new JBTable(new DefaultTableModel(
         EditorActionsUtil.toArray(configuration.getTableData()),
-        new String[] {
+        new String[]{
             CodeGPTBundle.get("configurationConfigurable.table.header.actionColumnLabel"),
             CodeGPTBundle.get("configurationConfigurable.table.header.promptColumnLabel")
         }));
@@ -101,12 +99,17 @@ public class ConfigurationComponent {
     systemPromptTextArea.setRows(3);
 
     openNewTabCheckBox = new JBCheckBox(
-        CodeGPTBundle.get("configurationConfigurable.openNewTabCheckBox.label"), false);
+        CodeGPTBundle.get("configurationConfigurable.openNewTabCheckBox.label"),
+        configuration.isCreateNewChatOnEachAction());
+    methodNameGenerationCheckBox = new JBCheckBox(
+        CodeGPTBundle.get("configurationConfigurable.disableMethodNameGeneration.label"),
+        configuration.isMethodNameGenerationEnabled());
 
     mainPanel = FormBuilder.createFormBuilder()
         .addComponent(tablePanel)
         .addVerticalGap(4)
         .addComponent(openNewTabCheckBox)
+        .addComponent(methodNameGenerationCheckBox)
         .addVerticalGap(4)
         .addComponent(new TitledSeparator(
             CodeGPTBundle.get("configurationConfigurable.section.assistant.title")))
@@ -133,7 +136,7 @@ public class ConfigurationComponent {
   private JPanel createTablePanel() {
     return ToolbarDecorator.createDecorator(table)
         .setPreferredSize(new Dimension(table.getPreferredSize().width, 140))
-        .setAddAction(anActionButton -> getModel().addRow(new Object[] {"", ""}))
+        .setAddAction(anActionButton -> getModel().addRow(new Object[]{"", ""}))
         .setRemoveAction(anActionButton -> getModel().removeRow(table.getSelectedRow()))
         .disableUpAction()
         .disableDownAction()
@@ -221,7 +224,7 @@ public class ConfigurationComponent {
   public void setTableData(Map<String, String> tableData) {
     var model = getModel();
     model.setNumRows(0);
-    tableData.forEach((action, prompt) -> model.addRow(new Object[] {action, prompt}));
+    tableData.forEach((action, prompt) -> model.addRow(new Object[]{action, prompt}));
   }
 
   public void setSystemPrompt(String systemPrompt) {
@@ -256,6 +259,14 @@ public class ConfigurationComponent {
     openNewTabCheckBox.setSelected(createNewChatOnEachAction);
   }
 
+  public boolean isMethodNameGenerationEnabled() {
+    return methodNameGenerationCheckBox.isSelected();
+  }
+
+  public void setDisableMethodNameGeneration(boolean disableMethodNameGeneration) {
+    methodNameGenerationCheckBox.setSelected(disableMethodNameGeneration);
+  }
+
   class RevertToDefaultsActionButton extends AnActionButton {
 
     RevertToDefaultsActionButton() {

src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java

@@ -38,7 +38,9 @@ public class ConfigurationConfigurable implements Configurable {
         || configurationComponent.getTemperature() != configuration.getTemperature()
         || !configurationComponent.getSystemPrompt().equals(configuration.getSystemPrompt())
         || configurationComponent.isCreateNewChatOnEachAction()
-        != configuration.isCreateNewChatOnEachAction();
+        != configuration.isCreateNewChatOnEachAction()
+        || configurationComponent.isMethodNameGenerationEnabled()
+        != configuration.isMethodNameGenerationEnabled();
   }
 
   @Override
@@ -50,6 +52,8 @@ public class ConfigurationConfigurable implements Configurable {
     configuration.setSystemPrompt(configurationComponent.getSystemPrompt());
     configuration.setCreateNewChatOnEachAction(
         configurationComponent.isCreateNewChatOnEachAction());
+    configuration.setMethodNameGenerationEnabled(
+        configurationComponent.isMethodNameGenerationEnabled());
     EditorActionsUtil.refreshActions();
   }
 
@@ -62,6 +66,8 @@ public class ConfigurationConfigurable implements Configurable {
     configurationComponent.setSystemPrompt(configuration.getSystemPrompt());
     configurationComponent.setCreateNewChatOnEachAction(
         configuration.isCreateNewChatOnEachAction());
+    configurationComponent.setDisableMethodNameGeneration(
+        configuration.isMethodNameGenerationEnabled());
     EditorActionsUtil.refreshActions();
   }
 

src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java

@@ -22,6 +22,7 @@ public class ConfigurationState implements PersistentStateComponent<Configuratio
   private double temperature = 0.2;
   private boolean createNewChatOnEachAction;
   private boolean ignoreGitCommitTokenLimit;
+  private boolean methodNameGenerationEnabled = true;
   private Map<String, String> tableData = EditorActionsUtil.DEFAULT_ACTIONS;
 
   public static ConfigurationState getInstance() {
@@ -86,4 +87,12 @@ public class ConfigurationState implements PersistentStateComponent<Configuratio
   public void setIgnoreGitCommitTokenLimit(boolean ignoreGitCommitTokenLimit) {
     this.ignoreGitCommitTokenLimit = ignoreGitCommitTokenLimit;
   }
+
+  public boolean isMethodNameGenerationEnabled() {
+    return methodNameGenerationEnabled;
+  }
+
+  public void setMethodNameGenerationEnabled(boolean methodNameGenerationEnabled) {
+    this.methodNameGenerationEnabled = methodNameGenerationEnabled;
+  }
 }