diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/completion/CodeCompletionParser.java similarity index 96% rename from codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java rename to codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/completion/CodeCompletionParser.java index 6b388f2d1..9f96f550a 100644 --- a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParser.java +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/completion/CodeCompletionParser.java @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.treesitter; +package ee.carlrobert.codegpt.treesitter.completion; import org.treesitter.TSLanguage; import org.treesitter.TSNode; diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParserFactory.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/completion/CodeCompletionParserFactory.java similarity index 98% rename from codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParserFactory.java rename to codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/completion/CodeCompletionParserFactory.java index 9bbe6dd43..ee9ba32d6 100644 --- a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParserFactory.java +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/completion/CodeCompletionParserFactory.java @@ -1,4 +1,4 @@ -package ee.carlrobert.codegpt.treesitter; +package ee.carlrobert.codegpt.treesitter.completion; import org.treesitter.TSLanguage; import org.treesitter.TreeSitterCSharp; diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/ProcessedTag.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/ProcessedTag.java new file mode 100644 index 000000000..1dd643fa2 --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/ProcessedTag.java @@ -0,0 +1,20 @@ +package ee.carlrobert.codegpt.treesitter.repository; + +public class ProcessedTag { + + private final String filePath; + private final String modifiedContent; + + public ProcessedTag(String filePath, String modifiedContent) { + this.filePath = filePath; + this.modifiedContent = modifiedContent; + } + + public String getFilePath() { + return filePath; + } + + public String getModifiedContent() { + return modifiedContent; + } +} \ No newline at end of file diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/QueryUtil.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/QueryUtil.java new file mode 100644 index 000000000..a0c9240d0 --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/QueryUtil.java @@ -0,0 +1,88 @@ +package ee.carlrobert.codegpt.treesitter.repository; + +import static java.util.stream.Collectors.toSet; + +import ee.carlrobert.codegpt.treesitter.repository.parser.RepositoryParser.FileMapping; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.function.Function; +import org.treesitter.TSLanguage; +import org.treesitter.TSNode; +import org.treesitter.TSParser; +import org.treesitter.TSQuery; +import org.treesitter.TSQueryCursor; +import org.treesitter.TSQueryMatch; +import org.treesitter.TSTree; + +public class QueryUtil { + + public static Set extractTagsFromFile( + TSLanguage language, + FileMapping fileMapping, + List queries, + TagType tagType) { + var fileContent = fileMapping.getContent(); + var rootNode = QueryUtil.getTree(language, fileContent).getRootNode(); + return queries.stream() + .map(query -> new TSQuery(language, query)) + .flatMap(query -> + query(query, rootNode, node -> new Tag( + fileMapping.getPath(), + fileMapping.getContent().substring(node.getStartByte(), node.getEndByte()), + language.symbolName(node.getSymbol()), + tagType)).stream()) + .collect(toSet()); + } + + public static List query(TSLanguage language, String code, String query) { + return query( + new TSQuery(language, query), + getTree(language, code).getRootNode(), + node -> new QueryResult( + code.substring(node.getStartByte(), node.getEndByte()), + language.symbolName(node.getSymbol()))); + } + + private static TSTree getTree(TSLanguage language, String input) { + var parser = new TSParser(); + parser.setLanguage(language); + return parser.parseString(null, input); + } + + private static List query(TSQuery query, TSNode rootNode, Function onCapture) { + var cursor = new TSQueryCursor(); + cursor.exec(query, rootNode); + var matches = new ArrayList(); + var match = new TSQueryMatch(); + while (cursor.nextMatch(match)) { + for (var capture : match.getCaptures()) { + try { + matches.add(onCapture.apply(capture.getNode())); + } catch (Throwable t) { + // todo: log + } + } + } + return matches; + } + + public static class QueryResult { + + private final String name; + private final String symbolName; + + public QueryResult(String name, String symbolName) { + this.name = name; + this.symbolName = symbolName; + } + + public String getName() { + return name; + } + + public String getSymbolName() { + return symbolName; + } + } +} diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/Tag.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/Tag.java new file mode 100644 index 000000000..0c025e836 --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/Tag.java @@ -0,0 +1,71 @@ +package ee.carlrobert.codegpt.treesitter.repository; + +import com.fasterxml.jackson.annotation.JsonCreator; +import com.fasterxml.jackson.annotation.JsonProperty; +import java.nio.file.Path; +import java.util.Objects; + +public class Tag { + + private final String filePath; + private final String name; + private final String symbolName; + private final TagType type; + + @JsonCreator + public Tag( + @JsonProperty("filePath") String filePath, + @JsonProperty("name") String name, + @JsonProperty("symbolName") String symbolName, + @JsonProperty("type") TagType type) { + this.filePath = filePath; + this.name = name; + this.symbolName = symbolName; + this.type = type; + } + + public String getFilePath() { + return filePath; + } + + public String getName() { + return name; + } + + public String getSymbolName() { + return symbolName; + } + + public TagType getType() { + return type; + } + + @Override + public String toString() { + return "Tag{" + + "fileName='" + Path.of(filePath).getFileName().toString() + '"' + + "name='" + name + '"' + + ", symbolName='" + symbolName + '"' + + '}' + "\n"; + } + + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof Tag)) { + return false; + } + Tag tag = (Tag) o; + return Objects.equals(filePath, tag.filePath) + && Objects.equals(name, tag.name) + && Objects.equals(symbolName, tag.symbolName); + } + + @Override + public int hashCode() { + return Objects.hash(filePath, name, symbolName); + } +} diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/TagType.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/TagType.java new file mode 100644 index 000000000..c25ab1218 --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/TagType.java @@ -0,0 +1,5 @@ +package ee.carlrobert.codegpt.treesitter.repository; + +public enum TagType { + DEFINITION, REFERENCE +} diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/JavaRepositoryParser.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/JavaRepositoryParser.java new file mode 100644 index 000000000..0fd584626 --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/JavaRepositoryParser.java @@ -0,0 +1,129 @@ +package ee.carlrobert.codegpt.treesitter.repository.parser; + +import static java.util.stream.Collectors.toList; +import static java.util.stream.Collectors.toSet; + +import com.intellij.openapi.project.Project; +import ee.carlrobert.codegpt.treesitter.repository.ProcessedTag; +import ee.carlrobert.codegpt.treesitter.repository.Tag; +import ee.carlrobert.codegpt.treesitter.repository.TagType; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Set; +import java.util.regex.Pattern; +import javax.annotation.Nullable; +import org.treesitter.TreeSitterJava; + +public class JavaRepositoryParser extends RepositoryParser { + + JavaRepositoryParser(Project project) { + super(project, new TreeSitterJava(), "java"); + } + + @Override + protected List getDefinitionQueries() { + return List.of( + "(class_declaration name: (identifier) @name.definition.class) @definition.class", + "(method_declaration name: (identifier) @name.definition.method) @definition.method", + "(interface_declaration name: (identifier) @name.definition.interface) @definition.interface", + "(package_declaration) @definition.package"); + } + + @Override + protected List getReferenceQueries() { + return List.of( + "(method_invocation name: (identifier) @name.reference.call arguments: (argument_list) @reference.call)", + "(type_list (type_identifier) @name.reference.implementation) @reference.implementation", + "(object_creation_expression type: (type_identifier) @name.reference.class) @reference.class", + "(superclass (type_identifier) @name.reference.class) @reference.class"); + } + + @Override + protected Set processMatchedTags( + Map> matchedTags, + Set definitionTags, + String codeSnippetPath) { + var dependencies = findDependencies(codeSnippetPath, definitionTags); + return matchedTags.entrySet().stream() + .filter(entry -> { + var className = Path.of(entry.getKey()).getFileName().toString().replace(".java", ""); + return dependencies.stream().anyMatch(dep -> dep.endsWith(className)); + }) + .map(this::processMatchEntry) + .limit(10) + .collect(toSet()); + } + + private @Nullable String getPackageDeclaration(String code) { + var result = query(code, "(package_declaration) @definition.package"); + if (result.isEmpty()) { + return null; + } + return result.get(0).getName(); + } + + private Set getImports(String code) { + return query(code, "(import_declaration) @definition.import").stream() + .map(it -> extractPackageName(it.getName())) + .collect(toSet()); + } + + private Set findDependencies(String codeSnippetPath, Set definitionTags) { + var fileContent = readFileContent(Path.of(codeSnippetPath)); + var packageDeclaration = getPackageDeclaration(fileContent); + var packageLevelDependencies = definitionTags.stream() + .filter( + defTag -> "package_declaration".equals(defTag.getSymbolName()) + && packageDeclaration != null + && defTag.getName().contentEquals(packageDeclaration)) + .map(declaration -> { + var fileName = Path.of(declaration.getFilePath()).getFileName().toString(); + return String.format( + "%s.%s", + extractPackageName(declaration.getName()), + fileName.replace(".java", "")); + }) + .collect(toSet()); + packageLevelDependencies.addAll(getImports(fileContent)); + return packageLevelDependencies; + } + + private ProcessedTag processMatchEntry(Entry> entry) { + var fileContent = readFileContent(Path.of(entry.getKey())); + var tags = getTags( + List.of(new FileMapping(entry.getKey(), fileContent)), + getDefinitionQueries(), + TagType.DEFINITION); + var methodDeclarations = tags + .stream() + .filter( + tag -> "method_declaration".equals(tag.getSymbolName()) + && entry.getValue().stream().anyMatch(target -> + tag.getName().contains(target.getName()))) + .map(Tag::getName) + .collect(toList()); + return new ProcessedTag(entry.getKey(), String.join("\n\n", methodDeclarations)); + } + + private static String extractPackageName(String line) { + var regex = "(?:import|package)\\s+([\\w.]+)\\.*"; + var pattern = Pattern.compile(regex); + var matcher = pattern.matcher(line); + if (matcher.find()) { + return matcher.group(1); + } + return null; + } + + private static String readFileContent(Path path) { + try { + return Files.readString(path); + } catch (IOException e) { + return ""; + } + } +} diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/PhpRepositoryParser.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/PhpRepositoryParser.java new file mode 100644 index 000000000..787036149 --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/PhpRepositoryParser.java @@ -0,0 +1,63 @@ +package ee.carlrobert.codegpt.treesitter.repository.parser; + +import com.intellij.openapi.project.Project; +import ee.carlrobert.codegpt.treesitter.repository.ProcessedTag; +import ee.carlrobert.codegpt.treesitter.repository.Tag; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.treesitter.TreeSitterPhp; + +public class PhpRepositoryParser extends RepositoryParser { + + PhpRepositoryParser(Project project) { + super(project, new TreeSitterPhp(), "php"); + } + + @Override + protected List getDefinitionQueries() { + return List.of( + "(namespace_definition\n" + + " name: (namespace_name) @name) @module", + "(interface_declaration\n" + + " name: (name) @name) @definition.interface", + "(trait_declaration\n" + + " name: (name) @name) @definition.interface", + "(class_declaration\n" + + " name: (name) @name) @definition.class", + "(class_interface_clause [(name) (qualified_name)] @name) @impl", + "(property_declaration\n" + + " (property_element (variable_name (name) @name))) @definition.field", + "(function_definition\n" + + " name: (name) @name) @definition.function", + "(method_declaration\n" + + " name: (name) @name) @definition.function"); + } + + @Override + protected List getReferenceQueries() { + return List.of( + "(object_creation_expression\n" + + " [\n" + + " (qualified_name (name) @name)\n" + + " (variable_name (name) @name)\n" + + " ]) @reference.class", + "(function_call_expression\n" + + " function: [\n" + + " (qualified_name (name) @name)\n" + + " (variable_name (name)) @name\n" + + " ]) @reference.call", + "(scoped_call_expression\n" + + " name: (name) @name) @reference.call", + "(member_call_expression\n" + + " name: (name) @name) @reference.call"); + } + + @Override + protected Set processMatchedTags( + Map> matchedTags, + Set definitionTags, + String codeSnippetPath) { + return Set.of(); + } +} diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/PythonRepositoryParser.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/PythonRepositoryParser.java new file mode 100644 index 000000000..f4213db9e --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/PythonRepositoryParser.java @@ -0,0 +1,45 @@ +package ee.carlrobert.codegpt.treesitter.repository.parser; + +import com.intellij.openapi.project.Project; +import ee.carlrobert.codegpt.treesitter.repository.ProcessedTag; +import ee.carlrobert.codegpt.treesitter.repository.Tag; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.treesitter.TreeSitterPython; + +public class PythonRepositoryParser extends RepositoryParser { + + PythonRepositoryParser(Project project) { + super(project, new TreeSitterPython(), "py"); + } + + @Override + protected List getDefinitionQueries() { + return List.of( + "(module (expression_statement (assignment left: (identifier) @name) @definition.constant))", + "(class_definition\n" + + " name: (identifier) @name) @definition.class", + "(function_definition\n" + + " name: (identifier) @name) @definition.function"); + } + + @Override + protected List getReferenceQueries() { + return List.of( + "(call\n" + + " function: [\n" + + " (identifier) @name\n" + + " (attribute\n" + + " attribute: (identifier) @name)\n" + + " ]) @reference.call"); + } + + @Override + protected Set processMatchedTags( + Map> matchedTags, + Set definitionTags, + String codeSnippetPath) { + return Set.of(); + } +} diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/RepositoryParser.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/RepositoryParser.java new file mode 100644 index 000000000..286def6a2 --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/RepositoryParser.java @@ -0,0 +1,192 @@ +package ee.carlrobert.codegpt.treesitter.repository.parser; + +import static java.util.stream.Collectors.toList; +import static java.util.stream.Collectors.toSet; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.intellij.openapi.project.Project; +import com.intellij.openapi.vcs.changes.ChangeListManager; +import com.intellij.openapi.vcs.changes.VcsIgnoreManager; +import com.intellij.openapi.vfs.VirtualFileManager; +import ee.carlrobert.codegpt.treesitter.repository.ProcessedTag; +import ee.carlrobert.codegpt.treesitter.repository.QueryUtil; +import ee.carlrobert.codegpt.treesitter.repository.QueryUtil.QueryResult; +import ee.carlrobert.codegpt.treesitter.repository.Tag; +import ee.carlrobert.codegpt.treesitter.repository.TagType; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.nio.file.Paths; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.treesitter.TSLanguage; + +public abstract class RepositoryParser { + + private final TSLanguage language; + private final List fileExtensions; + private final VcsIgnoreManager ignoreManager; + private final ChangeListManager changeListManager; + + public RepositoryParser(Project project, TSLanguage language, String fileExtensions) { + this(project, language, List.of(fileExtensions)); + } + + public RepositoryParser(Project project, TSLanguage language, List fileExtensions) { + this.language = language; + this.fileExtensions = fileExtensions; + this.ignoreManager = VcsIgnoreManager.getInstance(project); + this.changeListManager = ChangeListManager.getInstance(project); + } + + protected abstract List getDefinitionQueries(); + + protected abstract List getReferenceQueries(); + + protected abstract Set processMatchedTags( + Map> matchedTags, + Set definitionTags, + String codeSnippetPath); + + public boolean isDefinitionsFileExists(String projectName) { + return getExistingDefinitionsFilePath(projectName).toFile().exists(); + } + + public Set createRepositoryDefinitionTags(String projectName, String projectPath) { + var tags = getTags( + getFileMappings(projectPath), + getDefinitionQueries(), + TagType.DEFINITION); + saveRepositoryDefinitionTags(tags, getExistingDefinitionsFilePath(projectName)); + return tags; + } + + public Set processCode( + String projectName, + String projectPath, + String codeSnippetPath, + String codeSnippet) { + var definitionTags = getRepositoryDefinitionTags(projectName, projectPath); + var referenceTags = getCodeSnippetReferenceTags(codeSnippetPath, codeSnippet); + var matchedTags = getMatchedTags(referenceTags, definitionTags); + return processMatchedTags(matchedTags, definitionTags, codeSnippetPath); + } + + protected Set getTags(List fileMappings, List queries, TagType type) { + return fileMappings.parallelStream() + .map(mapping -> QueryUtil.extractTagsFromFile(language, mapping, queries, type)) + .flatMap(Set::stream) + .collect(toSet()); + } + + protected List query(String code, String queryString) { + return QueryUtil.query(language, code, queryString); + } + + private Set getCodeSnippetReferenceTags(String codeSnippetPath, String codeSnippet) { + return getTags( + List.of(new FileMapping(codeSnippetPath, codeSnippet)), + getReferenceQueries(), + TagType.REFERENCE); + } + + private Path getExistingDefinitionsFilePath(String projectName) { + return Paths.get( + System.getProperty("user.home"), + // TODO + String.format(".codegpt/%s/%s/definition-tags.json", projectName, fileExtensions.get(0))); + } + + private List getFileMappings(String pathString) { + try (var pathStream = Files.walk(Path.of(pathString))) { + return pathStream.filter(Files::isRegularFile) + .filter(path -> { + var virtualFile = VirtualFileManager.getInstance().findFileByNioPath(path); + var endsWithExtension = fileExtensions.stream() + .anyMatch(extension -> path.toString().endsWith(extension)); + return endsWithExtension + && virtualFile != null + && !changeListManager.isIgnoredFile(virtualFile) + && !ignoreManager.isPotentiallyIgnoredFile(virtualFile) + && virtualFile.getLength() < Math.pow(1024, 2); + }) + .map(FileMapping::new) + .collect(Collectors.toList()); + } catch (IOException e) { + throw new RuntimeException("Couldn't retrieve files.", e); + } + } + + private Set getRepositoryDefinitionTags(String projectName, String projectPath) { + if (projectPath == null || !isDefinitionsFileExists(projectName)) { + return new HashSet<>(); + } + + try { + var content = Files.readString(getExistingDefinitionsFilePath(projectName)); + return new ObjectMapper().readValue(content, new TypeReference<>() { + }); + } catch (IOException e) { + return new HashSet<>(); + } + } + + private void saveRepositoryDefinitionTags(Set tags, Path path) { + try { + Files.createDirectories(path.getParent()); + Files.writeString(path, new ObjectMapper().writeValueAsString(tags)); + } catch (IOException e) { + throw new RuntimeException("Unable to save repository definition tags"); + } + } + + private Map> getMatchedTags( + Set referenceTags, + Set definitionTags) { + return referenceTags.stream() + .filter(refTag -> List.of("identifier", "type_identifier").contains(refTag.getSymbolName())) + .flatMap(refTag -> definitionTags.stream() + .filter(defTag -> isMatchingIdentifierTag(defTag, refTag))) + .collect(Collectors.groupingBy(Tag::getFilePath, toList())); + } + + private boolean isMatchingIdentifierTag(Tag defTag, Tag refTag) { + return defTag != null && defTag.getName().equals(refTag.getName()); + } + + private String readFileContent(Path path) { + try { + return Files.readString(path); + } catch (IOException e) { + return ""; + } + } + + public class FileMapping { + + private final String path; + private final String content; + + public FileMapping(String path, String content) { + this.path = path; + this.content = content; + } + + public FileMapping(Path path) { + this.path = path.toAbsolutePath().toString(); + this.content = readFileContent(path); + } + + public String getPath() { + return path; + } + + public String getContent() { + return content; + } + } +} diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/RepositoryParserFactory.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/RepositoryParserFactory.java new file mode 100644 index 000000000..31e7b49a5 --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/RepositoryParserFactory.java @@ -0,0 +1,17 @@ +package ee.carlrobert.codegpt.treesitter.repository.parser; + +import com.intellij.openapi.project.Project; + +public class RepositoryParserFactory { + + public static RepositoryParser getParserForFileExtension(Project project, String extension) + throws IllegalArgumentException { + switch (extension) { + case "java": + return new JavaRepositoryParser(project); + // TODO: Add support for more languages + default: + throw new IllegalArgumentException("Unsupported file extension: " + extension); + } + } +} diff --git a/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/TypescriptRepositoryParser.java b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/TypescriptRepositoryParser.java new file mode 100644 index 000000000..8be64306e --- /dev/null +++ b/codegpt-treesitter/src/main/java/ee/carlrobert/codegpt/treesitter/repository/parser/TypescriptRepositoryParser.java @@ -0,0 +1,50 @@ +package ee.carlrobert.codegpt.treesitter.repository.parser; + +import com.intellij.openapi.project.Project; +import ee.carlrobert.codegpt.treesitter.repository.ProcessedTag; +import ee.carlrobert.codegpt.treesitter.repository.Tag; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.treesitter.TreeSitterTypescript; + +public class TypescriptRepositoryParser extends RepositoryParser { + + TypescriptRepositoryParser(Project project) { + super(project, new TreeSitterTypescript(), List.of("ts", "tsx")); + } + + @Override + protected List getDefinitionQueries() { + return List.of( + "(function_signature\n" + + " name: (identifier) @name) @definition.function", + "(method_signature\n" + + " name: (property_identifier) @name) @definition.method", + "(abstract_method_signature\n" + + " name: (property_identifier) @name) @definition.method", + "(abstract_class_declaration\n" + + " name: (type_identifier) @name) @definition.class", + "(module\n" + + " name: (identifier) @name) @definition.module", + "(interface_declaration\n" + + " name: (type_identifier) @name) @definition.interface"); + } + + @Override + protected List getReferenceQueries() { + return List.of( + "(type_annotation\n" + + " (type_identifier) @name) @reference.type", + "(new_expression\n" + + " constructor: (identifier) @name) @reference.class"); + } + + @Override + protected Set processMatchedTags( + Map> matchedTags, + Set definitionTags, + String codeSnippetPath) { + return Set.of(); + } +} diff --git a/codegpt-treesitter/src/test/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParserTest.java b/codegpt-treesitter/src/test/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParserTest.java index e1cc1cd8b..d9731a591 100644 --- a/codegpt-treesitter/src/test/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParserTest.java +++ b/codegpt-treesitter/src/test/java/ee/carlrobert/codegpt/treesitter/CodeCompletionParserTest.java @@ -2,6 +2,7 @@ import static org.assertj.core.api.Assertions.assertThat; +import ee.carlrobert.codegpt.treesitter.completion.CodeCompletionParserFactory; import org.junit.Test; public class CodeCompletionParserTest { diff --git a/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java b/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java index 4c780c352..8d24f5bf3 100644 --- a/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java +++ b/src/main/java/ee/carlrobert/codegpt/PluginStartupActivity.java @@ -11,7 +11,6 @@ import ee.carlrobert.codegpt.completions.you.auth.YouAuthenticationService; import ee.carlrobert.codegpt.completions.you.auth.response.YouAuthenticationResponse; import ee.carlrobert.codegpt.credentials.YouCredentialManager; -import ee.carlrobert.codegpt.settings.GeneralSettings; import ee.carlrobert.codegpt.settings.service.you.YouSettings; import ee.carlrobert.codegpt.ui.OverlayUtil; import org.jetbrains.annotations.NotNull; diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java index 9b9b9a204..bc4e216d4 100644 --- a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionEventListener.java @@ -11,7 +11,7 @@ import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator; import ee.carlrobert.codegpt.CodeGPTBundle; import ee.carlrobert.codegpt.actions.OpenSettingsAction; -import ee.carlrobert.codegpt.treesitter.CodeCompletionParserFactory; +import ee.carlrobert.codegpt.treesitter.completion.CodeCompletionParserFactory; import ee.carlrobert.codegpt.ui.OverlayUtil; import ee.carlrobert.codegpt.util.file.FileUtil; import ee.carlrobert.llm.client.openai.completion.ErrorDetails; diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionListenerBinder.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionListenerBinder.java index d5002d56e..2b5680109 100644 --- a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionListenerBinder.java +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionListenerBinder.java @@ -132,7 +132,7 @@ public void documentChangedNonBulk(@NotNull DocumentEvent event) { // ignore } } else { - codeCompletionService.handleCompletions(editor, caretOffset); + codeCompletionService.requestCompletions(editor, caretOffset); } }); } diff --git a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java index 4e80f0e26..0b1cd509d 100644 --- a/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java +++ b/src/main/java/ee/carlrobert/codegpt/codecompletions/CodeCompletionService.java @@ -72,7 +72,7 @@ public void cancelPreviousCall() { callDebouncer.cancelPreviousCall(); } - public void handleCompletions(Editor editor, int offset) { + public void requestCompletions(Editor editor, int offset) { PREVIOUS_INLAY_TEXT.set(editor, null); if (project.isDisposed() diff --git a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java index b2e4a8268..55a8a16c2 100644 --- a/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java +++ b/src/main/java/ee/carlrobert/codegpt/completions/CompletionRequestProvider.java @@ -40,6 +40,7 @@ import ee.carlrobert.llm.client.you.completion.YouCompletionRequest; import ee.carlrobert.llm.client.you.completion.YouCompletionRequestMessage; import java.nio.charset.StandardCharsets; +import java.nio.file.Path; import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -77,7 +78,7 @@ public CompletionRequestProvider(Conversation conversation) { } public static String getPromptWithContext(List referencedFiles, - String userPrompt) { + String userPrompt) { var includedFilesSettings = IncludedFilesSettings.getCurrentState(); var repeatableContext = referencedFiles.stream() .map(item -> includedFilesSettings.getRepeatableContext() @@ -285,9 +286,22 @@ private List buildMessages( messages.add(new OpenAIChatCompletionMessage("user", prompt)); } else { if (callParameters.getConversationType() == ConversationType.DEFAULT) { - messages.add(new OpenAIChatCompletionMessage( - "system", - ConfigurationSettings.getCurrentState().getSystemPrompt())); + var repositoryMapping = message.getRepositoryMapping(); + if (repositoryMapping != null && !repositoryMapping.isEmpty()) { + var relevantContext = repositoryMapping.stream().map( + processedTag -> ("// " + Path.of(processedTag.getFilePath()).getFileName().toString() + + "\n" + + processedTag.getModifiedContent()).replace( + "\n", "\n| ")).collect(joining("\n\n")); + + messages.add(new OpenAIChatCompletionMessage( + "system", + COMPLETION_SYSTEM_PROMPT + "\n\nRelevant context:\n\n" + relevantContext)); + } else { + messages.add(new OpenAIChatCompletionMessage( + "system", + ConfigurationSettings.getCurrentState().getSystemPrompt())); + } } if (callParameters.getConversationType() == ConversationType.FIX_COMPILE_ERRORS) { messages.add(new OpenAIChatCompletionMessage("system", FIX_COMPILE_ERRORS_SYSTEM_PROMPT)); diff --git a/src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java b/src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java index 5ff03f71a..490232b5f 100644 --- a/src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java +++ b/src/main/java/ee/carlrobert/codegpt/conversations/message/Message.java @@ -2,9 +2,11 @@ import com.fasterxml.jackson.annotation.JsonCreator; import com.fasterxml.jackson.annotation.JsonProperty; +import ee.carlrobert.codegpt.treesitter.repository.ProcessedTag; import ee.carlrobert.llm.client.you.completion.YouSerpResult; import java.util.List; import java.util.Objects; +import java.util.Set; import java.util.UUID; public class Message { @@ -15,6 +17,7 @@ public class Message { private String userMessage; private List serpResults; private List referencedFilePaths; + private Set repositoryMapping; public Message(String prompt, String response) { this(prompt); @@ -71,6 +74,14 @@ public void setReferencedFilePaths(List referencedFilePaths) { this.referencedFilePaths = referencedFilePaths; } + public Set getRepositoryMapping() { + return repositoryMapping; + } + + public void setRepositoryMapping(Set repositoryMapping) { + this.repositoryMapping = repositoryMapping; + } + @Override public boolean equals(Object obj) { if (obj == this) { diff --git a/src/main/java/ee/carlrobert/codegpt/indexes/FolderStructureTreePanel.java b/src/main/java/ee/carlrobert/codegpt/indexes/FolderStructureTreePanel.java index d6e13dba8..e34fad40c 100644 --- a/src/main/java/ee/carlrobert/codegpt/indexes/FolderStructureTreePanel.java +++ b/src/main/java/ee/carlrobert/codegpt/indexes/FolderStructureTreePanel.java @@ -233,6 +233,7 @@ public void customizeRenderer(JTree t, ".settings", "node_modules", "vendor", + "venv", "lib", "build", "target", diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java index 58afb1f2f..c7905d593 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ChatToolWindowTabPanel.java @@ -9,6 +9,7 @@ import com.intellij.openapi.application.ApplicationManager; import com.intellij.openapi.diagnostic.Logger; import com.intellij.openapi.editor.impl.EditorImpl; +import com.intellij.openapi.progress.impl.BackgroundableProcessIndicator; import com.intellij.openapi.project.Project; import com.intellij.ui.JBColor; import com.intellij.util.ui.JBUI; @@ -35,6 +36,7 @@ import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.TotalTokensPanel; import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.UserPromptTextArea; import ee.carlrobert.codegpt.toolwindow.chat.ui.textarea.UserPromptTextAreaHeader; +import ee.carlrobert.codegpt.treesitter.repository.parser.RepositoryParserFactory; import ee.carlrobert.codegpt.util.EditorUtil; import ee.carlrobert.codegpt.util.file.FileUtil; import ee.carlrobert.embedding.ReferencedFile; @@ -241,9 +243,28 @@ private void handleSubmit(String text) { var selectionModel = editor.getSelectionModel(); var selectedText = selectionModel.getSelectedText(); if (selectedText != null && !selectedText.isEmpty()) { - var fileExtension = FileUtil.getFileExtension( - ((EditorImpl) editor).getVirtualFile().getName()); - message = new Message(text + format("\n```%s\n%s\n```", fileExtension, selectedText)); + var fileName = ((EditorImpl) editor).getVirtualFile().getName(); + var fileExtension = FileUtil.getFileExtension(fileName); + try { + var parser = RepositoryParserFactory.getParserForFileExtension(project, fileExtension); + if (parser.isDefinitionsFileExists(project.getName())) { + var repositoryMapping = parser.processCode( + project.getName(), + project.getBasePath(), + ((EditorImpl) editor).getVirtualFile().getPath(), + selectedText); + message = new Message(text + format("\n```%s\n%s\n```", fileExtension, selectedText)); + message.setRepositoryMapping(repositoryMapping); + } else { + var processIndicator = createProgressIndicator(); + ApplicationManager.getApplication().executeOnPooledThread(() -> { + parser.createRepositoryDefinitionTags(project.getName(), project.getBasePath()); + processIndicator.processFinish(); + }); + } + } catch (IllegalArgumentException e) { + // ignore + } selectionModel.removeSelection(); } } @@ -251,6 +272,11 @@ private void handleSubmit(String text) { sendMessage(message, ConversationType.DEFAULT); } + private BackgroundableProcessIndicator createProgressIndicator() { + return new BackgroundableProcessIndicator(project, + "Indexing repository", null, null, true); + } + private JPanel createUserPromptPanel(ServiceType selectedService) { var panel = new JPanel(new BorderLayout()); panel.setBorder(JBUI.Borders.compound( diff --git a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/UserMessagePanel.java b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/UserMessagePanel.java index 5e723ed4c..71633e85f 100644 --- a/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/UserMessagePanel.java +++ b/src/main/java/ee/carlrobert/codegpt/toolwindow/chat/ui/UserMessagePanel.java @@ -1,5 +1,7 @@ package ee.carlrobert.codegpt.toolwindow.chat.ui; +import static java.util.stream.Collectors.toList; + import com.intellij.openapi.Disposable; import com.intellij.openapi.project.Project; import com.intellij.ui.ColorUtil; @@ -9,6 +11,7 @@ import com.intellij.util.ui.JBUI; import ee.carlrobert.codegpt.Icons; import ee.carlrobert.codegpt.conversations.message.Message; +import ee.carlrobert.codegpt.treesitter.repository.ProcessedTag; import ee.carlrobert.codegpt.settings.GeneralSettings; import java.awt.BorderLayout; import javax.swing.JPanel; @@ -28,12 +31,21 @@ public UserMessagePanel(Project project, Message message, Disposable parentDispo add(headerPanel, BorderLayout.NORTH); var referencedFilePaths = message.getReferencedFilePaths(); + var referencedRepositoryMapping = message.getRepositoryMapping(); if (referencedFilePaths != null && !referencedFilePaths.isEmpty()) { add(new SelectedFilesAccordion(project, referencedFilePaths), BorderLayout.CENTER); add(createResponseBody( project, message.getUserMessage(), parentDisposable), BorderLayout.SOUTH); + } else if (referencedRepositoryMapping != null && !referencedRepositoryMapping.isEmpty()) { + add(new SelectedFilesAccordion(project, referencedRepositoryMapping.stream() + .map(ProcessedTag::getFilePath) + .collect(toList())), BorderLayout.CENTER); + add(createResponseBody( + project, + message.getUserMessage(), + parentDisposable), BorderLayout.SOUTH); } else { add(createResponseBody(project, message.getPrompt(), parentDisposable), BorderLayout.SOUTH); }