Revert "Revert "feat: code completion improvements""

This reverts commit 7f586da0c1.

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

@@ -9,6 +9,7 @@ import com.knuddels.jtokkit.api.EncodingRegistry;
 import com.knuddels.jtokkit.api.EncodingType;
 import ee.carlrobert.codegpt.conversations.Conversation;
 import ee.carlrobert.llm.client.openai.completion.request.OpenAIChatCompletionMessage;
+import java.util.List;
 
 @Service
 public final class EncodingManager {
@@ -52,4 +53,19 @@ public final class EncodingManager {
       return 0;
     }
   }
+
+  /**
+   * Truncates the given text to the given number of tokens.
+   *
+   * @param text The text to truncate.
+   * @param maxTokens The maximum number of tokens to keep.
+   * @param fromStart Whether to truncate from the start or the end of the text.
+   * @return The truncated text.
+   */
+  public String truncateText(String text, int maxTokens, boolean fromStart) {
+    List<Integer> tokens = encoding.encode(text);
+    int tokensToRetrieve = Math.min(maxTokens, tokens.size());
+    int startIndex = fromStart ? 0 : tokens.size() - tokensToRetrieve;
+    return encoding.decode(tokens.subList(startIndex, startIndex + tokensToRetrieve));
+  }
 }

src/main/java/ee/carlrobert/codegpt/actions/DisableCompletionsAction.java

@@ -1,8 +1,8 @@
 package ee.carlrobert.codegpt.actions;
 
-import com.intellij.openapi.actionSystem.AnAction;
 import com.intellij.openapi.actionSystem.AnActionEvent;
 import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.project.DumbAwareAction;
 import ee.carlrobert.codegpt.codecompletions.CodeGPTEditorManager;
 import ee.carlrobert.codegpt.settings.configuration.ConfigurationState;
 import org.jetbrains.annotations.NotNull;
@@ -10,7 +10,7 @@ import org.jetbrains.annotations.NotNull;
 /**
  * Disables code-completion.<br/> Publishes message to {@link CodeCompletionEnabledListener#TOPIC}
  */
