Improve code context given to the LLM

[scastlara]

Is your feature request related to a problem? Please describe.
At the moment, we pass the LLM all the contents of all files in the PR.
Can we do better?

Describe the solution you'd like
There are many possible solutions 👇

1: Enriched context based on import dependency tree

• Given an import-dependency tree of the codebase, can we give the contents of all files at a particular distance of the files in the PR?
• Instead of passing all content of all these files (using the above approach) can we intelligently pass something more relevant? For instance, a context window around elements that appear in the PR (functions, classes, etc) that in files at the selected distance in the dependency tree.
• BONUS: Can we use tree-sitter or other similar tools to treat code not just as raw text but as "code" so that we pass relevant parts of it? (whole functions, classes etc?).

2: Enriched context based on grepping code elements

• Parse the ast or get code objects with something like tree-sitter in the PR diff/files of the PR.
• Select some elements from there based on some heuristic (number of appearances? some measure of centrality/importance?)
• Grep those in the whole codebase.
• Retrieve content using a context window around those elements.

Any more ideas?

[scastlara]

Quick way to parse functions and classes with tree-sitter

# Build the shared library with the Python grammar
import tree_sitter_python as tspython
from tree_sitter import Language, Parser

PY_LANGUAGE = Language(tspython.language())
parser = Parser(language=PY_LANGUAGE)
code = b"""
def hello(name):
    print(f"Hello, {name}")

class Foo:
    def method(self):
        pass
"""


tree = parser.parse(code)
root = tree.root_node


def extract_defs(node, source_code):
    results = []

    if node.type in ("function_definition", "class_definition"):
        name_node = node.child_by_field_name("name")
        if name_node is not None:
            name = source_code[name_node.start_byte : name_node.end_byte].decode("utf-8")
            line, col = name_node.start_point  # (line, column), zero-based
            results.append(
                {
                    "type": node.type.replace("_definition", ""),  # "function" or "class"
                    "name": name,
                    "line": line + 1,  # convert to 1-based line numbers
                    "column": col + 1,
                }
            )

    # Recurse into children
    for child in node.children:
        results.extend(extract_defs(child, source_code))

    return results


items = extract_defs(root, code)

for item in items:
    print(f"{item['type'].capitalize()} '{item['name']}' at line {item['line']} (col {item['column']})")