Skip to content

Commit f84a659

Browse files
carlrobertohcarlrobertoh
committed
feat: support multi-line code completions
1 parent 6fef9c2 commit f84a659

File tree

23 files changed

+650
-347
lines changed

23 files changed

+650
-347
lines changed

‎codegpt-treesitter/build.gradle.kts‎

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,33 +8,33 @@ dependencies {
88
implementation("io.github.bonede:tree-sitter-elixir:0.2.0")
99
implementation("io.github.bonede:tree-sitter-dockerfile:0.2.0")
1010
implementation("io.github.bonede:tree-sitter-dart:master-a")
11-
implementation("io.github.bonede:tree-sitter-css:0.21.0")
12-
implementation("io.github.bonede:tree-sitter-cpp:0.22.0a")
13-
implementation("io.github.bonede:tree-sitter-c-sharp:0.20.0a")
11+
implementation("io.github.bonede:tree-sitter-css:0.23.1")
12+
implementation("io.github.bonede:tree-sitter-cpp:0.23.4")
13+
implementation("io.github.bonede:tree-sitter-c-sharp:0.23.1")
1414
implementation("io.github.bonede:tree-sitter-fortran:master-a")
1515
implementation("io.github.bonede:tree-sitter-gitattributes:0.1.6")
16-
implementation("io.github.bonede:tree-sitter-go:0.21.0a")
16+
implementation("io.github.bonede:tree-sitter-go:0.23.3")
1717
implementation("io.github.bonede:tree-sitter-graphql:master-a")
18-
implementation("io.github.bonede:tree-sitter-html:0.20.3")
19-
implementation("io.github.bonede:tree-sitter-javascript:0.21.2")
20-
implementation("io.github.bonede:tree-sitter-json:0.21.0a")
21-
implementation("io.github.bonede:tree-sitter-kotlin:0.3.6")
18+
implementation("io.github.bonede:tree-sitter-html:0.23.2")
19+
implementation("io.github.bonede:tree-sitter-javascript:0.23.1")
20+
implementation("io.github.bonede:tree-sitter-json:0.23.0")
21+
implementation("io.github.bonede:tree-sitter-kotlin:0.3.8.1")
2222
implementation("io.github.bonede:tree-sitter-latex:0.3.0a")
2323
implementation("io.github.bonede:tree-sitter-lua:2.1.3a")
2424
implementation("io.github.bonede:tree-sitter-m68k:0.2.7a")
2525
implementation("io.github.bonede:tree-sitter-markdown:0.7.1a")
2626
implementation("io.github.bonede:tree-sitter-objc:main-a")
2727
implementation("io.github.bonede:tree-sitter-perl:1.1.0")
28-
implementation("io.github.bonede:tree-sitter-ruby:0.21.0")
29-
implementation("io.github.bonede:tree-sitter-rust:0.21.2")
30-
implementation("io.github.bonede:tree-sitter-scala:0.21.0a")
28+
implementation("io.github.bonede:tree-sitter-ruby:0.23.1")
29+
implementation("io.github.bonede:tree-sitter-rust:0.23.1")
30+
implementation("io.github.bonede:tree-sitter-scala:0.23.3")
3131
implementation("io.github.bonede:tree-sitter-scss:1.0.0a")
3232
implementation("io.github.bonede:tree-sitter-svelte:0.11.0a")
3333
implementation("io.github.bonede:tree-sitter-swift:0.5.0")
3434
implementation("io.github.bonede:tree-sitter-yaml:0.5.0a")
35-
implementation("io.github.bonede:tree-sitter-java:0.21.0a")
36-
implementation("io.github.bonede:tree-sitter-python:0.21.0a")
37-
implementation("io.github.bonede:tree-sitter-php:0.22.4")
35+
implementation("io.github.bonede:tree-sitter-java:0.23.4")
36+
implementation("io.github.bonede:tree-sitter-python:0.23.4")
37+
implementation("io.github.bonede:tree-sitter-php:0.23.11")
3838
implementation("io.github.bonede:tree-sitter-typescript:0.21.1")
3939
implementation("io.github.bonede:tree-sitter-query:0.3.0")
4040
}

‎codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java‎

Lines changed: 82 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,44 +1,116 @@
11
package ee.carlrobert.codegpt.treesitter;
22

3+
import java.nio.charset.StandardCharsets;
4+
import org.treesitter.TSInputEdit;
35
import org.treesitter.TSLanguage;
46
import org.treesitter.TSParser;
7+
import org.treesitter.TSPoint;
58
import org.treesitter.TSTree;
69

