diff --git a/src/main/java/ee/carlrobert/codegpt/EncodingManager.java b/src/main/java/ee/carlrobert/codegpt/EncodingManager.java index 44aef5fb..c3b4d44f 100644 --- a/src/main/java/ee/carlrobert/codegpt/EncodingManager.java +++ b/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 = diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt index e9fbd23b..dbdd2860 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt +++ b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillPromptTemplate.kt @@ -72,7 +72,7 @@ enum class InfillPromptTemplate(val label: String, val stopTokens: List? } } }, - STAR_CODER("StarCoder2", listOf("<|endoftext|>")) { + STAR_CODER("StarCoder2", listOf("<|endoftext|>", "")) { override fun buildPrompt(infillDetails: InfillRequest): String { // see https://huggingface.co/spaces/bigcode/bigcode-playground/blob/main/app.py val infillPrompt = diff --git a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt b/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt index 4225374e..db70cbc1 100644 --- a/src/main/kotlin/ee/carlrobert/codegpt/codecompletions/InfillRequestUtil.kt +++ b/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().getRepositoryForFile(project.workspaceFile) + val repositoryManager = project.service() + val gitRepository = repositoryManager.getRepositoryForFile(project.workspaceFile) + ?: repositoryManager.repositories.firstOrNull() if (service().state.autocompletionGitContextEnabled && gitRepository != null) { try { val stagedDiff = GitUtil.getStagedDiff(project, gitRepository) diff --git a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt index 53d6c761..3a55fb87 100644 --- a/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt +++ b/src/test/kotlin/ee/carlrobert/codegpt/codecompletions/CodeCompletionServiceTest.kt @@ -4,17 +4,156 @@ import com.intellij.codeInsight.inline.completion.session.InlineCompletionSessio import com.intellij.openapi.editor.VisualPosition import com.intellij.openapi.util.TextRange import com.intellij.testFramework.PlatformTestUtil +import ee.carlrobert.codegpt.CodeGPTKeys.REMAINING_EDITOR_COMPLETION 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() { - fun testApplyCompletionNextWordInlay() { + fun `test code completion with CodeGPT provider`() { + useCodeGPTService() + 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 + 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", prefix, suffix, "java") + 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 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 + ) + } + + fun `test apply next partial completion word`() { useLlamaService(true) myFixture.configureByText( "CompletionTest.java", diff --git a/src/test/kotlin/testsupport/mixin/ShortcutsTestMixin.kt b/src/test/kotlin/testsupport/mixin/ShortcutsTestMixin.kt index abe12273..dbabc145 100644 --- a/src/test/kotlin/testsupport/mixin/ShortcutsTestMixin.kt +++ b/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().state.selectedService = ServiceType.CODEGPT setCredential(CODEGPT_API_KEY, "TEST_API_KEY") - service().state.chatCompletionSettings.model = "TEST_MODEL" + service().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().state.selectedService = ServiceType.OPENAI setCredential(OPENAI_API_KEY, "TEST_API_KEY") - OpenAISettings.getCurrentState().model = model + service().state.run { + model = chatModel + isCodeCompletionsEnabled = true + } } fun useAzureService() {