From 845c7b4cee27a57cbb46dedbefaf0b09902d6c83 Mon Sep 17 00:00:00 2001 From: Carl-Robert Date: Sun, 19 Nov 2023 22:56:12 +0200 Subject: [PATCH] Support method name lookup generation (#280) --- .../codegpt.java-conventions.gradle.kts | 2 +- .../prompts/method-name-generator.txt | 2 + .../codegpt/PluginStartupActivity.java | 2 + .../completions/CompletionLookupService.java | 99 +++++++++++++++++++ .../CompletionRequestProvider.java | 23 +++++ .../configuration/ConfigurationComponent.java | 25 +++-- .../ConfigurationConfigurable.java | 8 +- .../configuration/ConfigurationState.java | 9 ++ .../resources/messages/codegpt.properties | 1 + 9 files changed, 162 insertions(+), 9 deletions(-) create mode 100644 codegpt-core/src/main/resources/prompts/method-name-generator.txt create mode 100644 src/main/java/ee/carlrobert/codegpt/completions/CompletionLookupService.java diff --git a/buildSrc/src/main/kotlin/codegpt.java-conventions.gradle.kts b/buildSrc/src/main/kotlin/codegpt.java-conventions.gradle.kts index d2142221..44700eac 100644 --- a/buildSrc/src/main/kotlin/codegpt.java-conventions.gradle.kts +++ b/buildSrc/src/main/kotlin/codegpt.java-conventions.gradle.kts @@ -23,7 +23,7 @@ checkstyle { } dependencies { - implementation("ee.carlrobert:llm-client:0.0.10") + implementation("ee.carlrobert:llm-client:0.0.11") } tasks { diff --git a/codegpt-core/src/main/resources/prompts/method-name-generator.txt b/codegpt-core/src/main/resources/prompts/method-name-generator.txt new file mode 100644 index 00000000..02e3d9b2 --- /dev/null +++ b/codegpt-core/src/main/resources/prompts/method-name-generator.txt @@ -0,0 +1,2 @@ +Given an existing function or method body, generate five alternative names for the function. +The response must be a comma-separated list of names. Exclude any additional information. \ No newline at end of file diff --git a/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java b/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java index 148727b4..2ae3f55e 100644 --- a/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java +++ b/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(); diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionLookupService.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionLookupService.java new file mode 100644 index 00000000..3bb90a84 --- /dev/null +++ b/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)); + } + })); + }; + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java index 0042a101..f3aa8156 100644 --- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java +++ b/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() diff --git a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java index 349fd09a..97fb8140 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java +++ b/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 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() { diff --git a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java index cae1cfe3..9f61bdf6 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java +++ b/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(); } diff --git a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java index 070bddd2..1a5bb1b9 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java +++ b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java @@ -22,6 +22,7 @@ public class ConfigurationState implements PersistentStateComponent tableData = EditorActionsUtil.DEFAULT_ACTIONS; public static ConfigurationState getInstance() { @@ -86,4 +87,12 @@ public class ConfigurationState implements PersistentStateComponent