test: add test coverage for code completions

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

@@ -20,6 +20,9 @@ import java.util.stream.Stream;
 @Service
 public final class EncodingManager {
 
+  private static final String SPECIAL_START = "<|";
+  private static final String SPECIAL_END = "|>";
+
   private static final Logger LOG = Logger.getInstance(EncodingManager.class);
 
   private final EncodingRegistry registry = Encodings.newDefaultEncodingRegistry();
@@ -76,7 +79,8 @@ public final class EncodingManager {
    * @return The truncated text.
    */
   public String truncateText(String text, int maxTokens, boolean fromStart) {
-    var tokens = encoding.encode(text);
+    var textWithSpecialEncodingsRemoved = text.replace(SPECIAL_START, "").replace(SPECIAL_END, "");
+    var tokens = encoding.encode(textWithSpecialEncodingsRemoved);
     int tokensToRetrieve = Math.min(maxTokens, tokens.size());
     int startIndex = fromStart ? 0 : tokens.size() - tokensToRetrieve;
     var truncatedList =

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt

@@ -72,7 +72,7 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List<String>?
             }
         }
     },
-    STAR_CODER("StarCoder2", listOf("<|endoftext|>")) {
+    STAR_CODER("StarCoder2", listOf("<|endoftext|>", "<file_sep>")) {
         override fun buildPrompt(infillDetails: InfillRequest): String {
             // see https://huggingface.co/spaces/bigcode/bigcode-playground/blob/main/app.py
             val infillPrompt =

src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt

@@ -27,8 +27,9 @@ object InfillRequestUtil {
             )
         val project = request.editor.project ?: return infillRequestBuilder.build()
 
-        val gitRepository =
-            project.service<GitRepositoryManager>().getRepositoryForFile(project.workspaceFile)
+        val repositoryManager = project.service<GitRepositoryManager>()
+        val gitRepository = repositoryManager.getRepositoryForFile(project.workspaceFile)
+            ?: repositoryManager.repositories.firstOrNull()
         if (service<ConfigurationSettings>().state.autocompletionGitContextEnabled && gitRepository != null) {
             try {
                 val stagedDiff = GitUtil.getStagedDiff(project, gitRepository)

src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt

@@ -1,59 +1,165 @@
 package ee.carlrobert.codegpt.codecompletions
 
+import com.intellij.codeInsight.inline.completion.session.InlineCompletionSession.Companion.getOrNull
 import com.intellij.openapi.editor.VisualPosition
+import com.intellij.testFramework.PlatformTestUtil
 import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION
-import ee.carlrobert.codegpt.settings.service.llama.LlamaSettings
-import ee.carlrobert.codegpt.util.file.FileUtil.getResourceContent
+import ee.carlrobert.codegpt.util.file.FileUtil
 import ee.carlrobert.llm.client.http.RequestEntity
 import ee.carlrobert.llm.client.http.exchange.StreamHttpExchange
-import ee.carlrobert.llm.client.util.JSONUtil.e
-import ee.carlrobert.llm.client.util.JSONUtil.jsonMapResponse
+import ee.carlrobert.llm.client.util.JSONUtil.*
 import org.assertj.core.api.Assertions.assertThat
 import testsupport.IntegrationTest
 
 class CodeCompletionServiceTest : IntegrationTest() {
 
-    private val cursorPosition = VisualPosition(3, 0)
-
-    fun testFetchCodeCompletionLlama() {
-        useLlamaService()
-        LlamaSettings.getCurrentState().isCodeCompletionsEnabled = true
+    fun `test code completion with CodeGPT provider`() {
+        useCodeGPTService()
         myFixture.configureByText(
             "CompletionTest.java",
-            getResourceContent("/codecompletions/code-completion-file.txt")
+            FileUtil.getResourceContent("/codecompletions/code-completion-file.txt")
         )
-        myFixture.editor.caretModel.moveToVisualPosition(cursorPosition)
-        val expectedCompletion = "TEST_OUTPUT"
+        myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(3, 0))
         val prefix = """
-         ${"z".repeat(245)}
-         [INPUT]
-         c
-         """.trimIndent() // 128 tokens
+             ${"z".repeat(245)}
+             [INPUT]
+             p
+             """.trimIndent() // 128 tokens
         val suffix = """
-         
-         [\INPUT]
-         ${"z".repeat(247)}
-         """.trimIndent() // 128 tokens
-        expectLlama(StreamHttpExchange { request: RequestEntity ->
-            assertThat(request.uri.path).isEqualTo("/completion")
+             
+             [\INPUT]
+             ${"z".repeat(247)}
+             """.trimIndent() // 128 tokens
+        expectCodeGPT(StreamHttpExchange { request: RequestEntity ->
+            assertThat(request.uri.path).isEqualTo("/v1/code/completions")
             assertThat(request.method).isEqualTo("POST")
             assertThat(request.body)
-                .extracting("prompt")
-                .isEqualTo(
-                    InfillPromptTemplate.CODE_LLAMA.buildPrompt(
-                        InfillRequest.Builder(prefix, suffix).build()
-                    )
-                )
+                .extracting("model", "prefix", "suffix", "fileExtension")
+                .containsExactly("TEST_CODE_MODEL", prefix, suffix, "java")
             listOf(
-                jsonMapResponse(
-                    e("content", expectedCompletion),
-                    e("stop", true)
-                )
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "ublic "))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "void"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", " main")))
             )
         })
 
-        myFixture.type('c')
+        myFixture.type('p')
 
-        waitExpecting { "TEST_OUTPUT" == REMAINING_EDITOR_COMPLETION[myFixture.editor] }
+        assertInlineSuggestion("Failed to display initial inline suggestion.") {
+            "ublic void main" == it
+        }
     }
-}
+
+    fun `test code completion with OpenAI provider`() {
+        useOpenAIService()
+        myFixture.configureByText(
+            "CompletionTest.java",
+            FileUtil.getResourceContent("/codecompletions/code-completion-file.txt")
+        )
+        myFixture.editor.caretModel.moveToVisualPosition(VisualPosition(3, 0))
+        val prefix = """
+             ${"z".repeat(245)}
+             [INPUT]
+             p
+             """.trimIndent() // 128 tokens
+        val suffix = """
+             
+             [\INPUT]
+             ${"z".repeat(247)}
+             """.trimIndent() // 128 tokens
+        expectOpenAI(StreamHttpExchange { request: RequestEntity ->
+            assertThat(request.uri.path).isEqualTo("/v1/completions")
+            assertThat(request.method).isEqualTo("POST")
+            assertThat(request.body)
+                .extracting("model", "prompt", "suffix", "max_tokens")
+                .containsExactly("gpt-3.5-turbo-instruct", prefix, suffix, 128)
+            listOf(
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "ublic "))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "void"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", " main")))
+            )
+        })
+
+        myFixture.type('p')
+
+        assertInlineSuggestion("Failed to display initial inline suggestion.") {
+            "ublic void main" == it
+        }
+    }
+
+    fun `test fetching next code completion`() {
+        useCodeGPTService()
+        myFixture.configureByText("CompletionTest.java", "")
+        expectCodeGPT(StreamHttpExchange { request: RequestEntity ->
+            assertThat(request.uri.path).isEqualTo("/v1/code/completions")
+            assertThat(request.method).isEqualTo("POST")
+            assertThat(request.body)
+                .extracting("model", "prefix", "suffix", "fileExtension")
+                .containsExactly("TEST_CODE_MODEL", "p", "", "java")
+            listOf(
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "ublic static"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", " void main(String"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "[] args) {\n"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "    System.out.print"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "ln(\"Hello, Worl"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "d!\");\n"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "}\n"))),
+                jsonMapResponse(
+                    "choices",
+                    jsonArray(jsonMap("text", "private static int getX() {\n"))
+                )
+            )
+        })
+        myFixture.type('p')
+        assertInlineSuggestion("Failed to display initial inline suggestion.") {
+            "ublic static void main(String[] args) {\n    System.out.println(\"Hello, World!\");" == it
+        }
+        assertThat(REMAINING_EDITOR_COMPLETION.get(myFixture.editor)).isEqualTo(
+            "ublic static void main(String[] args) {\n" +
+                    "    System.out.println(\"Hello, World!\");\n" +
+                    "}\n" +
+                    "private static int getX() {"
+        )
+        expectCodeGPT(StreamHttpExchange { request: RequestEntity ->
+            assertThat(request.uri.path).isEqualTo("/v1/code/completions")
+            assertThat(request.method).isEqualTo("POST")
+            assertThat(request.body)
+                .extracting("model", "prefix", "suffix", "fileExtension")
+                .containsExactly(
+                    "TEST_CODE_MODEL",
+                    "public static void main(String[] args) {\n" +
+                            "    System.out.println(\"Hello, World!\");\n" +
+                            "}\n" +
+                            "private static int getX() {",
+                    "",
+                    "java"
+                )
+            listOf(
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "\n    retur"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "n 10;\n"))),
+                jsonMapResponse("choices", jsonArray(jsonMap("text", "}\n"))),
+            )
+        })
+
+        myFixture.type('\t')
+
+        PlatformTestUtil.waitWithEventsDispatching(
+            "Failed to retrieve next completion",
+            {
+                REMAINING_EDITOR_COMPLETION.get(myFixture.editor) == "\n}\nprivate static int getX() {"
+            },
+            10
+        )
+    }
+
+    private fun assertInlineSuggestion(errorMessage: String, onAssert: (String) -> Boolean) {
+        PlatformTestUtil.waitWithEventsDispatching(
+            errorMessage,
+            {
+                val session = getOrNull(myFixture.editor) ?: return@waitWithEventsDispatching false
+                onAssert(session.context.textToInsert())
+            },
+            5
+        )
+    }
+}