-public class DisableCompletionsAction extends AnAction {
+public class DisableCompletionsAction extends DumbAwareAction {
 
   @Override
   public void actionPerformed(@NotNull AnActionEvent e) {

src/main/java/ee/carlrobert/codegpt/actions/EnableCompletionsAction.java

@@ -1,15 +1,15 @@
 package ee.carlrobert.codegpt.actions;
 
-import com.intellij.openapi.actionSystem.AnAction;
 import com.intellij.openapi.actionSystem.AnActionEvent;
 import com.intellij.openapi.application.ApplicationManager;
+import com.intellij.openapi.project.DumbAwareAction;
 import ee.carlrobert.codegpt.settings.configuration.ConfigurationState;
 import org.jetbrains.annotations.NotNull;
 
 /**
  * Enables code-completion.<br/> Publishes message to {@link CodeCompletionEnabledListener#TOPIC}
  */
-public class EnableCompletionsAction extends AnAction {
+public class EnableCompletionsAction extends DumbAwareAction {
 
   @Override
   public void actionPerformed(@NotNull AnActionEvent e) {

src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java

@@ -40,8 +40,7 @@ class CodeCompletionEventListener implements CompletionEventListener {
       progressIndicator.processFinish();
     }
 
-    var editorManager = CodeGPTEditorManager.getInstance();
-    editorManager.disposeEditorInlays(editor);
+    CodeGPTEditorManager.getInstance().disposeEditorInlays(editor);
 
     var inlayText = messageBuilder.toString();
     if (!inlayText.isEmpty()) {
@@ -60,7 +59,7 @@ class CodeCompletionEventListener implements CompletionEventListener {
     Notifications.Bus.notify(OverlayUtil.getDefaultNotification(
             String.format(
                 CodeGPTBundle.get("notification.completionError.description"),
-                ex.getMessage()),
+                error.getMessage()),
             NotificationType.ERROR)
         .addAction(new OpenSettingsAction()), editor.getProject());
   }

src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java

@@ -6,6 +6,7 @@ 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.codeInsight.lookup.LookupManager;
 import com.intellij.openapi.Disposable;
 import com.intellij.openapi.actionSystem.ActionManager;
 import com.intellij.openapi.actionSystem.AnAction;
@@ -24,13 +25,9 @@ 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;
@@ -38,14 +35,12 @@ 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)
@@ -53,7 +48,6 @@ import org.jetbrains.annotations.NotNull;
 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);
 
@@ -77,16 +71,13 @@ public final class CodeCompletionService implements Disposable {
     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)
+        || LookupManager.getActiveLookup(editor) != null
         || editor.isViewer()
         || editor.isOneLineMode()
     ) {
@@ -99,54 +90,16 @@ public final class CodeCompletionService implements Disposable {
       return;
     }
 
-    PsiElement elementAtCaret = ReadAction.compute(() -> psiFile.findElementAt(offset));
-    var completionService = CodeCompletionService.getInstance(project);
-    if (!completionService.isCompletionAllowed(elementAtCaret)) {
-      return;
-    }
-
+    var request = InfillRequestDetails.fromDocumentWithMaxOffset(document, offset);
     callDebouncer.debounce(
         Void.class,
-        (progressIndicator) -> completionService.fetchCodeCompletion(
-            elementAtCaret,
-            offset,
-            document,
+        (progressIndicator) -> CompletionRequestService.getInstance().getCodeCompletionAsync(
+            request,
             new CodeCompletionEventListener(editor, offset, progressIndicator)),
         500,
         TimeUnit.MILLISECONDS);
   }
 
-  /**
-   * Fetches code-completion (FIM) for the given position ({@code offsetInFile}) in the file. <br/>
-   * 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<String> linesList = inlayText.lines().collect(toList());
@@ -201,11 +154,13 @@ public final class CodeCompletionService implements Disposable {
     Document document = editor.getDocument();
     document.insertString(offset, text);
     editor.getCaretModel().moveToOffset(offset + text.length());
-    EditorUtil.reformatDocument(
-        requireNonNull(editor.getProject()),
-        document,
-        offset,
-        offset + text.length());
+    if (ConfigurationState.getInstance().isAutoFormattingEnabled()) {
+      EditorUtil.reformatDocument(
+          requireNonNull(editor.getProject()),
+          document,
+          offset,
+          offset + text.length());
+    }
   }
 
   @RequiresReadLock
@@ -246,20 +201,4 @@ public final class CodeCompletionService implements Disposable {
         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)));
-  }
 }

src/main/java/ee/carlrobert/codegpt/codecompletions/InfillRequestDetails.java

@@ -1,7 +1,14 @@
 package ee.carlrobert.codegpt.codecompletions;
 
+import com.intellij.openapi.editor.Document;
+import com.intellij.openapi.util.TextRange;
+import ee.carlrobert.codegpt.EncodingManager;
+
 public class InfillRequestDetails {
 
+  private static final int MAX_OFFSET = 10_000;
+  private static final int MAX_PROMPT_TOKENS = 512;
+
   private final String prefix;
   private final String suffix;
 
@@ -10,6 +17,12 @@ public class InfillRequestDetails {
     this.suffix = suffix;
   }
 
+  public static InfillRequestDetails fromDocumentWithMaxOffset(Document document, int caretOffset) {
+    int start = Math.max(0, caretOffset - MAX_OFFSET);
+    int end = Math.min(document.getTextLength(), caretOffset + MAX_OFFSET);
+    return fromDocumentWithCustomRange(document, caretOffset, start, end);
+  }
+
   public String getPrefix() {
     return prefix;
   }
@@ -17,4 +30,21 @@ public class InfillRequestDetails {
   public String getSuffix() {
     return suffix;
   }
+
+  private static InfillRequestDetails fromDocumentWithCustomRange(
+      Document document,
+      int caretOffset,
+      int start,
+      int end) {
+    var prefix = truncateText(document, start, caretOffset, false);
+    var suffix = truncateText(document, caretOffset, end, true);
+    return new InfillRequestDetails(prefix, suffix);
+  }
+
+  private static String truncateText(Document document, int start, int end, boolean fromStart) {
+    return EncodingManager.getInstance().truncateText(
+        document.getText(new TextRange(start, end)),
+        MAX_PROMPT_TOKENS,
+        fromStart);
+  }
 }