From 7387cf4536b140e3be9c7cc66bdb43fd0705679b Mon Sep 17 00:00:00 2001 From: Phil <39240633+PhilKes@users.noreply.github.com> Date: Wed, 31 Jan 2024 00:05:31 +0100 Subject: [PATCH] Inline Autocompletion Pt.2 (#333) * Add first draft of inline code completion with mock text * Adds InsertInlineTextAction for inserting autocomplete suggestion with tab - Changed to disable suggestions when text is selected - Adds and removes the insert action based on when it shows the inlay hint * Request inline code completion * Move inline completion prompt into txt file * Add inline completion settings to ConfigurationState * Fix code style * Use EditorTrackerListener instead of EditorFactoryListener to enable inline completion * Code completion requests synchronously without SSE * Use LlamaClient.getInfill() for inline code completion * support inlay block element rendering, clean up code * Use only enclosed Method or Class contents for code completion if possible * Refactor extracting PsiElement contents in code completion * bump llm-client * fix completion call from triggering on EDT, force method params to be nonnull by default * refactor request building, decrease delay value * Trigger code completion if cursor is not inside a word * Improve inlay rendering * Support cancellable infill requests * add statusbar widget, disable completions by default * Show error notification if code completion failed * Truely disable/enable EditorInlayHandler when completion is turned off/on * Add CodeCompletionEnabledListener Topic to control enabling/disabling code-completion * Add progress indicator for code-completion with option to cancel * Add CodeCompletionServiceTest + refactor inlay ElementRenderers * several improvements - replace timer implementation with call debouncing - use OpenAI /v1/completions API for completions - code refactoring * trigger progress indicator only for llama completions * fix tests --------- Co-authored-by: James Higgins Co-authored-by: Carl-Robert Linnupuu --- README.md | 5 + .../codegpt.java-conventions.gradle.kts | 2 +- .../prompts/inline-completion-prompt.txt | 1 + .../ee/carlrobert/codegpt/CodeGPTKeys.java | 9 +- .../carlrobert/codegpt/EncodingManager.java | 2 +- .../codegpt/PluginStartupActivity.java | 1 - .../CodeCompletionEnabledListener.java | 26 ++ .../actions/DisableCompletionsAction.java | 30 ++ .../actions/EnableCompletionsAction.java | 28 ++ .../codegpt/actions/OpenSettingsAction.java | 23 ++ .../codecompletions/CallDebouncer.java | 76 +++++ .../codegpt/codecompletions/CallRunnable.java | 11 + .../CodeCompletionEventListener.java | 75 +++++ .../CodeCompletionListenerBinder.java | 120 ++++++++ .../CodeCompletionRequestProvider.java | 38 +++ .../CodeCompletionService.java | 265 ++++++++++++++++++ .../CodeGPTEditorListener.java | 16 ++ .../codecompletions/CodeGPTEditorManager.java | 51 ++++ .../codecompletions/InfillPromptTemplate.java | 57 ++++ .../codecompletions/InfillRequestDetails.java | 20 ++ .../InlayBlockElementRenderer.java | 57 ++++ .../codecompletions/InlayElementRenderer.java | 38 +++ .../InlayInlineElementRenderer.java | 37 +++ .../completions/CompletionRequestHandler.java | 5 + .../CompletionRequestProvider.java | 10 +- .../completions/CompletionRequestService.java | 28 +- .../codegpt/completions/ConversationType.java | 1 + .../TotalUsageExceededException.java | 1 + .../configuration/ConfigurationComponent.java | 103 ++++++- .../ConfigurationConfigurable.java | 7 + .../configuration/ConfigurationState.java | 28 ++ .../settings/state/LlamaSettingsState.java | 10 + .../statusbar/CodeGPTStatusBarWidget.java | 52 ++++ .../CodeGPTStatusBarWidgetFactory.java | 33 +++ .../chat/standard/EditorActionEvent.java | 1 + .../carlrobert/codegpt/util/EditorUtil.java | 22 +- .../carlrobert/codegpt/util/MarkdownUtil.java | 10 +- src/main/resources/META-INF/plugin.xml | 27 ++ .../resources/messages/codegpt.properties | 13 +- .../CodeCompletionServiceTest.java | 138 +++++++++ .../codecompletions/code-completion-file.txt | 5 + 41 files changed, 1461 insertions(+), 21 deletions(-) create mode 100644 codegpt-core/src/main/resources/prompts/inline-completion-prompt.txt create mode 100644 src/main/java/ee/carlrobert/codegpt/actions/CodeCompletionEnabledListener.java create mode 100644 src/main/java/ee/carlrobert/codegpt/actions/DisableCompletionsAction.java create mode 100644 src/main/java/ee/carlrobert/codegpt/actions/EnableCompletionsAction.java create mode 100644 src/main/java/ee/carlrobert/codegpt/actions/OpenSettingsAction.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/CallDebouncer.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/CallRunnable.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionListenerBinder.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestProvider.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/CodeGPTEditorListener.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/CodeGPTEditorManager.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/InlayBlockElementRenderer.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/InlayElementRenderer.java create mode 100644 src/main/java/ee/carlrobert/codegpt/codecompletions/InlayInlineElementRenderer.java create mode 100644 src/main/java/ee/carlrobert/codegpt/statusbar/CodeGPTStatusBarWidget.java create mode 100644 src/main/java/ee/carlrobert/codegpt/statusbar/CodeGPTStatusBarWidgetFactory.java create mode 100644 src/test/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.java create mode 100644 src/test/resources/codecompletions/code-completion-file.txt diff --git a/README.md b/README.md index 92007b72..46e65055 100644 --- a/README.md +++ b/README.md @@ -164,6 +164,11 @@ git submodule update ./gradlew runIde -Penv=win-arm64 ``` +**Tailing logs** +```shell +tail -f build/idea-sandbox/system/log/idea.log +``` + ## Issues See the [open issues][open-issues] for a full list of proposed features (and known issues). diff --git a/buildSrc/src/main/kotlin/codegpt.java-conventions.gradle.kts b/buildSrc/src/main/kotlin/codegpt.java-conventions.gradle.kts index e5ac8d8f..0c6560df 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.2.0") + implementation("ee.carlrobert:llm-client:0.3.1") } tasks { diff --git a/codegpt-core/src/main/resources/prompts/inline-completion-prompt.txt b/codegpt-core/src/main/resources/prompts/inline-completion-prompt.txt new file mode 100644 index 00000000..391ad2c2 --- /dev/null +++ b/codegpt-core/src/main/resources/prompts/inline-completion-prompt.txt @@ -0,0 +1 @@ +{pre}{codeBefore}{suf}{codeAfter}{mid} \ No newline at end of file diff --git a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java index e5a99fe3..89c70747 100644 --- a/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java +++ b/src/main/java/ee/carlrobert/codegpt/CodeGPTKeys.java @@ -1,10 +1,17 @@ package ee.carlrobert.codegpt; +import com.intellij.openapi.editor.EditorCustomElementRenderer; +import com.intellij.openapi.editor.Inlay; import com.intellij.openapi.util.Key; import ee.carlrobert.embedding.ReferencedFile; import java.util.List; public class CodeGPTKeys { - public static final Key> SELECTED_FILES = Key.create("selectedFiles"); + public static final Key> SINGLE_LINE_INLAY = + Key.create("codegpt.editor.inlay.single-line"); + public static final Key> MULTI_LINE_INLAY = + Key.create("codegpt.editor.inlay.multi-line"); + public static final Key> SELECTED_FILES = + Key.create("codegpt.selectedFiles"); } diff --git a/src/main/java/ee/carlrobert/codegpt/EncodingManager.java b/src/main/java/ee/carlrobert/codegpt/EncodingManager.java index b25c60c3..d245f66b 100644 --- a/src/main/java/ee/carlrobert/codegpt/EncodingManager.java +++ b/src/main/java/ee/carlrobert/codegpt/EncodingManager.java @@ -48,7 +48,7 @@ public final class EncodingManager { try { return encoding.countTokens(text); } catch (Exception ex) { - LOG.error(ex); + LOG.warn(ex); return 0; } } diff --git a/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java b/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java index 824a84b6..5635825a 100644 --- a/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java +++ b/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java @@ -22,7 +22,6 @@ public class PluginStartupActivity implements StartupActivity { @Override public void runActivity(@NotNull Project project) { EditorActionsUtil.refreshActions(); - var authenticationResponse = YouUserManager.getInstance().getAuthenticationResponse(); if (authenticationResponse == null) { handleYouServiceAuthentication(); diff --git a/src/main/java/ee/carlrobert/codegpt/actions/CodeCompletionEnabledListener.java b/src/main/java/ee/carlrobert/codegpt/actions/CodeCompletionEnabledListener.java new file mode 100644 index 00000000..1e8a855c --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/actions/CodeCompletionEnabledListener.java @@ -0,0 +1,26 @@ +package ee.carlrobert.codegpt.actions; + +import com.intellij.util.messages.Topic; +import com.intellij.util.messages.Topic.BroadcastDirection; +import ee.carlrobert.codegpt.settings.configuration.ConfigurationState; +import java.util.EventListener; + +/** + * {@link EventListener} for changes of {@link ConfigurationState#isCodeCompletionsEnabled()}. + * + * @see EnableCompletionsAction + * @see DisableCompletionsAction + */ +public interface CodeCompletionEnabledListener extends EventListener { + + /** + * Topic for subscribing to {@link ConfigurationState#isCodeCompletionsEnabled()} changes.
+ * Broadcasts from Application-Level to all projects. + */ + @Topic.AppLevel + Topic TOPIC = new Topic<>(CodeCompletionEnabledListener.class, + BroadcastDirection.TO_DIRECT_CHILDREN); + + void onCodeCompletionsEnabledChange(boolean codeCompletionsEnabled); +} + diff --git a/src/main/java/ee/carlrobert/codegpt/actions/DisableCompletionsAction.java b/src/main/java/ee/carlrobert/codegpt/actions/DisableCompletionsAction.java new file mode 100644 index 00000000..13e444d6 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/actions/DisableCompletionsAction.java @@ -0,0 +1,30 @@ +package ee.carlrobert.codegpt.actions; + +import com.intellij.openapi.actionSystem.AnAction; +import com.intellij.openapi.actionSystem.AnActionEvent; +import com.intellij.openapi.application.ApplicationManager; +import ee.carlrobert.codegpt.codecompletions.CodeGPTEditorManager; +import ee.carlrobert.codegpt.settings.configuration.ConfigurationState; +import org.jetbrains.annotations.NotNull; + +/** + * Disables code-completion.
Publishes message to {@link CodeCompletionEnabledListener#TOPIC} + */ +public class DisableCompletionsAction extends AnAction { + + @Override + public void actionPerformed(@NotNull AnActionEvent e) { + ConfigurationState.getInstance().setCodeCompletionsEnabled(false); + CodeGPTEditorManager.getInstance().disposeAllInlays(e.getProject()); + ApplicationManager.getApplication() + .getMessageBus().syncPublisher(CodeCompletionEnabledListener.TOPIC) + .onCodeCompletionsEnabledChange(false); + } + + @Override + public void update(@NotNull AnActionEvent e) { + var codeCompletionEnabled = ConfigurationState.getInstance().isCodeCompletionsEnabled(); + e.getPresentation().setEnabled(codeCompletionEnabled); + e.getPresentation().setVisible(codeCompletionEnabled); + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/actions/EnableCompletionsAction.java b/src/main/java/ee/carlrobert/codegpt/actions/EnableCompletionsAction.java new file mode 100644 index 00000000..39fe132b --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/actions/EnableCompletionsAction.java @@ -0,0 +1,28 @@ +package ee.carlrobert.codegpt.actions; + +import com.intellij.openapi.actionSystem.AnAction; +import com.intellij.openapi.actionSystem.AnActionEvent; +import com.intellij.openapi.application.ApplicationManager; +import ee.carlrobert.codegpt.settings.configuration.ConfigurationState; +import org.jetbrains.annotations.NotNull; + +/** + * Enables code-completion.
Publishes message to {@link CodeCompletionEnabledListener#TOPIC} + */ +public class EnableCompletionsAction extends AnAction { + + @Override + public void actionPerformed(@NotNull AnActionEvent e) { + ConfigurationState.getInstance().setCodeCompletionsEnabled(true); + ApplicationManager.getApplication() + .getMessageBus().syncPublisher(CodeCompletionEnabledListener.TOPIC) + .onCodeCompletionsEnabledChange(true); + } + + @Override + public void update(@NotNull AnActionEvent e) { + var codeCompletionEnabled = ConfigurationState.getInstance().isCodeCompletionsEnabled(); + e.getPresentation().setEnabled(!codeCompletionEnabled); + e.getPresentation().setVisible(!codeCompletionEnabled); + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/actions/OpenSettingsAction.java b/src/main/java/ee/carlrobert/codegpt/actions/OpenSettingsAction.java new file mode 100644 index 00000000..8cc25e64 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/actions/OpenSettingsAction.java @@ -0,0 +1,23 @@ +package ee.carlrobert.codegpt.actions; + +import com.intellij.icons.AllIcons.General; +import com.intellij.openapi.actionSystem.AnAction; +import com.intellij.openapi.actionSystem.AnActionEvent; +import com.intellij.openapi.options.ShowSettingsUtil; +import ee.carlrobert.codegpt.CodeGPTBundle; +import ee.carlrobert.codegpt.settings.SettingsConfigurable; +import org.jetbrains.annotations.NotNull; + +public class OpenSettingsAction extends AnAction { + + public OpenSettingsAction() { + super(CodeGPTBundle.get("action.opensettings.title"), + CodeGPTBundle.get("action.opensettings.description"), + General.Settings); + } + + @Override + public void actionPerformed(@NotNull AnActionEvent e) { + ShowSettingsUtil.getInstance().showSettingsDialog(e.getProject(), SettingsConfigurable.class); + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CallDebouncer.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CallDebouncer.java new file mode 100644 index 00000000..d9085240 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CallDebouncer.java @@ -0,0 +1,76 @@ +package ee.carlrobert.codegpt.codecompletions; + +import static ee.carlrobert.codegpt.settings.service.ServiceType.LLAMA_CPP; + +import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator; +import com.intellij.openapi.project.Project; +import ee.carlrobert.codegpt.CodeGPTBundle; +import ee.carlrobert.codegpt.settings.state.SettingsState; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import okhttp3.sse.EventSource; + +public class CallDebouncer { + + private final Project project; + private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor(); + private final ConcurrentHashMap> delayedMap = new ConcurrentHashMap<>(); + private final AtomicReference currentCall = new AtomicReference<>(); + + public CallDebouncer(Project project) { + this.project = project; + } + + /** + * Implements a debounce mechanism for {@code callable} with a specified {@code delay}. This means + * the callable is set to execute after the given {@code delay} period. However, if this method is + * invoked again with the same key before the {@code delay} elapses, the scheduled execution will + * be aborted, and therefore the previous request will be cancelled. + */ + public void debounce(Object key, CallRunnable runnable, long delay, TimeUnit unit) { + Future prev = delayedMap.put(key, scheduler.schedule(() -> { + try { + cancelPreviousCall(); + var progressIndicator = LLAMA_CPP.equals(SettingsState.getInstance().getSelectedService()) + ? createProgressIndicator() + : null; + currentCall.set(runnable.call(progressIndicator)); + } finally { + delayedMap.remove(key); + } + }, delay, unit)); + + if (prev != null) { + prev.cancel(true); + } + } + + public void shutdown() { + cancelPreviousCall(); + scheduler.shutdownNow(); + } + + public void cancelPreviousCall() { + var call = currentCall.get(); + if (call != null) { + call.cancel(); + } + } + + private BackgroundableProcessIndicator createProgressIndicator() { + return new BackgroundableProcessIndicator(project, + CodeGPTBundle.get("codeCompletion.progress.title"), null, null, true) { + @Override + protected void onRunningChange() { + if (isCanceled()) { + cancelPreviousCall(); + CodeGPTEditorManager.getInstance().disposeAllInlays(project); + } + } + }; + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CallRunnable.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CallRunnable.java new file mode 100644 index 00000000..0cb91fc9 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CallRunnable.java @@ -0,0 +1,11 @@ +package ee.carlrobert.codegpt.codecompletions; + +import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator; +import okhttp3.sse.EventSource; +import org.jetbrains.annotations.Nullable; + +@FunctionalInterface +public interface CallRunnable { + + EventSource call(@Nullable BackgroundableProcessIndicator progressIndicator); +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java new file mode 100644 index 00000000..d20028a5 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java @@ -0,0 +1,75 @@ +package ee.carlrobert.codegpt.codecompletions; + +import static java.util.Objects.requireNonNull; + +import com.intellij.notification.NotificationType; +import com.intellij.notification.Notifications; +import com.intellij.openapi.application.ApplicationManager; +import com.intellij.openapi.diagnostic.Logger; +import com.intellij.openapi.editor.Editor; +import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator; +import ee.carlrobert.codegpt.CodeGPTBundle; +import ee.carlrobert.codegpt.actions.OpenSettingsAction; +import ee.carlrobert.codegpt.ui.OverlayUtil; +import ee.carlrobert.llm.client.openai.completion.ErrorDetails; +import ee.carlrobert.llm.completion.CompletionEventListener; +import javax.annotation.ParametersAreNonnullByDefault; +import org.jetbrains.annotations.Nullable; + +@ParametersAreNonnullByDefault +class CodeCompletionEventListener implements CompletionEventListener { + + private static final Logger LOG = Logger.getInstance(CodeCompletionEventListener.class); + + private final Editor editor; + private final int caretOffset; + private final BackgroundableProcessIndicator progressIndicator; + + public CodeCompletionEventListener( + Editor editor, + int caretOffset, + @Nullable BackgroundableProcessIndicator progressIndicator) { + this.editor = editor; + this.caretOffset = caretOffset; + this.progressIndicator = progressIndicator; + } + + @Override + public void onComplete(StringBuilder messageBuilder) { + if (progressIndicator != null) { + progressIndicator.processFinish(); + } + + var editorManager = CodeGPTEditorManager.getInstance(); + editorManager.disposeEditorInlays(editor); + + var inlayText = messageBuilder.toString(); + if (!inlayText.isEmpty()) { + ApplicationManager.getApplication().invokeLater(() -> + CodeCompletionService.getInstance(requireNonNull(editor.getProject())) + .addInlays(editor, caretOffset, inlayText)); + } + } + + @Override + public void onError(ErrorDetails error, Throwable ex) { + LOG.error(error.getMessage(), ex); + if (progressIndicator != null) { + progressIndicator.processFinish(); + } + Notifications.Bus.notify(OverlayUtil.getDefaultNotification( + String.format( + CodeGPTBundle.get("notification.completionError.description"), + ex.getMessage()), + NotificationType.ERROR) + .addAction(new OpenSettingsAction()), editor.getProject()); + } + + @Override + public void onCancelled(StringBuilder messageBuilder) { + LOG.info("Completion cancelled"); + if (progressIndicator != null) { + progressIndicator.processFinish(); + } + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionListenerBinder.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionListenerBinder.java new file mode 100644 index 00000000..6d4a8c48 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionListenerBinder.java @@ -0,0 +1,120 @@ +package ee.carlrobert.codegpt.codecompletions; + +import com.intellij.openapi.Disposable; +import com.intellij.openapi.application.ApplicationManager; +import com.intellij.openapi.editor.Editor; +import com.intellij.openapi.editor.event.BulkAwareDocumentListener; +import com.intellij.openapi.editor.event.CaretEvent; +import com.intellij.openapi.editor.event.CaretListener; +import com.intellij.openapi.editor.event.DocumentEvent; +import com.intellij.openapi.editor.event.SelectionEvent; +import com.intellij.openapi.editor.event.SelectionListener; +import ee.carlrobert.codegpt.actions.CodeCompletionEnabledListener; +import ee.carlrobert.codegpt.settings.configuration.ConfigurationState; +import org.jetbrains.annotations.NotNull; +import org.jetbrains.annotations.Nullable; + +public class CodeCompletionListenerBinder implements Disposable { + + private final Editor editor; + + private @Nullable EditorDocumentListener documentListener; + private @Nullable EditorSelectionListener selectionListener; + private @Nullable EditorCaretListener caretListener; + + public CodeCompletionListenerBinder(Editor editor) { + this.editor = editor; + + if (ConfigurationState.getInstance().isCodeCompletionsEnabled()) { + addListeners(); + } + + ApplicationManager.getApplication() + .getMessageBus() + .connect() + .subscribe( + CodeCompletionEnabledListener.TOPIC, + (CodeCompletionEnabledListener) (completionsEnabled) -> { + if (completionsEnabled) { + addListeners(); + if (editor.getProject() != null) { + CodeCompletionService.getInstance(editor.getProject()) + .handleCompletions(editor, editor.getCaretModel().getOffset()); + } + } else { + removeListeners(); + } + }); + } + + private void addListeners() { + if (documentListener == null) { + documentListener = new EditorDocumentListener(); + editor.getDocument().addDocumentListener(documentListener); + } + if (selectionListener == null) { + selectionListener = new EditorSelectionListener(); + editor.getSelectionModel().addSelectionListener(selectionListener); + } + if (caretListener == null) { + caretListener = new EditorCaretListener(); + editor.getCaretModel().addCaretListener(caretListener); + } + } + + private void removeListeners() { + if (documentListener != null) { + editor.getDocument().removeDocumentListener(documentListener); + documentListener = null; + } + if (selectionListener != null) { + editor.getSelectionModel().removeSelectionListener(selectionListener); + selectionListener = null; + } + if (caretListener != null) { + editor.getCaretModel().removeCaretListener(caretListener); + caretListener = null; + } + } + + @Override + public void dispose() { + removeListeners(); + } + + private class EditorSelectionListener implements SelectionListener { + + @Override + public void selectionChanged(@NotNull SelectionEvent event) { + CodeGPTEditorManager.getInstance().disposeEditorInlays(editor); + } + } + + private class EditorCaretListener implements CaretListener { + + @Override + public void caretPositionChanged(@NotNull CaretEvent event) { + var project = editor.getProject(); + if (event.getCaret() == null || project == null) { + return; + } + + CodeGPTEditorManager.getInstance().disposeEditorInlays(editor); + CodeCompletionService.getInstance(project) + .handleCompletions(editor, event.getCaret().getOffset()); + } + } + + private class EditorDocumentListener implements BulkAwareDocumentListener { + + @Override + public void documentChangedNonBulk(@NotNull DocumentEvent event) { + var project = editor.getProject(); + if (project != null) { + CodeGPTEditorManager.getInstance().disposeEditorInlays(editor); + CodeCompletionService.getInstance(project) + .handleCompletions(editor, editor.getCaretModel().getOffset()); + } + } + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestProvider.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestProvider.java new file mode 100644 index 00000000..ab941750 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionRequestProvider.java @@ -0,0 +1,38 @@ +package ee.carlrobert.codegpt.codecompletions; + +import ee.carlrobert.codegpt.settings.state.LlamaSettingsState; +import ee.carlrobert.llm.client.llama.completion.LlamaCompletionRequest; +import ee.carlrobert.llm.client.openai.completion.request.OpenAITextCompletionRequest; +import javax.annotation.ParametersAreNonnullByDefault; + +@ParametersAreNonnullByDefault +public class CodeCompletionRequestProvider { + + private static final int MAX_TOKENS = 256; + + private final InfillRequestDetails details; + + public CodeCompletionRequestProvider(InfillRequestDetails details) { + this.details = details; + } + + public OpenAITextCompletionRequest buildOpenAIRequest() { + return new OpenAITextCompletionRequest.Builder(details.getPrefix()) + .setSuffix(details.getSuffix()) + .setStream(true) + .setMaxTokens(MAX_TOKENS) + .setTemperature(0.1) + .build(); + } + + public LlamaCompletionRequest buildLlamaRequest() { + var promptTemplate = LlamaSettingsState.getInstance().getInfillPromptTemplate(); + var prompt = promptTemplate.buildPrompt(details.getPrefix(), details.getSuffix()); + return new LlamaCompletionRequest.Builder(prompt) + .setN_predict(MAX_TOKENS) + .setStream(true) + .setTemperature(0.1) + .setStop(promptTemplate.getStopTokens()) + .build(); + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java new file mode 100644 index 00000000..0b376346 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java @@ -0,0 +1,265 @@ +package ee.carlrobert.codegpt.codecompletions; + +import static com.intellij.openapi.components.Service.Level.PROJECT; +import static ee.carlrobert.codegpt.CodeGPTKeys.MULTI_LINE_INLAY; +import static ee.carlrobert.codegpt.CodeGPTKeys.SINGLE_LINE_INLAY; +import static java.util.Objects.requireNonNull; +import static java.util.stream.Collectors.toList; + +import com.intellij.openapi.Disposable; +import com.intellij.openapi.actionSystem.ActionManager; +import com.intellij.openapi.actionSystem.AnAction; +import com.intellij.openapi.actionSystem.AnActionEvent; +import com.intellij.openapi.actionSystem.KeyboardShortcut; +import com.intellij.openapi.application.ApplicationManager; +import com.intellij.openapi.application.ReadAction; +import com.intellij.openapi.command.WriteCommandAction; +import com.intellij.openapi.components.Service; +import com.intellij.openapi.diagnostic.Logger; +import com.intellij.openapi.editor.Document; +import com.intellij.openapi.editor.Editor; +import com.intellij.openapi.editor.EditorCustomElementRenderer; +import com.intellij.openapi.editor.Inlay; +import com.intellij.openapi.editor.InlayModel; +import com.intellij.openapi.keymap.KeymapManager; +import com.intellij.openapi.project.Project; +import com.intellij.openapi.util.TextRange; +import com.intellij.psi.PsiClass; +import com.intellij.psi.PsiDocumentManager; +import com.intellij.psi.PsiElement; +import com.intellij.psi.PsiFile; +import com.intellij.psi.PsiMethod; +import com.intellij.psi.PsiWhiteSpace; +import com.intellij.util.concurrency.annotations.RequiresBackgroundThread; +import com.intellij.util.concurrency.annotations.RequiresEdt; +import com.intellij.util.concurrency.annotations.RequiresReadLock; +import com.intellij.util.concurrency.annotations.RequiresWriteLock; +import ee.carlrobert.codegpt.actions.CodeCompletionEnabledListener; +import ee.carlrobert.codegpt.completions.CompletionRequestService; +import ee.carlrobert.codegpt.settings.configuration.ConfigurationState; +import ee.carlrobert.codegpt.util.EditorUtil; +import ee.carlrobert.llm.completion.CompletionEventListener; +import java.awt.event.KeyEvent; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.TimeUnit; +import javax.annotation.ParametersAreNonnullByDefault; +import javax.swing.KeyStroke; +import okhttp3.sse.EventSource; +import org.jetbrains.annotations.NotNull; + +@Service(PROJECT) +@ParametersAreNonnullByDefault +public final class CodeCompletionService implements Disposable { + + public static final String APPLY_INLAY_ACTION_ID = "ApplyInlayAction"; + public static final int MAX_OFFSET = 4000; + + private static final Logger LOG = Logger.getInstance(CodeCompletionService.class); + + private final CallDebouncer callDebouncer; + + private CodeCompletionService(Project project) { + this.callDebouncer = new CallDebouncer(project); + ApplicationManager.getApplication() + .getMessageBus() + .connect() + .subscribe( + CodeCompletionEnabledListener.TOPIC, + (CodeCompletionEnabledListener) (completionsEnabled) -> { + if (!completionsEnabled) { + callDebouncer.cancelPreviousCall(); + } + }); + } + + public static CodeCompletionService getInstance(Project project) { + return project.getService(CodeCompletionService.class); + } + + public boolean isCompletionAllowed(PsiElement elementAtCaret) { + return elementAtCaret instanceof PsiWhiteSpace; + } + + public void handleCompletions(Editor editor, int offset) { + Project project = editor.getProject(); + if (project == null + || project.isDisposed() + || !ConfigurationState.getInstance().isCodeCompletionsEnabled() + || !EditorUtil.isSelectedEditor(editor) + || editor.isViewer() + || editor.isOneLineMode() + ) { + return; + } + + var document = editor.getDocument(); + PsiFile psiFile = PsiDocumentManager.getInstance(project).getPsiFile(document); + if (psiFile == null) { + return; + } + + PsiElement elementAtCaret = ReadAction.compute(() -> psiFile.findElementAt(offset)); + var completionService = CodeCompletionService.getInstance(project); + if (!completionService.isCompletionAllowed(elementAtCaret)) { + return; + } + + callDebouncer.debounce( + Void.class, + (progressIndicator) -> completionService.fetchCodeCompletion( + elementAtCaret, + offset, + document, + new CodeCompletionEventListener(editor, offset, progressIndicator)), + 500, + TimeUnit.MILLISECONDS); + } + + /** + * Fetches code-completion (FIM) for the given position ({@code offsetInFile}) in the file.
+ * By default tries to find an enclosing {@link PsiMethod} or {@link PsiClass} for the given + * {@code offsetInFile} and only uses their content instead of the entire file's content. If no + * such enclosing {@link PsiElement} can be found, the file's entire content is used instead. + * + * @param elementAtCaret PsiElement at caret + * @param offsetInFile Global offset in the file. + * @param document If the offset is not enclosed in a {@link PsiMethod} nor a + * {@link PsiClass}, the entire file content is used for completion. + * @return Completion String + */ + @RequiresBackgroundThread + public EventSource fetchCodeCompletion( + PsiElement elementAtCaret, + int offsetInFile, + Document document, + CompletionEventListener eventListener) { + InfillRequestDetails requestDetails = tryFindEnclosingPsiElementTextRange( + List.of(PsiMethod.class, PsiClass.class), elementAtCaret) + .map(textRange -> createInfillRequest( + document, + offsetInFile, + textRange.getStartOffset(), + textRange.getEndOffset()) + ) + .orElse(createInfillRequest(document, offsetInFile)); + return CompletionRequestService.getInstance() + .getCodeCompletionAsync(requestDetails, eventListener); + } + + @RequiresEdt + public void addInlays(Editor editor, int caretOffset, String inlayText) { + List linesList = inlayText.lines().collect(toList()); + String firstLine = linesList.get(0); + String restOfLines = linesList.size() > 1 + ? String.join("\n", linesList.subList(1, linesList.size())) + : null; + InlayModel inlayModel = editor.getInlayModel(); + + if (!firstLine.isEmpty()) { + editor.putUserData(SINGLE_LINE_INLAY, inlayModel.addInlineElement( + caretOffset, + true, + Integer.MAX_VALUE, + new InlayInlineElementRenderer(firstLine))); + } + + if (restOfLines != null && !restOfLines.isEmpty()) { + editor.putUserData(MULTI_LINE_INLAY, inlayModel.addBlockElement( + caretOffset, + true, + false, + Integer.MAX_VALUE, + new InlayBlockElementRenderer(restOfLines))); + } + + registerApplyCompletionAction(() -> WriteCommandAction.runWriteCommandAction( + editor.getProject(), + () -> applyCompletion(editor, inlayText))); + } + + @RequiresWriteLock + private void applyCompletion(Editor editor, String text) { + if (editor.isDisposed()) { + LOG.warn("Editor is already disposed"); + return; + } + + var inlayKeys = List.of(SINGLE_LINE_INLAY, MULTI_LINE_INLAY); + for (var key : inlayKeys) { + Inlay inlay = editor.getUserData(key); + if (inlay != null) { + applyCompletion(editor, text, inlay.getOffset()); + CodeGPTEditorManager.getInstance().disposeEditorInlays(editor); + return; + } + } + } + + @RequiresWriteLock + private void applyCompletion(Editor editor, String text, int offset) { + Document document = editor.getDocument(); + document.insertString(offset, text); + editor.getCaretModel().moveToOffset(offset + text.length()); + EditorUtil.reformatDocument( + requireNonNull(editor.getProject()), + document, + offset, + offset + text.length()); + } + + @RequiresReadLock + private Optional tryFindEnclosingPsiElementTextRange( + List> types, + PsiElement elementAtCaret) { + return ReadAction.compute(() -> { + var element = elementAtCaret; + while (element != null) { + for (Class type : types) { + if (type.isInstance(element)) { + return Optional.of(element.getTextRange()); + } + } + element = element.getParent(); + } + + return Optional.empty(); + }); + } + + @Override + public void dispose() { + callDebouncer.shutdown(); + } + + private void registerApplyCompletionAction(Runnable onApply) { + var actionManager = ActionManager.getInstance(); + actionManager.registerAction( + APPLY_INLAY_ACTION_ID, + new AnAction() { + @Override + public void actionPerformed(@NotNull AnActionEvent e) { + onApply.run(); + } + }); + KeymapManager.getInstance().getActiveKeymap().addShortcut( + APPLY_INLAY_ACTION_ID, + new KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_TAB, 0), null)); + } + + private static InfillRequestDetails createInfillRequest(Document document, int offsetInFile) { + int begin = Integer.max(0, offsetInFile - MAX_OFFSET); + int end = Integer.min(document.getTextLength(), offsetInFile + MAX_OFFSET); + return createInfillRequest(document, offsetInFile, begin, end); + } + + private static InfillRequestDetails createInfillRequest( + Document document, + int caretOffset, + int start, + int end) { + return new InfillRequestDetails( + document.getText(new TextRange(start, caretOffset)), + document.getText(new TextRange(caretOffset, end))); + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeGPTEditorListener.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeGPTEditorListener.java new file mode 100644 index 00000000..0f19c88a --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeGPTEditorListener.java @@ -0,0 +1,16 @@ +package ee.carlrobert.codegpt.codecompletions; + +import com.intellij.openapi.editor.Editor; +import com.intellij.openapi.editor.event.EditorFactoryEvent; +import com.intellij.openapi.editor.event.EditorFactoryListener; +import com.intellij.openapi.editor.ex.util.EditorUtil; +import org.jetbrains.annotations.NotNull; + +public class CodeGPTEditorListener implements EditorFactoryListener { + + @Override + public void editorCreated(@NotNull EditorFactoryEvent event) { + Editor editor = event.getEditor(); + EditorUtil.disposeWithEditor(editor, new CodeCompletionListenerBinder(editor)); + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeGPTEditorManager.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeGPTEditorManager.java new file mode 100644 index 00000000..e9e99a69 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeGPTEditorManager.java @@ -0,0 +1,51 @@ +package ee.carlrobert.codegpt.codecompletions; + +import static ee.carlrobert.codegpt.CodeGPTKeys.MULTI_LINE_INLAY; +import static ee.carlrobert.codegpt.CodeGPTKeys.SINGLE_LINE_INLAY; + +import com.intellij.openapi.actionSystem.ActionManager; +import com.intellij.openapi.application.ApplicationManager; +import com.intellij.openapi.command.WriteCommandAction; +import com.intellij.openapi.components.Service; +import com.intellij.openapi.editor.Editor; +import com.intellij.openapi.editor.EditorCustomElementRenderer; +import com.intellij.openapi.editor.Inlay; +import com.intellij.openapi.fileEditor.FileEditor; +import com.intellij.openapi.fileEditor.FileEditorManager; +import com.intellij.openapi.fileEditor.TextEditor; +import com.intellij.openapi.project.Project; +import com.intellij.openapi.util.Key; + +@Service +public final class CodeGPTEditorManager { + + private CodeGPTEditorManager() { + } + + public static CodeGPTEditorManager getInstance() { + return ApplicationManager.getApplication().getService(CodeGPTEditorManager.class); + } + + public void disposeAllInlays(Project project) { + var allFileEditors = FileEditorManager.getInstance(project).getAllEditors(); + for (FileEditor fileEditor : allFileEditors) { + if (fileEditor instanceof TextEditor) { + disposeEditorInlays(((TextEditor) fileEditor).getEditor()); + } + } + } + + public void disposeEditorInlays(Editor editor) { + ActionManager.getInstance().unregisterAction(CodeCompletionService.APPLY_INLAY_ACTION_ID); + disposeInlay(editor, SINGLE_LINE_INLAY); + disposeInlay(editor, MULTI_LINE_INLAY); + } + + private void disposeInlay(Editor editor, Key> inlayKey) { + Inlay inlay = editor.getUserData(inlayKey); + if (inlay != null) { + WriteCommandAction.runWriteCommandAction(editor.getProject(), inlay::dispose); + editor.putUserData(inlayKey, null); + } + } +} diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.java new file mode 100644 index 00000000..f31e9ca7 --- /dev/null +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.java @@ -0,0 +1,57 @@ +package ee.carlrobert.codegpt.codecompletions; + +import static java.lang.String.format; + +import java.util.List; +import org.jetbrains.annotations.Nullable; + +public enum InfillPromptTemplate { + + OPENAI("OpenAI") { + @Override + public String buildPrompt(String prefix, String suffix) { + return format("<|fim_prefix|> %s <|fim_suffix|>%s <|fim_middle|>", prefix, suffix); + } + }, + LLAMA("Llama", List.of("")) { + @Override + public String buildPrompt(String prefix, String suffix) { + return format("
 %s %s ", prefix, suffix);
+    }
+  },
+  STABILITY("Stability AI", List.of("<|endoftext|>")) {
+    @Override
+    public String buildPrompt(String prefix, String suffix) {
+      return format("%s%s", prefix, suffix);
+    }
+  },
+  DEEPSEEK_CODER("DeepSeek Coder") {
+    @Override
+    public String buildPrompt(String prefix, String suffix) {
+      return format("<|fim_begin|>%s<|fim_hole|>%s<|fim_end|>", prefix, suffix);
+    }
+  };
+
+  private final String label;
+  private final @Nullable List stopTokens;
+
+  InfillPromptTemplate(String label) {
+    this(label, null);
+  }
+
+  InfillPromptTemplate(String label, @Nullable List stopTokens) {
+    this.label = label;
+    this.stopTokens = stopTokens;
+  }
+
+  public abstract String buildPrompt(String prefix, String suffix);
+
+  @Override
+  public String toString() {
+    return label;
+  }
+
+  public List getStopTokens() {
+    return stopTokens;
+  }
+}
diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.java
new file mode 100644
index 00000000..de65554c
--- /dev/null
+++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.java
@@ -0,0 +1,20 @@
+package ee.carlrobert.codegpt.codecompletions;
+
+public class InfillRequestDetails {
+
+  private final String prefix;
+  private final String suffix;
+
+  public InfillRequestDetails(String prefix, String suffix) {
+    this.prefix = prefix;
+    this.suffix = suffix;
+  }
+
+  public String getPrefix() {
+    return prefix;
+  }
+
+  public String getSuffix() {
+    return suffix;
+  }
+}
diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayBlockElementRenderer.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayBlockElementRenderer.java
new file mode 100644
index 00000000..b399596b
--- /dev/null
+++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayBlockElementRenderer.java
@@ -0,0 +1,57 @@
+package ee.carlrobert.codegpt.codecompletions;
+
+import static java.util.stream.Collectors.toList;
+
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.EditorCustomElementRenderer;
+import com.intellij.openapi.editor.Inlay;
+import com.intellij.openapi.editor.colors.EditorFontType;
+import com.intellij.openapi.editor.impl.FontInfo;
+import com.intellij.openapi.editor.markup.TextAttributes;
+import com.intellij.ui.JBColor;
+import java.awt.Font;
+import java.awt.Graphics2D;
+import java.awt.font.FontRenderContext;
+import java.awt.geom.Rectangle2D;
+import java.util.Comparator;
+import java.util.List;
+import org.jetbrains.annotations.NotNull;
+
+public class InlayBlockElementRenderer extends InlayElementRenderer {
+
+  protected InlayBlockElementRenderer(String inlayText) {
+    super(inlayText);
+  }
+
+  @Override
+  public void paint(
+      @NotNull Inlay inlay,
+      @NotNull Graphics2D g,
+      @NotNull Rectangle2D targetRegion,
+      @NotNull TextAttributes textAttributes) {
+    Editor editor = inlay.getEditor();
+    Font font = editor.getColorsScheme()
+        .getFont(EditorFontType.PLAIN)
+        .deriveFont(Font.ITALIC);
+    g.setFont(font);
+    g.setColor(JBColor.GRAY);
+
+    int x = (int) targetRegion.getX();
+    int fontBaseLineOffset = (int) calculateFontBaseLineOffset(font, editor);
+    List lines = inlayText.lines().collect(toList());
+    for (int i = 0; i < lines.size(); i++) {
+      String line = lines.get(i);
+      int y = (int) targetRegion.getY() + fontBaseLineOffset + i * editor.getLineHeight();
+      g.drawString(line, x, y);
+    }
+  }
+
+  public double calculateFontBaseLineOffset(Font font, Editor editor) {
+    FontRenderContext fontRenderContext =
+        FontInfo.getFontRenderContext(editor.getContentComponent());
+    Rectangle2D visualBounds = font.createGlyphVector(fontRenderContext, "Abc").getVisualBounds();
+    double fontBaseline = visualBounds.getHeight();
+    double linePadding = (editor.getLineHeight() - fontBaseline) / 2;
+    return Math.ceil(fontBaseline + linePadding);
+  }
+}
diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayElementRenderer.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayElementRenderer.java
new file mode 100644
index 00000000..e49e08ec
--- /dev/null
+++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayElementRenderer.java
@@ -0,0 +1,38 @@
+package ee.carlrobert.codegpt.codecompletions;
+
+import com.intellij.openapi.editor.EditorCustomElementRenderer;
+import com.intellij.openapi.editor.Inlay;
+import com.intellij.openapi.editor.colors.EditorFontType;
+import java.util.Comparator;
+import org.jetbrains.annotations.NotNull;
+
+public abstract class InlayElementRenderer implements EditorCustomElementRenderer {
+
+  protected final String inlayText;
+
+  protected InlayElementRenderer(String inlayText) {
+    this.inlayText = inlayText;
+  }
+
+  @Override
+  public int calcWidthInPixels(@NotNull Inlay inlay) {
+    var longestLine = getInlayText().lines()
+        .max(Comparator.comparingInt(String::length))
+        .orElse("");
+    var editor = inlay.getEditor();
+    return editor.getContentComponent()
+        .getFontMetrics(editor.getColorsScheme().getFont(EditorFontType.PLAIN))
+        .stringWidth(longestLine);
+  }
+
+  @Override
+  public int calcHeightInPixels(@NotNull Inlay inlay) {
+    int lineHeight = inlay.getEditor().getLineHeight();
+    int linesCount = (int) inlayText.lines().count();
+    return lineHeight * linesCount;
+  }
+
+  public String getInlayText() {
+    return inlayText;
+  }
+}
diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayInlineElementRenderer.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayInlineElementRenderer.java
new file mode 100644
index 00000000..74ff771d
--- /dev/null
+++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/InlayInlineElementRenderer.java
@@ -0,0 +1,37 @@
+package ee.carlrobert.codegpt.codecompletions;
+
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.EditorCustomElementRenderer;
+import com.intellij.openapi.editor.Inlay;
+import com.intellij.openapi.editor.colors.EditorFontType;
+import com.intellij.openapi.editor.markup.TextAttributes;
+import com.intellij.ui.JBColor;
+import java.awt.Font;
+import java.awt.Graphics2D;
+import java.awt.geom.Rectangle2D;
+import org.jetbrains.annotations.NotNull;
+
+public class InlayInlineElementRenderer extends InlayElementRenderer {
+
+  protected InlayInlineElementRenderer(String inlayText) {
+    super(inlayText);
+  }
+
+  @Override
+  public void paint(
+      @NotNull Inlay inlay,
+      @NotNull Graphics2D g,
+      @NotNull Rectangle2D targetRegion,
+      @NotNull TextAttributes textAttributes) {
+    Editor editor = inlay.getEditor();
+    Font font = editor.getColorsScheme()
+        .getFont(EditorFontType.PLAIN)
+        .deriveFont(Font.ITALIC);
+    g.setFont(font);
+    g.setColor(JBColor.GRAY);
+    g.drawString(
+        inlayText,
+        (int) targetRegion.getX(),
+        (int) targetRegion.getY() + editor.getAscent());
+  }
+}
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java
index d338cdf4..9a39c79b 100644
--- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java
+++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestHandler.java
@@ -109,6 +109,11 @@ public class CompletionRequestHandler {
         completionResponseEventListener.handleCompleted(messageBuilder.toString(), callParameters);
       }
 
+      @Override
+      public void onCancelled(StringBuilder messageBuilder) {
+        completionResponseEventListener.handleCompleted(messageBuilder.toString(), callParameters);
+      }
+
       @Override
       public void onError(ErrorDetails error, Throwable ex) {
         try {
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java
index 8e41322f..e64bd838 100644
--- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java
+++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java
@@ -27,7 +27,6 @@ import ee.carlrobert.embedding.EmbeddingsService;
 import ee.carlrobert.embedding.ReferencedFile;
 import ee.carlrobert.llm.client.llama.completion.LlamaCompletionRequest;
 import ee.carlrobert.llm.client.openai.completion.OpenAIChatCompletionModel;
-import ee.carlrobert.llm.client.openai.completion.OpenAICompletionRequest;
 import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionMessage;
 import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionRequest;
 import ee.carlrobert.llm.client.you.completion.YouCompletionRequest;
@@ -52,6 +51,9 @@ public class CompletionRequestProvider {
   public static final String FIX_COMPILE_ERRORS_SYSTEM_PROMPT = getResourceContent(
       "/prompts/fix-compile-errors.txt");
 
+  public static final String INLINE_COMPLETION_PROMPT = getResourceContent(
+      "/prompts/inline-completion-prompt.txt");
+
   private final EncodingManager encodingManager = EncodingManager.getInstance();
   private final EmbeddingsService embeddingsService;
   private final Conversation conversation;
@@ -80,7 +82,7 @@ public class CompletionRequestProvider {
         .replace("{QUESTION}", userPrompt);
   }
 
-  public static OpenAICompletionRequest buildOpenAILookupCompletionRequest(
+  public static OpenAIChatCompletionRequest buildOpenAILookupCompletionRequest(
       String context) {
     return new OpenAIChatCompletionRequest.Builder(
         List.of(
@@ -147,17 +149,19 @@ public class CompletionRequestProvider {
       CallParameters callParameters,
       boolean useContextualSearch,
       @Nullable String overriddenPath) {
+
     var builder = new OpenAIChatCompletionRequest.Builder(
         buildMessages(model, callParameters, useContextualSearch))
         .setModel(model)
         .setMaxTokens(ConfigurationState.getInstance().getMaxTokens())
+        .setStream(true)
         .setTemperature(ConfigurationState.getInstance().getTemperature());
 
     if (overriddenPath != null) {
       builder.setOverriddenPath(overriddenPath);
     }
 
-    return (OpenAIChatCompletionRequest) builder.build();
+    return builder.build();
   }
 
   public List buildMessages(
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java
index d11334e2..f27347e9 100644
--- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java
+++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestService.java
@@ -1,15 +1,17 @@
 package ee.carlrobert.codegpt.completions;
 
+import static ee.carlrobert.codegpt.settings.service.ServiceType.AZURE;
 import static ee.carlrobert.codegpt.settings.service.ServiceType.LLAMA_CPP;
 import static ee.carlrobert.codegpt.settings.service.ServiceType.OPENAI;
 import static ee.carlrobert.codegpt.settings.service.ServiceType.YOU;
 
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.components.Service;
+import ee.carlrobert.codegpt.codecompletions.CodeCompletionRequestProvider;
+import ee.carlrobert.codegpt.codecompletions.InfillRequestDetails;
 import ee.carlrobert.codegpt.credentials.AzureCredentialsManager;
 import ee.carlrobert.codegpt.credentials.OpenAICredentialsManager;
 import ee.carlrobert.codegpt.settings.configuration.ConfigurationState;
-import ee.carlrobert.codegpt.settings.service.ServiceType;
 import ee.carlrobert.codegpt.settings.state.AzureSettingsState;
 import ee.carlrobert.codegpt.settings.state.OpenAISettingsState;
 import ee.carlrobert.codegpt.settings.state.SettingsState;
@@ -69,6 +71,22 @@ public final class CompletionRequestService {
     }
   }
 
+  public EventSource getCodeCompletionAsync(
+      InfillRequestDetails requestDetails,
+      CompletionEventListener eventListener) {
+    var requestProvider = new CodeCompletionRequestProvider(requestDetails);
+    switch (SettingsState.getInstance().getSelectedService()) {
+      case OPENAI:
+        return CompletionClientProvider.getOpenAIClient()
+            .getCompletionAsync(requestProvider.buildOpenAIRequest(), eventListener);
+      case LLAMA_CPP:
+        return CompletionClientProvider.getLlamaClient()
+            .getChatCompletionAsync(requestProvider.buildLlamaRequest(), eventListener);
+      default:
+        throw new IllegalArgumentException("Code completion not supported for selected service");
+    }
+  }
+
   public void generateCommitMessageAsync(
       String prompt,
       CompletionEventListener eventListener) {
@@ -79,10 +97,10 @@ public final class CompletionRequestService {
         .setModel(OpenAISettingsState.getInstance().getModel())
         .build();
     var selectedService = SettingsState.getInstance().getSelectedService();
-    if (selectedService == ServiceType.OPENAI) {
+    if (selectedService == OPENAI) {
       CompletionClientProvider.getOpenAIClient().getChatCompletionAsync(request, eventListener);
     }
-    if (selectedService == ServiceType.AZURE) {
+    if (selectedService == AZURE) {
       CompletionClientProvider.getAzureClient().getChatCompletionAsync(request, eventListener);
     }
   }
@@ -106,10 +124,10 @@ public final class CompletionRequestService {
 
   public boolean isRequestAllowed() {
     var selectedService = SettingsState.getInstance().getSelectedService();
-    if (selectedService == ServiceType.AZURE) {
+    if (selectedService == AZURE) {
       return AzureCredentialsManager.getInstance().isCredentialSet();
     }
-    if (selectedService == ServiceType.OPENAI) {
+    if (selectedService == OPENAI) {
       return OpenAICredentialsManager.getInstance().isApiKeySet();
     }
     return true;
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/ConversationType.java b/src/main/java/ee/carlrobert/codegpt/completions/ConversationType.java
index 59e4ee76..efff0bbf 100644
--- a/src/main/java/ee/carlrobert/codegpt/completions/ConversationType.java
+++ b/src/main/java/ee/carlrobert/codegpt/completions/ConversationType.java
@@ -6,4 +6,5 @@ public enum ConversationType {
   EDITOR_ACTION,
   FIX_COMPILE_ERRORS,
   MULTI_FILE,
+  INLINE_COMPLETION,
 }
diff --git a/src/main/java/ee/carlrobert/codegpt/completions/TotalUsageExceededException.java b/src/main/java/ee/carlrobert/codegpt/completions/TotalUsageExceededException.java
index ffdfbf08..3143455f 100644
--- a/src/main/java/ee/carlrobert/codegpt/completions/TotalUsageExceededException.java
+++ b/src/main/java/ee/carlrobert/codegpt/completions/TotalUsageExceededException.java
@@ -1,4 +1,5 @@
 package ee.carlrobert.codegpt.completions;
 
 class TotalUsageExceededException extends RuntimeException {
+
 }
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 ab51ac2d..5e6d7958 100644
--- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java
+++ b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java
@@ -48,8 +48,10 @@ public class ConfigurationComponent {
   private final JBCheckBox autoFormattingCheckBox;
   private final JTextArea systemPromptTextArea;
   private final JTextArea commitMessagePromptTextArea;
+  private final JTextArea inlineCompletionPromptTextArea;
   private final IntegerField maxTokensField;
   private final JBTextField temperatureField;
+  private final JBTextField inlineDelayField;
 
   public ConfigurationComponent(Disposable parentDisposable, ConfigurationState configuration) {
     table = new JBTable(new DefaultTableModel(
@@ -68,7 +70,8 @@ public class ConfigurationComponent {
     temperatureField = new JBTextField(12);
     temperatureField.setText(String.valueOf(configuration.getTemperature()));
 
-    var temperatureFieldValidator = createInputValidator(parentDisposable, temperatureField);
+    var temperatureFieldValidator = createTemperatureInputValidator(parentDisposable,
+        temperatureField);
     temperatureField.getDocument().addDocumentListener(new DocumentListener() {
       @Override
       public void insertUpdate(DocumentEvent e) {
@@ -106,6 +109,33 @@ public class ConfigurationComponent {
     commitMessagePromptTextArea.setLineWrap(true);
     commitMessagePromptTextArea.setBorder(JBUI.Borders.empty(8, 4));
 
+    inlineCompletionPromptTextArea = new JTextArea(configuration.getInlineCompletionPrompt(), 3,
+        60);
+    inlineCompletionPromptTextArea.setLineWrap(true);
+    inlineCompletionPromptTextArea.setBorder(JBUI.Borders.empty(8, 4));
+
+    inlineDelayField = new JBTextField(12);
+    inlineDelayField.setText(String.valueOf(configuration.getTemperature()));
+
+    var inlineDelayFieldValidator = createInlineDelayInputValidator(parentDisposable,
+        inlineDelayField);
+    inlineDelayField.getDocument().addDocumentListener(new DocumentListener() {
+      @Override
+      public void insertUpdate(DocumentEvent e) {
+        inlineDelayFieldValidator.revalidate();
+      }
+
+      @Override
+      public void removeUpdate(DocumentEvent e) {
+        inlineDelayFieldValidator.revalidate();
+      }
+
+      @Override
+      public void changedUpdate(DocumentEvent e) {
+        inlineDelayFieldValidator.revalidate();
+      }
+    });
+
     checkForPluginUpdatesCheckBox = new JBCheckBox(
         CodeGPTBundle.get("configurationConfigurable.checkForPluginUpdates.label"),
         configuration.isCheckForPluginUpdates());
@@ -135,6 +165,10 @@ public class ConfigurationComponent {
             CodeGPTBundle.get("configurationConfigurable.section.commitMessage.title")))
         .addComponent(createCommitMessageConfigurationForm())
         .addComponentFillVertically(new JPanel(), 0)
+        .addComponent(new TitledSeparator(
+            CodeGPTBundle.get("configurationConfigurable.section.inlineCompletion.title")))
+        .addComponent(createInlineCompletionConfigurationForm())
+        .addComponentFillVertically(new JPanel(), 0)
         .getPanel();
   }
 
@@ -227,7 +261,28 @@ public class ConfigurationComponent {
     return form;
   }
 
-  private ComponentValidator createInputValidator(
+  private JPanel createInlineCompletionConfigurationForm() {
+    var formBuilder = FormBuilder.createFormBuilder();
+    addAssistantFormLabeledComponent(
+        formBuilder,
+        "configurationConfigurable.section.inlineCompletion.systemPromptField.label",
+        "configurationConfigurable.section.inlineCompletion.systemPromptField.comment",
+        JBUI.Panels
+            .simplePanel(inlineCompletionPromptTextArea)
+            .withBorder(JBUI.Borders.customLine(
+                JBUI.CurrentTheme.CustomFrameDecorations.separatorForeground())));
+    formBuilder.addVerticalGap(8);
+    addAssistantFormLabeledComponent(
+        formBuilder,
+        "configurationConfigurable.section.inlineCompletion.delay.label",
+        "configurationConfigurable.section.inlineCompletion.delay.comment",
+        inlineDelayField);
+    var form = formBuilder.getPanel();
+    form.setBorder(JBUI.Borders.emptyLeft(16));
+    return form;
+  }
+
+  private ComponentValidator createTemperatureInputValidator(
       Disposable parentDisposable,
       JBTextField component) {
     var validator = new ComponentValidator(parentDisposable)
@@ -254,6 +309,33 @@ public class ConfigurationComponent {
     return validator;
   }
 
+  private ComponentValidator createInlineDelayInputValidator(
+      Disposable parentDisposable,
+      JBTextField component) {
+    var validator = new ComponentValidator(parentDisposable)
+        .withValidator(() -> {
+          var valueText = component.getText();
+          try {
+            var value = Integer.parseInt(valueText);
+            if (value <= 0) {
+              return new ValidationInfo(
+                  CodeGPTBundle.get("validation.error.mustBeGreaterThanZero"),
+                  component);
+            }
+          } catch (NumberFormatException e) {
+            return new ValidationInfo(
+                CodeGPTBundle.get("validation.error.mustBeNumber"),
+                component);
+          }
+
+          return null;
+        })
+        .andStartOnFocusLost()
+        .installOn(component);
+    validator.enableValidation();
+    return validator;
+  }
+
   private DefaultTableModel getModel() {
     return (DefaultTableModel) table.getModel();
   }
@@ -280,6 +362,23 @@ public class ConfigurationComponent {
     return commitMessagePromptTextArea.getText();
   }
 
+  public void setInlineCompletionPrompt(String inlineCompletionPrompt) {
+    inlineCompletionPromptTextArea.setText(inlineCompletionPrompt);
+  }
+
+  public String getInlineCompletionPrompt() {
+    return inlineCompletionPromptTextArea.getText();
+  }
+
+
+  public int getInlineDelay() {
+    return Integer.parseInt(inlineDelayField.getText());
+  }
+
+  public void setInlineDelay(int inlineDelay) {
+    inlineDelayField.setText(String.valueOf(inlineDelay));
+  }
+
   public double getTemperature() {
     return Double.parseDouble(temperatureField.getText());
   }
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 baa62886..a225f578 100644
--- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java
+++ b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationConfigurable.java
@@ -36,9 +36,12 @@ public class ConfigurationConfigurable implements Configurable {
     return !configurationComponent.getTableData().equals(configuration.getTableData())
         || configurationComponent.getMaxTokens() != configuration.getMaxTokens()
         || configurationComponent.getTemperature() != configuration.getTemperature()
+        || configurationComponent.getInlineDelay() != configuration.getInlineDelay()
         || !configurationComponent.getSystemPrompt().equals(configuration.getSystemPrompt())
         || !configurationComponent.getCommitMessagePrompt()
         .equals(configuration.getCommitMessagePrompt())
+        || !configurationComponent.getInlineCompletionPrompt()
+        .equals(configuration.getInlineCompletionPrompt())
         || configurationComponent.isCheckForPluginUpdates()
         != configuration.isCheckForPluginUpdates()
         || configurationComponent.isCreateNewChatOnEachAction()
@@ -57,6 +60,8 @@ public class ConfigurationConfigurable implements Configurable {
     configuration.setTemperature(configurationComponent.getTemperature());
     configuration.setSystemPrompt(configurationComponent.getSystemPrompt());
     configuration.setCommitMessagePrompt(configurationComponent.getCommitMessagePrompt());
+    configuration.setInlineCompletionPrompt(configurationComponent.getInlineCompletionPrompt());
+    configuration.setInlineDelay(configurationComponent.getInlineDelay());
     configuration.setCheckForPluginUpdates(configurationComponent.isCheckForPluginUpdates());
     configuration.setCreateNewChatOnEachAction(
         configurationComponent.isCreateNewChatOnEachAction());
@@ -74,6 +79,8 @@ public class ConfigurationConfigurable implements Configurable {
     configurationComponent.setTemperature(configuration.getTemperature());
     configurationComponent.setSystemPrompt(configuration.getSystemPrompt());
     configurationComponent.setCommitMessagePrompt(configuration.getCommitMessagePrompt());
+    configurationComponent.setInlineCompletionPrompt(configuration.getInlineCompletionPrompt());
+    configurationComponent.setInlineDelay(configuration.getInlineDelay());
     configurationComponent.setCheckForPluginUpdates(configuration.isCheckForPluginUpdates());
     configurationComponent.setCreateNewChatOnEachAction(
         configuration.isCreateNewChatOnEachAction());
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 3904433f..6c10f548 100644
--- a/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java
+++ b/src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationState.java
@@ -2,6 +2,7 @@ package ee.carlrobert.codegpt.settings.configuration;
 
 import static ee.carlrobert.codegpt.completions.CompletionRequestProvider.COMPLETION_SYSTEM_PROMPT;
 import static ee.carlrobert.codegpt.completions.CompletionRequestProvider.GENERATE_COMMIT_MESSAGE_SYSTEM_PROMPT;
+import static ee.carlrobert.codegpt.completions.CompletionRequestProvider.INLINE_COMPLETION_PROMPT;
 
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.components.PersistentStateComponent;
@@ -20,14 +21,17 @@ public class ConfigurationState implements PersistentStateComponent tableData = EditorActionsUtil.DEFAULT_ACTIONS;
 
   public static ConfigurationState getInstance() {
@@ -61,6 +65,22 @@ public class ConfigurationState implements PersistentStateComponent splitCodeBlocks(String inputMarkdown) {
     List result = new ArrayList<>();
diff --git a/src/main/resources/META-INF/plugin.xml b/src/main/resources/META-INF/plugin.xml
index 66707add..1f2564b0 100644
--- a/src/main/resources/META-INF/plugin.xml
+++ b/src/main/resources/META-INF/plugin.xml
@@ -36,9 +36,12 @@
         
         
         
+        
         
         
+        
     
 
     messages.codegpt
@@ -70,6 +73,30 @@
               anchor="before"/>
             
         
+
+        
+            
+            
+        
+        
+            
+            
+        
+        
+            
+            
+        
+
+        
+            
+            
+            
+            
+        
+
         
             
             The {pre}, {suf} and {mid} are replaced depending on the used Model's FIM template.
+configurationConfigurable.section.inlineCompletion.delay.label=Delay:
+configurationConfigurable.section.inlineCompletion.delay.comment=Inline completion is requested if user is idle for x milliseconds
 advancedSettingsConfigurable.displayName=CodeGPT: Advanced Settings
 advancedSettingsConfigurable.proxy.title=HTTP/SOCKS Proxy
 advancedSettingsConfigurable.proxy.typeComboBoxField.label=Proxy:
@@ -149,6 +156,7 @@ validation.error.fieldRequired=This field is required.
 validation.error.invalidEmail=The email you entered is invalid.
 validation.error.mustBeNumber=Value must be number.
 validation.error.mustBeBetweenZeroAndOne=Value must be between 0 and 1.
+validation.error.mustBeGreaterThanZero=Value must be greater than 0
 checkForUpdatesTask.title=Checking for CodeGPT update...
 checkForUpdatesTask.notification.message=An update for CodeGPT is available.
 checkForUpdatesTask.notification.installButton=Install update
@@ -157,5 +165,8 @@ llamaServerAgent.buildingProject.description=Building llama.cpp...
 llamaServerAgent.serverBootup.description=Booting up server...
 notification.compilationError.description=CodeGPT has detected a compilation error. Would you like assistance in resolving it?
 notification.compilationError.okLabel=Resolve errors
+notification.completionError.description=Completion failed:
%s
+statusBar.widget.tooltip=Status
 shared.promptTemplate=Prompt template:
-shared.port=Port:
\ No newline at end of file
+shared.port=Port:
+codeCompletion.progress.title=Code completion in progress
\ No newline at end of file
diff --git a/src/test/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.java b/src/test/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.java
new file mode 100644
index 00000000..87c46ac7
--- /dev/null
+++ b/src/test/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.java
@@ -0,0 +1,138 @@
+package ee.carlrobert.codegpt.codecompletions;
+
+import static ee.carlrobert.codegpt.CodeGPTKeys.MULTI_LINE_INLAY;
+import static ee.carlrobert.codegpt.CodeGPTKeys.SINGLE_LINE_INLAY;
+import static ee.carlrobert.codegpt.codecompletions.CodeCompletionService.APPLY_INLAY_ACTION_ID;
+import static ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent;
+import static ee.carlrobert.llm.client.util.JSONUtil.e;
+import static ee.carlrobert.llm.client.util.JSONUtil.jsonMapResponse;
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+
+import com.intellij.openapi.actionSystem.ActionManager;
+import com.intellij.openapi.actionSystem.AnAction;
+import com.intellij.openapi.application.ReadAction;
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.editor.Editor;
+import com.intellij.openapi.editor.EditorCustomElementRenderer;
+import com.intellij.openapi.editor.Inlay;
+import com.intellij.openapi.editor.VisualPosition;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.PsiFile;
+import ee.carlrobert.llm.client.http.exchange.StreamHttpExchange;
+import ee.carlrobert.llm.completion.CompletionEventListener;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
+import testsupport.IntegrationTest;
+
+public class CodeCompletionServiceTest extends IntegrationTest {
+
+  private final VisualPosition cursorPosition = new VisualPosition(2, 8);
+
+  public void testFetchCodeCompletionLlama() {
+    useLlamaService();
+    var codeCompletionService = CodeCompletionService.getInstance(getProject());
+    String fileContents = getResourceContent(
+        "/codecompletions/code-completion-file.txt");
+    PsiFile psiFile = myFixture.configureByText("CompletionTest.java", fileContents);
+    Editor editor = myFixture.getEditor();
+    Document document = editor.getDocument();
+    editor.getCaretModel().moveToVisualPosition(cursorPosition);
+    var prefix = "public static int gcd(int x, int y){\n";
+    var suffix = "\n"
+        + "    }";
+    var expectedCompletion = "return xyz;";
+    expectLlama((StreamHttpExchange) request -> {
+      assertThat(request.getUri().getPath()).isEqualTo("/completion");
+      assertThat(request.getMethod()).isEqualTo("POST");
+      assertThat(request.getBody())
+          .extracting("prompt")
+          .isEqualTo(InfillPromptTemplate.LLAMA.buildPrompt(prefix, suffix));
+      return List.of(jsonMapResponse(e("content", expectedCompletion), e("stop", true)));
+    });
+
+    int caretOffset = editor.getCaretModel().getOffset();
+    PsiElement elementAtCaret = ReadAction.compute(() -> psiFile.findElementAt(caretOffset));
+
+    StringBuilder actualCompletion = new StringBuilder();
+    codeCompletionService.fetchCodeCompletion(elementAtCaret, caretOffset, document,
+        new CompletionEventListener() {
+          @Override
+          public void onComplete(StringBuilder messageBuilder) {
+            actualCompletion.append(messageBuilder);
+          }
+        });
+    await().atMost(2, SECONDS)
+        .until(() -> actualCompletion.length() > 0);
+    assertEquals(expectedCompletion, actualCompletion.toString());
+  }
+
+  public void testAddInlaysSingleLine() {
+    var codeCompletionService = setupTestCodeCompletion();
+    Editor editor = myFixture.getEditor();
+    editor.getCaretModel().moveToVisualPosition(cursorPosition);
+    var expectedInlay = "        return xyz;";
+    int caretOffset = editor.getCaretModel().getOffset();
+
+    codeCompletionService.addInlays(editor, caretOffset, expectedInlay);
+
+    checkInlay(editor.getUserData(SINGLE_LINE_INLAY), InlayInlineElementRenderer.class,
+        expectedInlay, caretOffset);
+    checkPerformInlayAction(editor.getDocument(), cursorPosition.line, cursorPosition.line,
+        expectedInlay);
+    ActionManager.getInstance().unregisterAction(APPLY_INLAY_ACTION_ID);
+  }
+
+  public void testAddInlaysMultiLine() {
+    var codeCompletionService = setupTestCodeCompletion();
+    Editor editor = myFixture.getEditor();
+    editor.getCaretModel().moveToVisualPosition(cursorPosition);
+    var expectedInlay = "        int z = 1;\n        z = 2 + 3;\n        return xyz;";
+    int caretOffset = editor.getCaretModel().getOffset();
+
+    codeCompletionService.addInlays(editor, caretOffset, expectedInlay);
+
+    // First line of inlay
+    checkInlay(editor.getUserData(SINGLE_LINE_INLAY), InlayInlineElementRenderer.class,
+        expectedInlay.substring(0, expectedInlay.indexOf("\n")), caretOffset);
+    // Other lines of inlay
+    checkInlay(editor.getUserData(MULTI_LINE_INLAY), InlayBlockElementRenderer.class,
+        expectedInlay.substring(expectedInlay.indexOf("\n") + 1), caretOffset);
+    checkPerformInlayAction(editor.getDocument(), cursorPosition.line, cursorPosition.line + 2,
+        expectedInlay);
+    ActionManager.getInstance().unregisterAction(APPLY_INLAY_ACTION_ID);
+  }
+
+  private CodeCompletionService setupTestCodeCompletion() {
+    useLlamaService();
+    var codeCompletionService = CodeCompletionService.getInstance(getProject());
+    String fileContents = getResourceContent(
+        "/codecompletions/code-completion-file.txt");
+    myFixture.configureByText("CompletionTest.java", fileContents);
+    return codeCompletionService;
+  }
+
+  private void checkInlay(Inlay inlay,
+      Class clazz, String expectedText, int expectedOffset) {
+    assertNotNull(inlay);
+    assertTrue(clazz.isInstance(inlay.getRenderer()));
+    InlayElementRenderer renderer = (InlayElementRenderer) inlay.getRenderer();
+    assertEquals(expectedText, renderer.getInlayText());
+    assertEquals(expectedOffset, inlay.getOffset());
+  }
+
+  private void checkPerformInlayAction(Document document, int startLine, int endLine,
+      String expectedText) {
+    AnAction applyInlayAction = ActionManager.getInstance().getAction(APPLY_INLAY_ACTION_ID);
+    assertNotNull(applyInlayAction);
+    myFixture.performEditorAction(APPLY_INLAY_ACTION_ID);
+
+    TextRange inlayTextRange = new TextRange(document.getLineStartOffset(startLine),
+        document.getLineEndOffset(endLine));
+    assertEquals(expectedText, document.getText(inlayTextRange));
+  }
+
+
+}
diff --git a/src/test/resources/codecompletions/code-completion-file.txt b/src/test/resources/codecompletions/code-completion-file.txt
new file mode 100644
index 00000000..801cb3c1
--- /dev/null
+++ b/src/test/resources/codecompletions/code-completion-file.txt
@@ -0,0 +1,5 @@
+public class CompletionTest {
+    public static int gcd(int x, int y){
+
+    }
+}
\ No newline at end of file