710
public class CodeCompletionParser {
811

9-
protected final TSLanguage language;
12+
private final TSParser parser;
1013

1114
public CodeCompletionParser(TSLanguage language) {
12-
this.language = language;
15+
parser = new TSParser();
16+
parser.setLanguage(language);
1317
}
1418

1519
public String parse(String prefix, String suffix, String output) {
1620
var result = new StringBuilder(output);
21+
String input = prefix + result + suffix;
22+
TSTree currentTree = parser.parseString(null, input);
23+
1724
while (!result.isEmpty()) {
18-
if (containsError(prefix + result + suffix)) {
25+
if (containsError(currentTree)) {
26+
int deletionIndex = prefix.length() + result.length() - 1;
27+
Position pos = getPosition(input, deletionIndex);
28+
29+
int startByte = pos.byteOffset;
30+
int oldEndByte = startByte + getByteLength(result.substring(result.length() - 1));
31+
32+
TSPoint startPoint = pos.point;
33+
TSPoint oldEndPoint = computeOldEndPoint(startPoint, result.charAt(result.length() - 1));
34+
35+
currentTree.edit(
36+
new TSInputEdit(startByte, oldEndByte, startByte, startPoint, oldEndPoint, startPoint));
37+
1938
result.deleteCharAt(result.length() - 1);
39+
40+
if (result.length() > 1 && result.charAt(result.length() - 1) == '{') {
41+
long bracketCount = result.chars().filter(ch -> ch == '{').count();
42+
if (bracketCount == 1) {
43+
var newTree = parser.parseString(currentTree, prefix + result + "}" + suffix);
44+
var treeString = newTree.getRootNode().toString();
45+
if (!treeString.contains("ERROR")) {
46+
return result + "}";
47+
}
48+
}
49+
}
50+
51+
input = prefix + result + suffix;
52+
53+
currentTree = parser.parseString(currentTree, input);
2054
} else {
2155
return result.toString();
2256
}
2357
}
2458

2559
if (output.contains("\n")) {
26-
return parse(prefix, suffix, output.substring(0, output.indexOf("\n")));
60+
var finalResult = output.substring(0, output.indexOf("\n"));
61+
if (finalResult.charAt(finalResult.length() - 1) == '{') {
62+
return finalResult + "}";
63+
}
64+
return finalResult;
2765
}
2866

2967
return output;
3068
}
3169

32-
private boolean containsError(String input) {
33-
var treeString = getTree(input).getRootNode().toString();
70+
private boolean containsError(TSTree tree) {
71+
var treeString = tree.getRootNode().toString();
3472
return treeString.contains("ERROR")
3573
|| treeString.contains("MISSING \"}\"")
3674
|| treeString.contains("MISSING \")\"");
3775
}
3876

39-
private TSTree getTree(String input) {
40-
var parser = new TSParser();
41-
parser.setLanguage(language);
42-
return parser.parseString(null, input);
77+
private Position getPosition(String input, int index) {
78+
int row = 0;
79+
int col = 0;
80+
int byteOffset = 0;
81+
for (int i = 0; i < index; i++) {
82+
char c = input.charAt(i);
83+
int charByteLength = getByteLength(String.valueOf(c));
84+
byteOffset += charByteLength;
85+
86+
if (c == '\n') {
87+
row++;
88+
col = 0;
89+
} else {
90+
col++;
91+
}
92+
}
93+
return new Position(new TSPoint(row, col), byteOffset);
94+
}
95+
96+
private int getByteLength(String str) {
97+
return str.getBytes(StandardCharsets.UTF_8).length;
98+
}
99+
100+
private TSPoint computeOldEndPoint(TSPoint startPoint, char deletedChar) {
101+
int row = startPoint.getRow();
102+
int col = startPoint.getColumn();
103+
104+
if (deletedChar == '\n') {
105+
row++;
106+
col = 0;
107+
} else {
108+
col++;
109+
}
110+
return new TSPoint(row, col);
111+
}
112+
113+
private record Position(TSPoint point, int byteOffset) {
114+
43115
}
44116
}

‎codegpt-treesitter/src/test/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParserTest.java‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -38,9 +38,7 @@ class Main {
3838
return 10;
3939
}
4040
}""";
41-
var output = """
42-
prevNumber) {
43-
if() {""";
41+
var output = "prevNumber);";
4442

4543
var parsedResponse = CodeCompletionParserFactory
4644
.getParserForFileExtension("java")

‎gradle/libs.versions.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ jsoup = "1.17.2"
1212
jtokkit = "1.1.0"
1313
junit = "5.11.0"
1414
kotlin = "2.0.0"
15-
llm-client = "0.8.28"
15+
llm-client = "0.8.29"
1616
okio = "3.9.0"
17-
tree-sitter = "0.22.6a"
17+
tree-sitter = "0.24.4"
1818

1919
[libraries]
2020
analytics = { module = "com.rudderstack.sdk.java.analytics:analytics", version.ref = "analytics" }

‎src/main/java/ee/carlrobert/codegpt/settings/configuration/ConfigurationComponent.java‎

Lines changed: 8 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -25,11 +25,9 @@ public class ConfigurationComponent {
2525
private final JBCheckBox checkForNewScreenshotsCheckBox;
2626
private final JBCheckBox methodNameGenerationCheckBox;
2727
private final JBCheckBox autoFormattingCheckBox;
28-
private final JBCheckBox autocompletionPostProcessingCheckBox;
29-
private final JBCheckBox autocompletionContextAwareCheckBox;
30-
private final JBCheckBox autocompletionGitContextCheckBox;
3128
private final IntegerField maxTokensField;
3229
private final JBTextField temperatureField;
30+
private final CodeCompletionConfigurationForm codeCompletionForm;
3331

3432
public ConfigurationComponent(
3533
Disposable parentDisposable,
@@ -72,31 +70,21 @@ public void changedUpdate(DocumentEvent e) {
7270
autoFormattingCheckBox = new JBCheckBox(
7371
CodeGPTBundle.get("configurationConfigurable.autoFormatting.label"),
7472
configuration.getAutoFormattingEnabled());
75-
autocompletionPostProcessingCheckBox = new JBCheckBox(
76-
CodeGPTBundle.get("configurationConfigurable.autocompletionPostProcessing.label"),
77-
configuration.getAutocompletionPostProcessingEnabled()
78-
);
79-
autocompletionContextAwareCheckBox = new JBCheckBox(
80-
CodeGPTBundle.get("configurationConfigurable.autocompletionContextAwareCheckBox.label"),
81-
configuration.getAutocompletionContextAwareEnabled()
82-
);
83-
autocompletionGitContextCheckBox = new JBCheckBox(
84-
CodeGPTBundle.get("configurationConfigurable.autocompletionGitContextCheckBox.label"),
85-
configuration.getAutocompletionGitContextEnabled()
86-
);
73+
74+
codeCompletionForm = new CodeCompletionConfigurationForm();
8775

8876
mainPanel = FormBuilder.createFormBuilder()
8977
.addComponent(checkForPluginUpdatesCheckBox)
9078
.addComponent(checkForNewScreenshotsCheckBox)
9179
.addComponent(methodNameGenerationCheckBox)
9280
.addComponent(autoFormattingCheckBox)
93-
.addComponent(autocompletionPostProcessingCheckBox)
94-
.addComponent(autocompletionContextAwareCheckBox)
95-
.addComponent(autocompletionGitContextCheckBox)
9681
.addVerticalGap(4)
9782
.addComponent(new TitledSeparator(
9883
CodeGPTBundle.get("configurationConfigurable.section.assistant.title")))
9984
.addComponent(createAssistantConfigurationForm())
85+
.addComponent(new TitledSeparator(
86+
CodeGPTBundle.get("configurationConfigurable.section.codeCompletion.title")))
87+
.addComponent(codeCompletionForm.createPanel())
10088
.addComponentFillVertically(new JPanel(), 0)
10189
.getPanel();
10290
}
@@ -113,9 +101,7 @@ public ConfigurationSettingsState getCurrentFormState() {
113101
state.setCheckForNewScreenshots(checkForNewScreenshotsCheckBox.isSelected());
114102
state.setMethodNameGenerationEnabled(methodNameGenerationCheckBox.isSelected());
115103
state.setAutoFormattingEnabled(autoFormattingCheckBox.isSelected());
116-
state.setAutocompletionPostProcessingEnabled(autocompletionPostProcessingCheckBox.isSelected());
117-
state.setAutocompletionContextAwareEnabled(autocompletionContextAwareCheckBox.isSelected());
118-
state.setAutocompletionGitContextEnabled(autocompletionGitContextCheckBox.isSelected());
104+
state.setCodeCompletionSettings(codeCompletionForm.getFormState());
119105
return state;
120106
}
121107

@@ -127,13 +113,7 @@ public void resetForm() {
127113
checkForNewScreenshotsCheckBox.setSelected(configuration.getCheckForNewScreenshots());
128114
methodNameGenerationCheckBox.setSelected(configuration.getMethodNameGenerationEnabled());
129115
autoFormattingCheckBox.setSelected(configuration.getAutoFormattingEnabled());
130-
autocompletionPostProcessingCheckBox.setSelected(
131-
configuration.getAutocompletionPostProcessingEnabled());
132-
autocompletionContextAwareCheckBox.setSelected(
133-
configuration.getAutocompletionContextAwareEnabled());
134-
autocompletionGitContextCheckBox.setSelected(
135-
configuration.getAutocompletionGitContextEnabled()
136-
);
116+
codeCompletionForm.resetForm(configuration.getCodeCompletionSettings());
137117
}
138118

139119
// Formatted keys are not referenced in the messages bundle file

‎src/main/java/ee/carlrobert/codegpt/settings/service/llama/form/InfillPromptTemplatePanel.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package ee.carlrobert.codegpt.settings.service.llama.form;
22

3+
import ee.carlrobert.codegpt.codecompletions.CompletionType;
34
import ee.carlrobert.codegpt.codecompletions.InfillPromptTemplate;
45
import ee.carlrobert.codegpt.codecompletions.InfillRequest;
56

@@ -17,6 +18,8 @@ public InfillPromptTemplatePanel(
1718

1819
@Override
1920
protected String buildPromptDescription(InfillPromptTemplate template) {
20-
return template.buildPrompt(new InfillRequest.Builder("PREFIX", "SUFFIX", 0).build());
21+
return template.buildPrompt(new InfillRequest
22+
.Builder("PREFIX", "SUFFIX", 0, CompletionType.MULTI_LINE)
23+
.build());
2124
}
2225
}

0 commit comments

Comments
 (0)