Building Better Code RAG: Semantic Chunking with Tree-Sitter

When I first built the retrieval pipeline for AskRepo, I did what everyone does when they need to feed code into an embedding model → split it every N characters and move on. It worked, in the sense that it ran without errors. It did not work in the sense that mattered, the search results were bad.

This post is about why that happened, and how switching to AST-aware chunking with Tree-Sitter fixed it.

A quick primer, if RAG is new to you

RAG (Retrieval-Augmented Generation) is the pattern behind "chat with your codebase" tools like AskRepo. The idea is simple: instead of asking an LLM to answer a question from memory alone, you first retrieve relevant pieces of your codebase, then hand those pieces to the LLM as context before it answers.

To make retrieval possible, you break your codebase into small pieces ("chunks"), convert each chunk into a vector ("embedding") that captures its meaning, and store those vectors in a database. When a user asks a question, you embed the question too, and pull out the chunks whose vectors are closest to it.

The whole system lives or dies on one decision that's easy to get wrong: how you chunk the code in the first place.

The problem with naive chunking

Say you split a Python file every 500 characters. You'll eventually get something like this:

def calculate_total(items):
    """Calculate the total price of all items with tax."""
    subtotal = 0
    for item in items:
        subtotal += item.price * item.qu

That chunk ends mid-token. The embedding for it has no idea what qu was going to be, whether the function returns anything, or what "total" even means in context. If someone searches "calculate total with tax," this chunk might not even surface and if it does, it's useless once retrieved.

Character-based and even token-based splitting treats code like prose. But code isn't prose. It has structure functions, classes, blocks, scopes and a chunk boundary that ignores that structure is basically guaranteed to cut something important in half.

Enter the AST