src/test/kotlin/testsupport/mixin/ShortcutsTestMixin.kt

@@ -19,19 +19,22 @@ import java.util.function.BooleanSupplier
 interface ShortcutsTestMixin {
 
   fun useCodeGPTService() {
-    GeneralSettings.getCurrentState().selectedService = ServiceType.CODEGPT
+    service<GeneralSettings>().state.selectedService = ServiceType.CODEGPT
     setCredential(CODEGPT_API_KEY, "TEST_API_KEY")
-    service<CodeGPTServiceSettings>().state.chatCompletionSettings.model = "TEST_MODEL"
+    service<CodeGPTServiceSettings>().state.run {
+      chatCompletionSettings.model = "TEST_MODEL"
+      codeCompletionSettings.model = "TEST_CODE_MODEL"
+      codeCompletionSettings.codeCompletionsEnabled = true
+    }
   }
 
-  fun useOpenAIService() {
-    useOpenAIService("gpt-4")
-  }
-
-  fun useOpenAIService(model: String? = "gpt-4") {
-    GeneralSettings.getCurrentState().selectedService = ServiceType.OPENAI
+  fun useOpenAIService(chatModel: String? = "gpt-4") {
+    service<GeneralSettings>().state.selectedService = ServiceType.OPENAI
     setCredential(OPENAI_API_KEY, "TEST_API_KEY")
-    OpenAISettings.getCurrentState().model = model
+    service<OpenAISettings>().state.run {
+      model = chatModel
+      isCodeCompletionsEnabled = true
+    }
   }
 
   fun useAzureService() {