From 2ce05a50af8e17fd0ba754aa906e9af485b40775 Mon Sep 17 00:00:00 2001 From: Carl-Robert Linnupuu Date: Mon, 26 Aug 2024 17:39:06 +0300 Subject: [PATCH] feat: add git context to code completions --- gradle/libs.versions.toml | 2 +- .../GenerateGitCommitMessageAction.java | 183 +++++++----------- .../configuration/ConfigurationComponent.java | 10 + .../llama/form/InfillPromptTemplatePanel.java | 4 +- .../codegpt/util/CommitWorkflowChanges.java | 37 ++++ .../CodeCompletionRequestFactory.kt | 48 ++--- .../codecompletions/CodeCompletionService.kt | 4 +- .../CodeGPTInlineCompletionProvider.kt | 110 +---------- .../CodeGPTInlineCompletionSuggestion.kt | 172 ++++++++++++++++ .../codecompletions/CompletionSplitter.kt | 45 +++++ .../codecompletions/InfillPromptTemplate.kt | 20 +- .../codegpt/codecompletions/InfillRequest.kt | 50 +++++ .../codecompletions/InfillRequestDetails.kt | 94 --------- .../configuration/ConfigurationSettings.kt | 1 + .../CodeCompletionConfigurationForm.kt | 7 +- .../service/codegpt/CodeGPTAvailableModels.kt | 1 + .../form/CustomServiceCodeCompletionForm.kt | 19 +- .../ee/carlrobert/codegpt/util/GitUtil.kt | 45 +++++ .../resources/messages/codegpt.properties | 1 + .../CodeCompletionServiceTest.kt | 17 +- 20 files changed, 494 insertions(+), 376 deletions(-) create mode 100644 src/main/java/ee/carlrobert/codegpt/util/CommitWorkflowChanges.java create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionSuggestion.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CompletionSplitter.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt delete mode 100644 src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.kt create mode 100644 src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index b8bd24e9..66b15b55 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -12,7 +12,7 @@ jsoup = "1.17.2" jtokkit = "1.1.0" junit = "5.11.0" kotlin = "2.0.0" -llm-client = "0.8.12" +llm-client = "0.8.14" okio = "3.9.0" tree-sitter = "0.22.6a" diff --git a/src/main/java/ee/carlrobert/codegpt/actions/GenerateGitCommitMessageAction.java b/src/main/java/ee/carlrobert/codegpt/actions/GenerateGitCommitMessageAction.java index 5ec6b964..9052cae2 100644 --- a/src/main/java/ee/carlrobert/codegpt/actions/GenerateGitCommitMessageAction.java +++ b/src/main/java/ee/carlrobert/codegpt/actions/GenerateGitCommitMessageAction.java @@ -12,10 +12,9 @@ import com.intellij.openapi.actionSystem.AnActionEvent; import com.intellij.openapi.application.ApplicationManager; import com.intellij.openapi.command.WriteCommandAction; import com.intellij.openapi.project.Project; -import com.intellij.openapi.vcs.FilePath; +import com.intellij.openapi.util.Computable; import com.intellij.openapi.vcs.VcsDataKeys; -import com.intellij.openapi.vcs.changes.Change; -import com.intellij.openapi.vfs.VirtualFile; +import com.intellij.openapi.vcs.VcsException; import com.intellij.vcs.commit.CommitWorkflowUi; import ee.carlrobert.codegpt.CodeGPTBundle; import ee.carlrobert.codegpt.EncodingManager; @@ -23,19 +22,19 @@ import ee.carlrobert.codegpt.Icons; import ee.carlrobert.codegpt.completions.CompletionRequestService; import ee.carlrobert.codegpt.settings.configuration.CommitMessageTemplate; import ee.carlrobert.codegpt.ui.OverlayUtil; +import ee.carlrobert.codegpt.util.CommitWorkflowChanges; +import ee.carlrobert.codegpt.util.GitUtil; import ee.carlrobert.llm.client.openai.completion.ErrorDetails; import ee.carlrobert.llm.completion.CompletionEventListener; -import java.io.BufferedReader; -import java.io.File; +import git4idea.repo.GitRepositoryManager; import java.io.IOException; -import java.io.InputStreamReader; import java.nio.file.Files; import java.nio.file.Path; -import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.concurrent.ExecutionException; import okhttp3.sse.EventSource; import org.jetbrains.annotations.NotNull; @@ -54,18 +53,20 @@ public class GenerateGitCommitMessageAction extends AnAction { @Override public void update(@NotNull AnActionEvent event) { - var commitWorkflowUi = event.getData(VcsDataKeys.COMMIT_WORKFLOW_UI); - if (commitWorkflowUi == null) { - event.getPresentation().setVisible(false); - return; - } + ApplicationManager.getApplication().invokeLater(() -> { + var commitWorkflowUi = event.getData(VcsDataKeys.COMMIT_WORKFLOW_UI); + if (commitWorkflowUi == null) { + event.getPresentation().setVisible(false); + return; + } - var callAllowed = CompletionRequestService.isRequestAllowed(); - event.getPresentation().setEnabled(callAllowed - && new CommitWorkflowChanges(commitWorkflowUi).isFilesSelected()); - event.getPresentation().setText(CodeGPTBundle.get(callAllowed - ? "action.generateCommitMessage.title" - : "action.generateCommitMessage.missingCredentials")); + var callAllowed = CompletionRequestService.isRequestAllowed(); + event.getPresentation().setEnabled(callAllowed + && new CommitWorkflowChanges(commitWorkflowUi).isFilesSelected()); + event.getPresentation().setText(CodeGPTBundle.get(callAllowed + ? "action.generateCommitMessage.title" + : "action.generateCommitMessage.missingCredentials")); + }); } @Override @@ -75,7 +76,7 @@ public class GenerateGitCommitMessageAction extends AnAction { return; } - var gitDiff = getGitDiff(event, project); + var gitDiff = getDiff(event, project); var tokenCount = encodingManager.countTokens(gitDiff); if (tokenCount > MAX_TOKEN_COUNT_WARNING && OverlayUtil.showTokenSoftLimitWarningDialog(tokenCount) != OK) { @@ -97,6 +98,57 @@ public class GenerateGitCommitMessageAction extends AnAction { return ActionUpdateThread.EDT; } + private String getDiff(AnActionEvent event, Project project) { + var commitWorkflowUi = Optional.ofNullable(event.getData(VcsDataKeys.COMMIT_WORKFLOW_UI)) + .orElseThrow(() -> new IllegalStateException("Could not retrieve commit workflow ui.")); + var changes = new CommitWorkflowChanges(commitWorkflowUi); + var projectBasePath = project.getBasePath(); + + try { + return ApplicationManager.getApplication().executeOnPooledThread(() -> { + try { + var repository = GitRepositoryManager.getInstance(project) + .getRepositoryForFile(project.getWorkspaceFile()); + if (repository == null) { + return ""; + } + + var stagedGitDiff = String.join("\n", GitUtil.getStagedDiff(project, repository)); + var unstagedGitDiff = String.join("\n", GitUtil.getUnstagedDiff(project, repository)); + var newFilesContent = + getNewFilesDiff(projectBasePath, changes.getIncludedUnversionedFilePaths()); + return Map.of( + "Unstaged git diff", unstagedGitDiff, + "Staged git diff", stagedGitDiff, + "New files", newFilesContent) + .entrySet().stream() + .filter(entry -> !entry.getValue().isEmpty()) + .map(entry -> "%s:%n%s".formatted(entry.getKey(), entry.getValue())) + .collect(joining("\n\n")); + } catch (VcsException e) { + throw new RuntimeException("Unable to get staged diff", e); + } + }).get(); + } catch (InterruptedException | ExecutionException e) { + throw new RuntimeException(e); + } + } + + private static String getNewFilesDiff(String projectPath, List filePaths) { + return filePaths.stream() + .map(pathString -> { + var filePath = Path.of(pathString); + var relativePath = Path.of(projectPath).relativize(filePath); + try { + return "New file '" + relativePath + "' content:\n" + Files.readString(filePath); + } catch (IOException ignored) { + return null; + } + }) + .filter(Objects::nonNull) + .collect(joining("\n")); + } + private CompletionEventListener getEventListener( Project project, CommitWorkflowUi commitWorkflowUi) { @@ -123,97 +175,4 @@ public class GenerateGitCommitMessageAction extends AnAction { } }; } - - private String getGitDiff(AnActionEvent event, Project project) { - var commitWorkflowUi = Optional.ofNullable(event.getData(VcsDataKeys.COMMIT_WORKFLOW_UI)) - .orElseThrow(() -> new IllegalStateException("Could not retrieve commit workflow ui.")); - var changes = new CommitWorkflowChanges(commitWorkflowUi); - var projectBasePath = project.getBasePath(); - var gitDiff = getGitDiff(projectBasePath, changes.getIncludedVersionedFilePaths(), false); - var stagedGitDiff = getGitDiff(projectBasePath, changes.getIncludedVersionedFilePaths(), true); - var newFilesContent = - getNewFilesDiff(projectBasePath, changes.getIncludedUnversionedFilePaths()); - - return Map.of( - "Git diff", gitDiff, - "Staged git diff", stagedGitDiff, - "New files", newFilesContent) - .entrySet().stream() - .filter(entry -> !entry.getValue().isEmpty()) - .map(entry -> "%s:%n%s".formatted(entry.getKey(), entry.getValue())) - .collect(joining("\n\n")); - } - - private String getGitDiff(String projectPath, List filePaths, boolean cached) { - if (filePaths.isEmpty()) { - return ""; - } - - var process = createGitDiffProcess(projectPath, filePaths, cached); - return new BufferedReader(new InputStreamReader(process.getInputStream())) - .lines() - .collect(joining("\n")); - } - - private String getNewFilesDiff(String projectPath, List filePaths) { - return filePaths.stream() - .map(pathString -> { - var filePath = Path.of(pathString); - var relativePath = Path.of(projectPath).relativize(filePath); - try { - return "New file '" + relativePath + "' content:\n" + Files.readString(filePath); - } catch (IOException ignored) { - return null; - } - }) - .filter(Objects::nonNull) - .collect(joining("\n")); - } - - private Process createGitDiffProcess(String projectPath, List filePaths, boolean cached) { - var command = new ArrayList(); - command.add("git"); - command.add("diff"); - if (cached) { - command.add("--cached"); - } - command.addAll(filePaths); - - var processBuilder = new ProcessBuilder(command); - processBuilder.directory(new File(projectPath)); - try { - return processBuilder.start(); - } catch (IOException ex) { - throw new RuntimeException("Unable to start git diff process", ex); - } - } - - static class CommitWorkflowChanges { - - private final List includedVersionedFilePaths; - private final List includedUnversionedFilePaths; - - CommitWorkflowChanges(CommitWorkflowUi commitWorkflowUi) { - includedVersionedFilePaths = commitWorkflowUi.getIncludedChanges().stream() - .map(Change::getVirtualFile) - .filter(Objects::nonNull) - .map(VirtualFile::getPath) - .toList(); - includedUnversionedFilePaths = commitWorkflowUi.getIncludedUnversionedFiles().stream() - .map(FilePath::getPath) - .toList(); - } - - public List getIncludedVersionedFilePaths() { - return includedVersionedFilePaths; - } - - public List getIncludedUnversionedFilePaths() { - return includedUnversionedFilePaths; - } - - public boolean isFilesSelected() { - return !includedVersionedFilePaths.isEmpty() || !includedUnversionedFilePaths.isEmpty(); - } - } } 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 6f921848..a1853f09 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java +++ b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java @@ -48,6 +48,7 @@ public class ConfigurationComponent { private final JBCheckBox autoFormattingCheckBox; private final JBCheckBox autocompletionPostProcessingCheckBox; private final JBCheckBox autocompletionContextAwareCheckBox; + private final JBCheckBox autocompletionGitContextCheckBox; private final JTextArea commitMessagePromptTextArea; private final IntegerField maxTokensField; private final JBTextField temperatureField; @@ -122,6 +123,10 @@ public class ConfigurationComponent { CodeGPTBundle.get("configurationConfigurable.autocompletionContextAwareCheckBox.label"), configuration.getAutocompletionContextAwareEnabled() ); + autocompletionGitContextCheckBox = new JBCheckBox( + CodeGPTBundle.get("configurationConfigurable.autocompletionGitContextCheckBox.label"), + configuration.getAutocompletionGitContextEnabled() + ); mainPanel = FormBuilder.createFormBuilder() .addComponent(tablePanel) @@ -133,6 +138,7 @@ public class ConfigurationComponent { .addComponent(autoFormattingCheckBox) .addComponent(autocompletionPostProcessingCheckBox) .addComponent(autocompletionContextAwareCheckBox) + .addComponent(autocompletionGitContextCheckBox) .addVerticalGap(4) .addComponent(new TitledSeparator( CodeGPTBundle.get("configurationConfigurable.section.assistant.title"))) @@ -161,6 +167,7 @@ public class ConfigurationComponent { state.setAutoFormattingEnabled(autoFormattingCheckBox.isSelected()); state.setAutocompletionPostProcessingEnabled(autocompletionPostProcessingCheckBox.isSelected()); state.setAutocompletionContextAwareEnabled(autocompletionContextAwareCheckBox.isSelected()); + state.setAutocompletionGitContextEnabled(autocompletionGitContextCheckBox.isSelected()); return state; } @@ -179,6 +186,9 @@ public class ConfigurationComponent { configuration.getAutocompletionPostProcessingEnabled()); autocompletionContextAwareCheckBox.setSelected( configuration.getAutocompletionContextAwareEnabled()); + autocompletionGitContextCheckBox.setSelected( + configuration.getAutocompletionGitContextEnabled() + ); } private Map getTableData() { diff --git a/src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java b/src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java index 75ac30eb..19873f99 100644 --- a/src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java +++ b/src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java @@ -1,7 +1,7 @@ package ee.carlrobert.codegpt.settings.service.llama.form; import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate; -import ee.carlrobert.codegpt.codecompletions.InfillRequestDetails; +import ee.carlrobert.codegpt.codecompletions.InfillRequest; public class InfillPromptTemplatePanel extends BasePromptTemplatePanel { @@ -17,6 +17,6 @@ public class InfillPromptTemplatePanel extends BasePromptTemplatePanel includedVersionedFilePaths; + private final List includedUnversionedFilePaths; + + public CommitWorkflowChanges(CommitWorkflowUi commitWorkflowUi) { + includedVersionedFilePaths = commitWorkflowUi.getIncludedChanges().stream() + .map(Change::getVirtualFile) + .filter(Objects::nonNull) + .map(VirtualFile::getPath) + .toList(); + includedUnversionedFilePaths = commitWorkflowUi.getIncludedUnversionedFiles().stream() + .map(FilePath::getPath) + .toList(); + } + + public List getIncludedVersionedFilePaths() { + return includedVersionedFilePaths; + } + + public List getIncludedUnversionedFilePaths() { + return includedUnversionedFilePaths; + } + + public boolean isFilesSelected() { + return !includedVersionedFilePaths.isEmpty() || !includedUnversionedFilePaths.isEmpty(); + } +} diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt index 9c91e081..89c2e219 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestFactory.kt @@ -12,6 +12,7 @@ import ee.carlrobert.codegpt.settings.service.custom.CustomServiceSettings import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings import ee.carlrobert.codegpt.settings.service.llama.LlamaSettingsState import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings +import ee.carlrobert.llm.client.codegpt.request.CodeCompletionRequest import ee.carlrobert.llm.client.llama.completion.LlamaCompletionRequest import ee.carlrobert.llm.client.ollama.completion.request.OllamaCompletionRequest import ee.carlrobert.llm.client.ollama.completion.request.OllamaParameters @@ -23,30 +24,34 @@ import java.nio.charset.StandardCharsets object CodeCompletionRequestFactory { + private const val MAX_TOKENS = 128 + @JvmStatic - fun buildCodeGPTRequest(details: InfillRequestDetails): OpenAITextCompletionRequest { + fun buildCodeGPTRequest(details: InfillRequest): CodeCompletionRequest { val settings = service().state.codeCompletionSettings - return OpenAITextCompletionRequest.Builder(details.prefix) - .setSuffix(details.suffix) - .setStream(true) + return CodeCompletionRequest.Builder() .setModel(settings.model) - .setMaxTokens(getMaxTokens(details.prefix, details.suffix)) - .setTemperature(0.4) + .setPrefix(details.prefix) + .setSuffix(details.suffix) + .setFileExtension(details.fileDetails?.fileExtension) + .setFileContent(details.fileDetails?.fileContent) + .setStagedDiff(details.vcsDetails?.stagedDiff) + .setUnstagedDiff(details.vcsDetails?.unstagedDiff) .build() } @JvmStatic - fun buildOpenAIRequest(details: InfillRequestDetails): OpenAITextCompletionRequest { + fun buildOpenAIRequest(details: InfillRequest): OpenAITextCompletionRequest { return OpenAITextCompletionRequest.Builder(details.prefix) .setSuffix(details.suffix) .setStream(true) - .setMaxTokens(getMaxTokens(details.prefix, details.suffix)) + .setMaxTokens(MAX_TOKENS) .setTemperature(0.4) .build() } @JvmStatic - fun buildCustomRequest(details: InfillRequestDetails): Request { + fun buildCustomRequest(details: InfillRequest): Request { val settings = service().state.codeCompletionSettings val credential = getCredential(CredentialKey.CUSTOM_SERVICE_API_KEY) return buildCustomRequest( @@ -61,7 +66,7 @@ object CodeCompletionRequestFactory { @JvmStatic fun buildCustomRequest( - details: InfillRequestDetails, + details: InfillRequest, url: String, headers: Map, body: Map, @@ -93,21 +98,19 @@ object CodeCompletionRequestFactory { } @JvmStatic - fun buildLlamaRequest(details: InfillRequestDetails): LlamaCompletionRequest { + fun buildLlamaRequest(details: InfillRequest): LlamaCompletionRequest { val settings = LlamaSettings.getCurrentState() val promptTemplate = getLlamaInfillPromptTemplate(settings) val prompt = promptTemplate.buildPrompt(details) - println("PROMPT: ") - println(prompt) return LlamaCompletionRequest.Builder(prompt) - .setN_predict(getMaxTokens(details.prefix, details.suffix)) + .setN_predict(MAX_TOKENS) .setStream(true) .setTemperature(0.4) .setStop(promptTemplate.stopTokens) .build() } - fun buildOllamaRequest(details: InfillRequestDetails): OllamaCompletionRequest { + fun buildOllamaRequest(details: InfillRequest): OllamaCompletionRequest { val settings = service().state return OllamaCompletionRequest.Builder( settings.model, @@ -116,7 +119,7 @@ object CodeCompletionRequestFactory { .setOptions( OllamaParameters.Builder() .stop(settings.fimTemplate.stopTokens) - .numPredict(getMaxTokens(details.prefix, details.suffix)) + .numPredict(MAX_TOKENS) .temperature(0.4) .build() ) @@ -137,7 +140,7 @@ object CodeCompletionRequestFactory { private fun transformValue( value: Any, template: InfillPromptTemplate, - details: InfillRequestDetails + details: InfillRequest ): Any { if (value !is String) return value @@ -152,15 +155,4 @@ object CodeCompletionRequestFactory { } } } - - private fun getMaxTokens(prefix: String, suffix: String): Int { - if ((prefix.isNotEmpty() && isBoundaryCharacter(prefix[prefix.length - 1])) - || (suffix.isNotEmpty() && isBoundaryCharacter(suffix[0])) - ) { - return 16 - } - return 36 - } - - private fun isBoundaryCharacter(c: Char): Boolean = c in "()[]{}<>~!@#$%^&*-+=|\\;:'\",./?" } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt index 723717cd..40526045 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.kt @@ -35,12 +35,12 @@ class CodeCompletionService { } fun getCodeCompletionAsync( - requestDetails: InfillRequestDetails, + requestDetails: InfillRequest, eventListener: CompletionEventListener ): EventSource = when (val selectedService = GeneralSettings.getSelectedService()) { CODEGPT -> CompletionClientProvider.getCodeGPTClient() - .getCompletionAsync(buildCodeGPTRequest(requestDetails), eventListener) + .getCodeCompletionAsync(buildCodeGPTRequest(requestDetails), eventListener) OPENAI -> CompletionClientProvider.getOpenAIClient() .getCompletionAsync(buildOpenAIRequest(requestDetails), eventListener) diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionProvider.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionProvider.kt index 6fbe8d06..e0310413 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionProvider.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionProvider.kt @@ -1,44 +1,25 @@ package ee.carlrobert.codegpt.codecompletions import com.intellij.codeInsight.inline.completion.* -import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement -import com.intellij.notification.NotificationType -import com.intellij.openapi.application.EDT import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.thisLogger -import com.intellij.openapi.util.TextRange -import ee.carlrobert.codegpt.CodeGPTKeys -import ee.carlrobert.codegpt.codecompletions.psi.CompletionContextService import ee.carlrobert.codegpt.settings.GeneralSettings -import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings import ee.carlrobert.codegpt.settings.service.ServiceType import ee.carlrobert.codegpt.settings.service.codegpt.CodeGPTServiceSettings import ee.carlrobert.codegpt.settings.service.custom.CustomServiceSettings import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings import ee.carlrobert.codegpt.settings.service.ollama.OllamaSettings import ee.carlrobert.codegpt.settings.service.openai.OpenAISettings -import ee.carlrobert.codegpt.treesitter.CodeCompletionParserFactory -import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification -import ee.carlrobert.llm.client.openai.completion.ErrorDetails -import ee.carlrobert.llm.completion.CompletionEventListener -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.channels.awaitClose -import kotlinx.coroutines.flow.channelFlow -import kotlinx.coroutines.launch -import kotlinx.coroutines.withContext -import okhttp3.sse.EventSource -import java.util.concurrent.atomic.AtomicReference import kotlin.time.Duration import kotlin.time.DurationUnit import kotlin.time.toDuration class CodeGPTInlineCompletionProvider : DebouncedInlineCompletionProvider() { + companion object { private val logger = thisLogger() } - private val currentCall = AtomicReference(null) - override val id: InlineCompletionProviderID get() = InlineCompletionProviderID("CodeGPTInlineCompletionProvider") @@ -49,70 +30,7 @@ class CodeGPTInlineCompletionProvider : DebouncedInlineCompletionProvider() { logger.error("Could not find project") return InlineCompletionSuggestion.empty() } - - return InlineCompletionSuggestion.Default(channelFlow { - val caretOffset = withContext(Dispatchers.EDT) { editor.caretModel.offset } - val infillContext = - if (service().state.autocompletionContextAwareEnabled) - service().findContext(editor, caretOffset) - else null - val infillRequest = if (infillContext == null) { - val (prefix, suffix) = withContext(Dispatchers.EDT) { - val prefix = - request.document.getText(TextRange(0, caretOffset)) - val suffix = - request.document.getText( - TextRange( - caretOffset, - request.document.textLength - ) - ) - Pair(prefix, suffix) - } - InfillRequestDetails.withoutContext(prefix, suffix) - } else { - // TODO: truncate contextElements if too long? - InfillRequestDetails.withContext( - infillContext, - caretOffset - ) - } - - currentCall.set( - project.service().getCodeCompletionAsync( - infillRequest, - CodeCompletionEventListener { - val settings = service().state - try { - var inlineText = it.toString() - if (settings.autocompletionPostProcessingEnabled) { - inlineText = CodeCompletionParserFactory - .getParserForFileExtension(request.file.virtualFile.extension) - .parse( - // TODO: ? - infillRequest.prefix, - infillRequest.suffix, - inlineText - ) - } - - editor.putUserData(CodeGPTKeys.PREVIOUS_INLAY_TEXT, inlineText) - launch { - try { - trySend(InlineCompletionGrayTextElement(inlineText)) - } catch (e: Exception) { - logger.error("Failed to send inline completion suggestion", e) - } - } - } catch (t: Throwable) { - logger.error(t) - settings.autocompletionPostProcessingEnabled = false - } - } - ) - ) - awaitClose { cancelCurrentCall() } - }) + return CodeGPTInlineCompletionSuggestion(project, request) } override suspend fun getDebounceDelay(request: InlineCompletionRequest): Duration { @@ -134,28 +52,4 @@ class CodeGPTInlineCompletionProvider : DebouncedInlineCompletionProvider() { } return event is InlineCompletionEvent.DocumentChange && codeCompletionsEnabled } - - private fun cancelCurrentCall() { - currentCall.getAndSet(null)?.cancel() - } - - class CodeCompletionEventListener( - private val completed: (StringBuilder) -> Unit - ) : CompletionEventListener { - - override fun onComplete(messageBuilder: StringBuilder) { - completed(messageBuilder) - } - - override fun onCancelled(messageBuilder: StringBuilder) { - completed(messageBuilder) - } - - override fun onError(error: ErrorDetails, ex: Throwable) { - if (ex.message == null || (ex.message != null && ex.message != "Canceled")) { - showNotification(error.message, NotificationType.ERROR) - logger.error(error.message, ex) - } - } - } } \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionSuggestion.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionSuggestion.kt new file mode 100644 index 00000000..d75b9005 --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CodeGPTInlineCompletionSuggestion.kt @@ -0,0 +1,172 @@ +package ee.carlrobert.codegpt.codecompletions + +import com.intellij.codeInsight.inline.completion.InlineCompletionRequest +import com.intellij.codeInsight.inline.completion.InlineCompletionSuggestion +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionElement +import com.intellij.codeInsight.inline.completion.elements.InlineCompletionGrayTextElement +import com.intellij.notification.NotificationType +import com.intellij.openapi.application.EDT +import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.components.service +import com.intellij.openapi.diagnostic.thisLogger +import com.intellij.openapi.project.Project +import com.intellij.openapi.util.TextRange +import com.intellij.openapi.vcs.VcsException +import com.intellij.refactoring.suggested.startOffset +import ee.carlrobert.codegpt.EncodingManager +import ee.carlrobert.codegpt.codecompletions.psi.CompletionContextService +import ee.carlrobert.codegpt.codecompletions.psi.readText +import ee.carlrobert.codegpt.settings.configuration.ConfigurationSettings +import ee.carlrobert.codegpt.ui.OverlayUtil.showNotification +import ee.carlrobert.codegpt.util.GitUtil +import ee.carlrobert.llm.client.openai.completion.ErrorDetails +import ee.carlrobert.llm.completion.CompletionEventListener +import git4idea.repo.GitRepositoryManager +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.channels.awaitClose +import kotlinx.coroutines.flow.Flow +import kotlinx.coroutines.flow.channelFlow +import kotlinx.coroutines.launch +import kotlinx.coroutines.withContext +import okhttp3.sse.EventSource +import java.util.concurrent.atomic.AtomicReference + +class CodeGPTInlineCompletionSuggestion( + private val project: Project, + private val inlineCompletionRequest: InlineCompletionRequest +) : InlineCompletionSuggestion() { + + private val currentCall = AtomicReference(null) + + companion object { + private val logger = thisLogger() + private const val MAX_PROMPT_TOKENS = 128 + } + + override val suggestionFlow: Flow + get() = channelFlow { + val infillRequest = buildInfillRequest(inlineCompletionRequest) + launch { + val completionCall = + project.service().getCodeCompletionAsync( + infillRequest, + CodeCompletionEventListener { + try { + runInEdt { + trySend(InlineCompletionGrayTextElement(it.toString())) + } + } catch (e: Exception) { + logger.error("Failed to send inline completion suggestion", e) + } + } + ) + currentCall.set(completionCall) + } + awaitClose { currentCall.getAndSet(null)?.cancel() } + } + + private class CodeCompletionEventListener( + private val completed: (StringBuilder) -> Unit + ) : CompletionEventListener { + + override fun onComplete(messageBuilder: StringBuilder) { + completed(messageBuilder) + } + + override fun onCancelled(messageBuilder: StringBuilder) { + completed(messageBuilder) + } + + override fun onError(error: ErrorDetails, ex: Throwable) { + if (ex.message == null || (ex.message != null && ex.message != "Canceled")) { + showNotification(error.message, NotificationType.ERROR) + logger.error(error.message, ex) + } + } + } + + private suspend fun buildInfillRequest(request: InlineCompletionRequest): InfillRequest { + val caretOffset = withContext(Dispatchers.EDT) { request.editor.caretModel.offset } + val configurationState = service().state + val (prefix, suffix) = withContext(Dispatchers.EDT) { + val prefix = + request.document.getText(TextRange(0, caretOffset)) + val suffix = + request.document.getText( + TextRange( + caretOffset, + request.document.textLength + ) + ) + Pair( + prefix.truncateText(MAX_PROMPT_TOKENS, false), + suffix.truncateText(MAX_PROMPT_TOKENS) + ) + } + val fileExtension = request.file.virtualFile.extension + val fileContent = request.document.text + val infillRequestBuilder = InfillRequest.Builder(prefix, suffix) + .fileDetails(InfillRequest.FileDetails(fileContent, fileExtension)) + val project = request.editor.project ?: return infillRequestBuilder.build() + + val gitRepository = + project.service().getRepositoryForFile(project.workspaceFile) + if (configurationState.autocompletionGitContextEnabled && gitRepository != null) { + try { + val stagedDiff = GitUtil.getStagedDiff(project, gitRepository) + val unstagedDiff = GitUtil.getUnstagedDiff(project, gitRepository) + if (stagedDiff.isNotEmpty() || unstagedDiff.isNotEmpty()) { + infillRequestBuilder.vcsDetails( + InfillRequest.VcsDetails( + stagedDiff.joinToString("\n"), + unstagedDiff.joinToString("\n") + ) + ) + } + } catch (e: VcsException) { + logger.error("Failed to get git context", e) + } + } + + getInfillContext(request, caretOffset)?.let { infillRequestBuilder.context(it) } + + return infillRequestBuilder.build() + } + + private fun getInfillContext( + request: InlineCompletionRequest, + caretOffset: Int + ): InfillContext? { + val infillContext = + if (service().state.autocompletionContextAwareEnabled) + service().findContext(request.editor, caretOffset) + else null + + if (infillContext == null) { + return null + } + + val caretInEnclosingElement = + caretOffset - infillContext.enclosingElement.psiElement.startOffset + val entireText = infillContext.enclosingElement.psiElement.readText() + val prefix = entireText.take(caretInEnclosingElement) + val suffix = + if (entireText.length < caretInEnclosingElement) "" else entireText.takeLast( + entireText.length - caretInEnclosingElement + ) + return truncateContext(prefix + suffix, infillContext) + } + + private fun String.truncateText(maxTokens: Int, fromStart: Boolean = true): String { + return service().truncateText(this, maxTokens, fromStart) + } + + private fun truncateContext(prompt: String, infillContext: InfillContext): InfillContext { + var promptTokens = EncodingManager.getInstance().countTokens(prompt) + val truncatedContextElements = infillContext.contextElements.takeWhile { + promptTokens += it.tokens + promptTokens <= MAX_PROMPT_TOKENS + }.toSet() + return InfillContext(infillContext.enclosingElement, truncatedContextElements) + } +} \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CompletionSplitter.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CompletionSplitter.kt new file mode 100644 index 00000000..d786150f --- /dev/null +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/CompletionSplitter.kt @@ -0,0 +1,45 @@ +package ee.carlrobert.codegpt.codecompletions + +class CompletionSplitter { + companion object { + fun split(completionText: String): String { + val boundaryPredicates = listOf<(Char) -> Boolean>( + { Character.isWhitespace(it) }, + { it.isBoundaryCharacter() }, + { !it.isBoundaryCharacter() && !Character.isWhitespace(it) } + ) + for (predicate in boundaryPredicates) { + val blockIndex = + findContinuousBlock(completionText, predicate) + if (blockIndex != -1) { + return completionText.substring(0, blockIndex) + } + } + return completionText + } + + private fun findContinuousBlock( + fullCompletion: String, + isBoundaryCharacter: BoundaryFinder + ): Int { + if (!isBoundaryCharacter.isBoundaryCharacter(fullCompletion[0])) { + return -1 + } + + var endIndex = 0 + while (endIndex < fullCompletion.length + && isBoundaryCharacter.isBoundaryCharacter(fullCompletion[endIndex]) + ) { + ++endIndex + } + + return endIndex + } + + internal fun interface BoundaryFinder { + fun isBoundaryCharacter(c: Char): Boolean + } + } +} + +fun Char?.isBoundaryCharacter(): Boolean = this != null && this in "()[]{}<>~!@#$%^&*-+=|\\;:'\",./?" \ No newline at end of file diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt index f7d17705..0f81de33 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt @@ -3,14 +3,14 @@ package ee.carlrobert.codegpt.codecompletions enum class InfillPromptTemplate(val label: String, val stopTokens: List?) { OPENAI("OpenAI", null) { - override fun buildPrompt(infillDetails: InfillRequestDetails): String { + override fun buildPrompt(infillDetails: InfillRequest): String { val infillPrompt = "<|fim_prefix|> ${infillDetails.prefix} <|fim_suffix|>${infillDetails.suffix} <|fim_middle|>" return createDefaultMultiFilePrompt(infillDetails, infillPrompt) } }, CODE_LLAMA("Code Llama", listOf("")) { - override fun buildPrompt(infillDetails: InfillRequestDetails): String { + override fun buildPrompt(infillDetails: InfillRequest): String { val infillPrompt = "
 ${infillDetails.prefix} ${infillDetails.suffix} "
             return createDefaultMultiFilePrompt(infillDetails, infillPrompt)
         }
@@ -19,7 +19,7 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List?
         "CodeGemma Instruct",
         listOf("<|file_separator|>", "<|fim_prefix|>", "<|fim_suffix|>", "<|fim_middle|>", "")
     ) {
-        override fun buildPrompt(infillDetails: InfillRequestDetails): String {
+        override fun buildPrompt(infillDetails: InfillRequest): String {
             // see https://huggingface.co/google/codegemma-7b#for-code-completion
             val infillPrompt =
                 "<|fim_prefix|>${infillDetails.prefix}<|fim_suffix|>${infillDetails.suffix}<|fim_middle|>"
@@ -36,7 +36,7 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List?
         }
     },
     CODE_QWEN("CodeQwen1.5", listOf("<|endoftext|>")) {
-        override fun buildPrompt(infillDetails: InfillRequestDetails): String {
+        override fun buildPrompt(infillDetails: InfillRequest): String {
             // see https://github.com/QwenLM/CodeQwen1.5?tab=readme-ov-file#2-file-level-code-completion-fill-in-the-middle
             val infillPrompt =
                 "${infillDetails.prefix}${infillDetails.suffix}"
@@ -54,14 +54,14 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List?
         }
     },
     STABILITY("Stability AI", listOf("<|endoftext|>")) {
-        override fun buildPrompt(infillDetails: InfillRequestDetails): String {
+        override fun buildPrompt(infillDetails: InfillRequest): String {
             val infillPrompt =
                 "${infillDetails.prefix}${infillDetails.suffix}"
             return createDefaultMultiFilePrompt(infillDetails, infillPrompt)
         }
     },
     DEEPSEEK_CODER("DeepSeek Coder", listOf("<|EOT|>")) {
-        override fun buildPrompt(infillDetails: InfillRequestDetails): String {
+        override fun buildPrompt(infillDetails: InfillRequest): String {
             // see https://github.com/deepseek-ai/DeepSeek-Coder?tab=readme-ov-file#2-code-insertion
             val infillPrompt =
                 "<｜fim▁begin｜>${infillDetails.prefix}<｜fim▁hole｜>${infillDetails.suffix}<｜fim▁end｜>"
@@ -76,7 +76,7 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List?
         }
     },
     STAR_CODER("StarCoder2", listOf("<|endoftext|>")) {
-        override fun buildPrompt(infillDetails: InfillRequestDetails): String {
+        override fun buildPrompt(infillDetails: InfillRequest): String {
             // see https://huggingface.co/spaces/bigcode/bigcode-playground/blob/main/app.py
             val infillPrompt =
                 "${infillDetails.prefix}  ${infillDetails.suffix}"
@@ -94,14 +94,14 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List?
         }
     },
     CODESTRAL("Codestral", listOf("")) {
-        override fun buildPrompt(infillDetails: InfillRequestDetails): String {
+        override fun buildPrompt(infillDetails: InfillRequest): String {
             // see https://github.com/mistralai/mistral-common/blob/master/src/mistral_common/tokens/tokenizers/base.py
             val infillPrompt = "[SUFFIX]${infillDetails.suffix}[PREFIX]${infillDetails.prefix}[MIDDLE]"
             return createDefaultMultiFilePrompt(infillDetails, infillPrompt)
         }
     };
 
-    abstract fun buildPrompt(infillDetails: InfillRequestDetails): String
+    abstract fun buildPrompt(infillDetails: InfillRequest): String
 
     override fun toString(): String {
         return label
@@ -109,7 +109,7 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List?
 
     companion object {
         private fun createDefaultMultiFilePrompt(
-            infillDetails: InfillRequestDetails,
+            infillDetails: InfillRequest,
             infillPrompt: String
         ): String {
             val context = infillDetails.context
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt
new file mode 100644
index 00000000..dd6d028a
--- /dev/null
+++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequest.kt
@@ -0,0 +1,50 @@
+package ee.carlrobert.codegpt.codecompletions
+
+import com.intellij.psi.PsiElement
+import ee.carlrobert.codegpt.codecompletions.psi.filePath
+import ee.carlrobert.codegpt.codecompletions.psi.readText
+
+class InfillRequest private constructor(
+    val prefix: String,
+    val suffix: String,
+    val fileDetails: FileDetails?,
+    val vcsDetails: VcsDetails?,
+    val context: InfillContext?
+) {
+
+    data class VcsDetails(val stagedDiff: String? = null, val unstagedDiff: String? = null)
+    data class FileDetails(val fileContent: String, val fileExtension: String? = null)
+
+    class Builder(private val prefix: String, private val suffix: String) {
+        private var fileDetails: FileDetails? = null
+        private var vcsDetails: VcsDetails? = null
+        private var context: InfillContext? = null
+
+        fun fileDetails(fileDetails: FileDetails) = apply { this.fileDetails = fileDetails }
+        fun vcsDetails(vcsDetails: VcsDetails) = apply { this.vcsDetails = vcsDetails }
+        fun context(context: InfillContext) = apply { this.context = context }
+
+        fun build() =
+            InfillRequest(prefix, suffix, fileDetails, vcsDetails, context)
+    }
+
+    companion object {
+        fun builder(prefix: String, suffix: String) = Builder(prefix, suffix)
+    }
+}
+
+class InfillContext(
+    val enclosingElement: ContextElement,
+    // TODO: Add some kind of ranking, which contextElements are more important than others
+    val contextElements: Set
+) {
+
+    fun getRepoName(): String = enclosingElement.psiElement.project.name
+}
+
+class ContextElement(val psiElement: PsiElement) {
+    var tokens: Int = -1
+
+    fun filePath() = this.psiElement.filePath()
+    fun text() = this.psiElement.readText()
+}
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.kt
deleted file mode 100644
index c07b00b1..00000000
--- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.kt
+++ /dev/null
@@ -1,94 +0,0 @@
-package ee.carlrobert.codegpt.codecompletions
-
-import com.intellij.codeInsight.navigation.ImplementationSearcher
-import com.intellij.psi.PsiElement
-import com.intellij.refactoring.suggested.startOffset
-import ee.carlrobert.codegpt.EncodingManager
-import ee.carlrobert.codegpt.codecompletions.psi.filePath
-import ee.carlrobert.codegpt.codecompletions.psi.readText
-
-
-class InfillRequestDetails(val prefix: String, val suffix: String, val context: InfillContext?) :
-    ImplementationSearcher() {
-    companion object {
-        private const val MAX_OFFSET = 10_000
-        private const val MAX_PROMPT_TOKENS = 128
-        private const val MAX_INFILL_PROMPT_TOKENS = 1_000
-
-        fun withoutContext(
-            prefix: String,
-            suffix: String
-        ): InfillRequestDetails {
-            val truncatedPrefix = prefix.takeLast(MAX_OFFSET)
-            val truncatedSuffix = suffix.take(MAX_OFFSET)
-            return InfillRequestDetails(
-                truncateText(truncatedPrefix, false),
-                truncateText(truncatedSuffix, true),
-                null
-            )
-        }
-
-        fun withContext(
-            infillContext: InfillContext,
-            caretOffsetInFile: Int,
-        ): InfillRequestDetails {
-            val caretInEnclosingElement =
-                caretOffsetInFile - infillContext.enclosingElement.psiElement.startOffset
-            val entireText = infillContext.enclosingElement.psiElement.readText()
-            val prefix = truncateText(entireText.take(caretInEnclosingElement), false)
-            val suffix = truncateText(
-                if (entireText.length < caretInEnclosingElement) "" else entireText.takeLast(
-                    entireText.length - caretInEnclosingElement
-                ), true
-            )
-            return InfillRequestDetails(
-                prefix,
-                suffix,
-                truncateContext(prefix + suffix, infillContext)
-            )
-        }
-
-        private fun truncateContext(prompt: String, infillContext: InfillContext): InfillContext {
-            var promptTokens = EncodingManager.getInstance().countTokens(prompt)
-            val truncatedContextElements = infillContext.contextElements.takeWhile {
-                promptTokens += it.tokens
-                promptTokens <= MAX_INFILL_PROMPT_TOKENS
-            }.toSet()
-            return InfillContext(infillContext.enclosingElement, truncatedContextElements)
-        }
-
-        private fun truncateText(
-            text: String,
-            fromStart: Boolean
-        ): String {
-            return EncodingManager.getInstance().truncateText(
-                text,
-                MAX_PROMPT_TOKENS,
-                fromStart
-            )
-        }
-    }
-}
-
-class InfillContext(
-    val enclosingElement: ContextElement,
-    // TODO: Add some kind of ranking, which contextElements are more important than others
-    val contextElements: Set
-) {
-
-    fun getRepoName(): String = enclosingElement.psiElement.project.name
-}
-
-
-class ContextElement {
-    val psiElement: PsiElement
-    var tokens: Int
-
-    constructor(psiElement: PsiElement) {
-        this.psiElement = psiElement
-        this.tokens = -1
-    }
-
-    fun filePath() = this.psiElement.filePath()
-    fun text() = this.psiElement.readText()
-}
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt
index c1892f8a..a330cdb2 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/configuration/ConfigurationSettings.kt
@@ -34,6 +34,7 @@ class ConfigurationSettingsState : BaseState() {
     var autoFormattingEnabled by property(true)
     var autocompletionPostProcessingEnabled by property(false)
     var autocompletionContextAwareEnabled by property(false)
+    var autocompletionGitContextEnabled by property(true)
     var tableData by map()
 
     init {
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt
index 3570a27a..ac381e65 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/CodeCompletionConfigurationForm.kt
@@ -9,7 +9,7 @@ import com.intellij.ui.components.JBLabel
 import com.intellij.util.ui.FormBuilder
 import ee.carlrobert.codegpt.CodeGPTBundle
 import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate
-import ee.carlrobert.codegpt.codecompletions.InfillRequestDetails
+import ee.carlrobert.codegpt.codecompletions.InfillRequest
 import org.apache.commons.text.StringEscapeUtils
 import java.awt.FlowLayout
 import javax.swing.Box
@@ -64,8 +64,9 @@ class CodeCompletionConfigurationForm(
     private fun updatePromptTemplateHelpTooltip(template: InfillPromptTemplate) {
         promptTemplateHelpText.setToolTipText(null)
 
-        val description = StringEscapeUtils.escapeHtml4(template.buildPrompt(
-            InfillRequestDetails("PREFIX", "SUFFIX", null)))
+        val description = StringEscapeUtils.escapeHtml4(
+            template.buildPrompt(InfillRequest.Builder("PREFIX", "SUFFIX").build())
+        )
         HelpTooltip()
             .setTitle(template.toString())
             .setDescription("
$description
")
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt
index 8c03573a..27c07a22 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/codegpt/CodeGPTAvailableModels.kt
@@ -59,6 +59,7 @@ object CodeGPTAvailableModels {
 
     @JvmStatic
     val CODE_MODELS: List = listOf(
+        CodeGPTModel("Codestral", "codestral", Icons.OpenAI, INDIVIDUAL),
         CodeGPTModel("GPT-3.5 Turbo Instruct", "gpt-3.5-turbo-instruct", Icons.OpenAI, INDIVIDUAL),
         CodeGPTModel("StarCoder (16B)", "starcoder-16b", Icons.CodeGPTModel, FREE),
         CodeGPTModel("StarCoder (7B) - FREE", "starcoder-7b", Icons.CodeGPTModel, FREE),
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt
index 7d4a439f..5a411570 100644
--- a/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt
+++ b/src/main/kotlin/ee/carlrobert/codegpt/settings/service/custom/form/CustomServiceCodeCompletionForm.kt
@@ -13,7 +13,7 @@ import com.intellij.util.ui.FormBuilder
 import ee.carlrobert.codegpt.CodeGPTBundle
 import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestFactory
 import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate
-import ee.carlrobert.codegpt.codecompletions.InfillRequestDetails
+import ee.carlrobert.codegpt.codecompletions.InfillRequest
 import ee.carlrobert.codegpt.completions.CompletionRequestService
 import ee.carlrobert.codegpt.settings.configuration.Placeholder
 import ee.carlrobert.codegpt.settings.service.custom.CustomServiceCodeCompletionSettingsState
@@ -29,7 +29,6 @@ import java.awt.FlowLayout
 import javax.swing.Box
 import javax.swing.JButton
 import javax.swing.JPanel
-import javax.swing.SwingUtilities
 
 class CustomServiceCodeCompletionForm(
     state: CustomServiceCodeCompletionSettingsState,
@@ -110,7 +109,13 @@ class CustomServiceCodeCompletionForm(
                 }
             )
             .addComponent(tabbedPane)
-            .addComponent(ComponentPanelBuilder.createCommentComponent(getHtmlDescription(), true, 100))
+            .addComponent(
+                ComponentPanelBuilder.createCommentComponent(
+                    getHtmlDescription(),
+                    true,
+                    100
+                )
+            )
             .addComponentFillVertically(JPanel(), 0)
             .panel
 
@@ -145,7 +150,7 @@ class CustomServiceCodeCompletionForm(
     private fun testConnection() {
         CompletionRequestService.getInstance().getCustomOpenAICompletionAsync(
             CodeCompletionRequestFactory.buildCustomRequest(
-                InfillRequestDetails("Hello", "!", null),
+                InfillRequest.Builder("Hello", "!").build(),
                 urlField.text,
                 tabbedPane.headers,
                 tabbedPane.body,
@@ -188,11 +193,7 @@ class CustomServiceCodeCompletionForm(
 
         val description = StringEscapeUtils.escapeHtml4(
             template.buildPrompt(
-                InfillRequestDetails(
-                    "PREFIX",
-                    "SUFFIX",
-                    null
-                )
+                InfillRequest.Builder("PREFIX", "SUFFIX").build(),
             )
         )
         HelpTooltip()
diff --git a/src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt
new file mode 100644
index 00000000..8212d55c
--- /dev/null
+++ b/src/main/kotlin/ee/carlrobert/codegpt/util/GitUtil.kt
@@ -0,0 +1,45 @@
+package ee.carlrobert.codegpt.util
+
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.vcs.VcsException
+import git4idea.commands.Git
+import git4idea.commands.GitCommand
+import git4idea.commands.GitLineHandler
+import git4idea.repo.GitRepository
+
+object GitUtil {
+
+    @Throws(VcsException::class)
+    @JvmStatic
+    fun getStagedDiff(project: Project, gitRepository: GitRepository): List {
+        return getGitDiff(project, gitRepository, true)
+    }
+
+    @Throws(VcsException::class)
+    @JvmStatic
+    fun getUnstagedDiff(project: Project, gitRepository: GitRepository): List {
+        return getGitDiff(project, gitRepository, false)
+    }
+
+    private fun getGitDiff(
+        project: Project,
+        gitRepository: GitRepository,
+        staged: Boolean
+    ): List {
+        val handler = GitLineHandler(project, gitRepository.root, GitCommand.DIFF)
+        if (staged) {
+            handler.addParameters("--cached")
+        }
+        handler.addParameters(
+            "--unified=2",
+            "--diff-filter=AM",
+            "--no-prefix",
+            "--no-color",
+        )
+
+        val commandResult = Git.getInstance().runCommand(handler)
+        return commandResult.output.filter {
+            listOf("diff --git", "index ", "---", "- ", "+++").none { prefix -> it.startsWith(prefix) }
+        }
+    }
+}
diff --git a/src/main/resources/messages/codegpt.properties b/src/main/resources/messages/codegpt.properties
index e1b7e81b..5e2b7778 100644
--- a/src/main/resources/messages/codegpt.properties
+++ b/src/main/resources/messages/codegpt.properties
@@ -116,6 +116,7 @@ configurationConfigurable.enableMethodNameGeneration.label=Enable method name lo
 configurationConfigurable.autoFormatting.label=Enable automatic code formatting
 configurationConfigurable.autocompletionPostProcessing.label=Enable code completion post processing
 configurationConfigurable.autocompletionContextAwareCheckBox.label=Enable project context aware code completion
+configurationConfigurable.autocompletionGitContextCheckBox.label=Enable Git context for code completions
 configurationConfigurable.section.assistant.title=Assistant Configuration
 configurationConfigurable.section.assistant.systemPromptField.label=System prompt:
 configurationConfigurable.section.assistant.systemPromptField.comment=The system message helps to set the behaviour of the assistant
diff --git a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt
index d4476330..fbf2d0e5 100644
--- a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt
+++ b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt
@@ -39,14 +39,17 @@ class CodeCompletionServiceTest : IntegrationTest() {
             assertThat(request.method).isEqualTo("POST")
             assertThat(request.body)
                 .extracting("prompt")
-                .isEqualTo(InfillPromptTemplate.CODE_LLAMA.buildPrompt(
-                    InfillRequestDetails(
-                        prefix,
-                        suffix,
-                        null
+                .isEqualTo(
+                    InfillPromptTemplate.CODE_LLAMA.buildPrompt(
+                        InfillRequest.Builder(prefix, suffix).build()
                     )
-                ))
-            listOf(jsonMapResponse(e("content", expectedCompletion), e("stop", true)))
+                )
+            listOf(
+                jsonMapResponse(
+                    e("content", expectedCompletion),
+                    e("stop", true)
+                )
+            )
         })
 
         myFixture.type('c')