An Abstract Syntax Tree (AST) is the structural representation a compiler (or your editor's syntax highlighter) builds from source code. Every function, every class, every import statement becomes a typed node with a defined start and end a function_declaration, a class_body, an import_statement each with exact byte ranges in the original file.

For chunking, this is exactly what you want: boundaries that respect meaning, not arbitrary character counts. That's what Tree-Sitter gives you a fast, incremental parser that produces this kind of tree for dozens of languages, with a consistent API across all of them.

Mapping code to tree nodes

Tree-Sitter parses text into a tree of language-specific node types, but "meaningful unit" isn't the same thing across languages. A function in Python is a function_definition, but in Rust it's a function_item. In TypeScript, you also want interfaces and type aliases, which don't exist as concepts in Python at all.

So the first step is deciding, per language, which node types actually count as a "symbol" worth chunking on its own:

# backend/lib/ast_parser/language_config.py
SYMBOL_NODE_TYPES = {
    "python": {
        "function_definition",
        "class_definition",
    },
    "typescript": {
        "function_declaration",
        "class_declaration",
        "method_definition",
        "interface_declaration",
        "type_alias_declaration",
        # ...
    },
    "rust": {
        "function_item",
        "impl_item",
        "struct_item",
        "trait_item",
    },
}

When a file is parsed, we look up its language from the file extension, grab this set of target types, and any node matching one of them gets pulled out as a symbol chunk.

Walking the AST: scopes and hierarchy

Real code is nested. A method lives inside a class, which might live inside a namespace. If retrieval returns a method chunk in isolation, it's useful to know which class that method belongs to otherwise you've handed the LLM a fragment with no home.

To preserve this, we walk the AST recursively. When we hit a node of interest lets say, a class definition we extract its metadata and pass its name down as the parent_name for anything nested inside it, like its methods:

# backend/lib/ast_parser/ast_walker.py
def walk(node, source_lines, language, target_types, parent_name=""):
    symbols = []
 
    if node.type in target_types:
        name = get_name(node, language)
        docstring = extract_docstring(node, language, source_lines)
        start, end = node.start_point[0], node.end_point[0]
        source = "\n".join(source_lines[start : end + 1])
 
        symbol = Symbol(
            name=name,
            symbol_type=node.type,
            start_line=start + 1,
            end_line=end + 1,
            source=source,
            docstring=docstring,
            parent_name=parent_name,
        )
 
        # Walk children, setting current symbol as their parent
        for child in node.children:
            nested = walk(child, source_lines, language, target_types, parent_name=name)
            symbol.children.extend(nested)
 
        symbols.append(symbol)
    else:
        # Keep traversing down
        for child in node.children:
            symbols.extend(walk(child, source_lines, language, target_types, parent_name))
 
    return symbols

This gives us a tree of symbols, which we then flatten into a list. Class definitions and method definitions become separate, independently searchable chunks but their parent-child relationship is preserved in the metadata, so a retrieved method still knows which class it came from.

The devil in the details: language-specific extraction

Walking the tree sounds simple in theory. In practice, extracting names and docstrings gets messy fast, because every language's syntax tree is shaped differently. A few examples:

1. Go methods. A method receiver like func (s *Server) Start() {} buries the interesting part you have to parse the receiver itself to figure out that Server is the parent, not just Start.
2. JavaScript/TypeScript variables. A declaration like const calculateTotal = ... means the actual identifier is nested inside a variable_declarator node, not sitting at the top level.
3. Docstrings and comments. Python stores its docstring as a string literal expression inside the function body. Go, Rust, and Java instead put a comment (// or /** ... */) directly above the block. There's no single rule that works for both.

We handle this with extractors that route based on language:

# backend/lib/ast_parser/extractors.py
def extract_docstring(node, lang, source_lines):
    match lang:
        case "python":
            return extract_python_comment(node)
        case "go":
            return extract_go_comment(node, source_lines)
        case "java" | "cpp" | "c" | "javascript" | "typescript" | "rust":
            return extract_leading_comment(node, source_lines)
        case _:
            return ""

None of this logic is exciting on its own, but skipping it means silently losing context (or names) for entire languages.

Don't lose the context: handling "orphans"

If you only extract functions and classes, what happens to everything else import statements, global constants, module-level variables? They're not "symbols" by our definition, but they're often crucial context. If a search needs to know which file imports jwt, and you've thrown away the imports block because it wasn't inside a function, retrieval quietly fails.

To fix this, we compute orphan lines: any line in the file not covered by an extracted symbol.

# backend/lib/ast_parser/ast_walker.py
def find_orphan_lines(source_lines, symbols):
    covered = set()
 
    for sym in flatten(symbols):
        for line_no in range(sym.start_line, sym.end_line + 1):
            covered.add(line_no)
 
    orphan_lines = [
        line for i, line in enumerate(source_lines, start=1)
        if i not in covered
    ]
    return "\n".join(orphan_lines).strip()

These orphan lines get chunked and indexed as their own module_context chunk type, so the file stays fully searchable not just the parts that happened to be inside a function or class.

The guardrail: hybrid token splitting

AST chunking solves the "cut mid-function" problem, but it introduces a new one: what if a single function is enormous? A 2000-line legacy function is still one AST node, and it can easily blow past an embedding model's token limit (say, 512 tokens).

So symbols are the default chunk boundary, but they're not the only rule. If a symbol fits comfortably within the token limit, it becomes one clean, structured chunk. If it doesn't, we fall back to splitting it with a sliding token window (using tiktoken), with some overlap between windows so context isn't lost at the seams.

# backend/utils/chunker.py
def chunk_symbol(symbol: Symbol, file_path: str, repo_name: str, commit_sha: str):
    splits = _token_split(symbol.source)  # Splits by token window if source is too large
    chunks = []
    for i, (text, token_count, rel_start, rel_end) in enumerate(splits):
        chunks.append(Chunk(
            id=_make_id(repo_name, file_path, i, text),
            text=text,
            metadata={
                "file_path": file_path,
                "start_line": symbol.start_line + rel_start,
                "end_line": symbol.start_line + rel_end,
                "symbol_name": symbol.name,
                "symbol_type": symbol.symbol_type,
                "docstring": symbol.docstring,
            }
        ))
    return chunks

This is the piece that ties everything together: symbols, orphan lines, and plain-text files all pass through the same chunker and come out the other side as Chunk objects with full metadata attached ready to be embedded and stored.

Where this leaves us

None of this requires exotic technology Tree-Sitter, a bit of tree-walking, and some careful bookkeeping around edge cases. But the payoff is chunks that actually correspond to something a developer would recognize as "one idea": a whole function, a whole class, a whole method instead of an arbitrary slice of text that happens to be 500 characters long.

If you're building retrieval over code and your results feel off, the chunking step is usually where to look first. It's not the most glamorous part of a RAG pipeline, but it's the part everything else depends on.

If you want to see the full implementation, the code above is from AskRepo a browser extension that lets you chat with any GitHub repository using natural language.