Chunk Code by Function With tree-sitter, Not by 512 Tokens

Retrieval results in codebase assistants often start mid-function at the except clause, with neither the def line nor the class name in sight—only the file path hints at where it came from. Yet the model will confidently write an explanation of a function whose name it has never seen. The root of the problem isn't the model; it's the chunker: it slices every 512 tokens the way it would handle prose, and the cut lands squarely in the middle of a function.

Code has natural seams at def and class, and a token counter knows nothing about them. The approach in this article: parse Python with tree-sitter, chunk at function and class boundaries, write the file path and enclosing symbol into metadata, and split oversized definitions further along statement boundaries. Then score both chunkers with an annotated query set so you can see the difference on your own machine.

What fixed windows do to functions

A fixed window has no idea where a function begins. If a function starts at token 470 of a 512-token window, that chunk only gets the decorators, the signature, and four lines of body; the next chunk gets the rest of the body with no name attached. When you ask "where's the function that parses config files," the retriever matches a signature with no logic, or logic with no signature.

Carnegie Mellon University and Augment Code's cAST paper found the same problem: "existing line-based chunking heuristics often break semantic structure, splitting a function or merging unrelated code together." Their AST-based chunker splits oversized syntax nodes and merges small siblings under a size limit, achieving a 4.3 percentage point Recall@5 improvement on RepoEval retrieval and a 2.67 percentage point Pass@1 improvement on SWE-bench generation.

Frameworks are moving in the same direction. LlamaIndex offers a CodeSplitter that parses with tree-sitter and, by default in source, caps chunks at 1500 characters. The cAST authors released astchunk, which measures chunk size by non-whitespace characters and supports Python, Java, C#, and TypeScript. These tools are ready to use, but writing your own teaches you what they're doing for you and how to tune them when your codebase has a 600-line function.

tree-sitter is an incremental parser generator and runtime library that produces a concrete syntax tree (CST) for source code rather than an abstract syntax tree (AST). It supports dozens of languages, parses tolerantly, and is extremely fast. The key point: tree-sitter compiles each language's grammar rules into a C dynamic library, and the Python runtime loads them through separate wheels like tree_sitter_python. The node type names it exposes vary by language—in Python a function definition is a function_definition and a class is a class_definition, while in JavaScript they are function_declaration and class_declaration. So the chunker needs its own DEFS set per language. tree-sitter's error tolerance means it still produces a usable syntax tree when the code has syntax errors, and the has_error attribute on nodes can be used to identify damaged regions.

Two chunkers, the same token budget

Install dependencies: py-tree-sitter 0.26.0 requires Python 3.10 or later, and the Python grammar ships as a separate wheel. The embedding model is BAAI/bge-small-en-v1.5, 384-dimensional, with a max sequence length of 512 tokens. Both chunkers get a 480-token budget, counted with the embedding model's own tokenizer, so that short headers and special tokens also fit under 512 and no content is silently truncated—that's a fair "512-token chunk."

pip install tree-sitter tree-sitter-python \
  sentence-transformers numpy

The first file, chunkers.py, starts with the Chunk type, and embed_text() puts a path-and-symbol breadcrumb like # json/decoder.py :: JSONDecoder.raw_decode on the first line of the embedded text. The baseline chunker fixed_chunks accumulates tokens line by line, emits a chunk when the budget would be exceeded, then backs off about 64 tokens for overlap. It aligns to line ends and overlaps, making it more forgiving than a naive token splitter, but it's missing exactly two things: syntax-aware boundaries and symbol names.

**BAAI/bge-small-en-v1.5** is a general-purpose English embedding model released by BAAI, 384-dimensional, based on the BERT architecture, with a max sequence length of 512 tokens. It uses a WordPiece tokenizer—which is not ideal for code: WordPiece splits on subwords, breaking `snake_case` and `camelCase` into meaningless fragments and chopping `raw_decode` into `raw`, `##_`, `decode`, and so on. The embedding model's max length means inputs over 512 tokens are silently truncated, so the example keeps the chunk budget at 480 tokens, leaving headroom for the path/symbol breadcrumb at the top of the embedded text and special tokens. In production you should switch to an embedding model trained on code corpora (like CodeBERT, jina-code, etc.) and validate the swap against the same evaluation set.
# chunkers.py
from dataclasses import dataclass
from typing import Callable

import tree_sitter_python as tspython
from tree_sitter import Language, Parser

PARSER = Parser(Language(tspython.language()))
DEFS = {"function_definition", "class_definition"}
Count = Callable[[str], int]

@dataclass
class Chunk:
    path: str
    symbol: str  # "JSONDecoder.raw_decode", "<module>", or ""
    start: int   # 1-based first line
    end: int     # 1-based last line, inclusive
    text: str

    def embed_text(self) -> str:
        head = f"# {self.path}"
        if self.symbol:
            head += f" :: {self.symbol}"
        return f"{head}\n{self.text}"

Using src.split("\n") rather than splitlines() is deliberate: tree-sitter counts lines only by \n, while splitlines() also splits on form feeds, and some older Python files contain form feeds—mixing the two would shift every line number after the first form feed by one.

The AST chunker is in the same file. block walks the statements of a module or class body, handing each definition to define, while statements between definitions (imports, constants, class docstrings) are collected into loose and packed into <module> chunks or class-named chunks. define emits a definition whole when it fits; decorators count as part of the definition because tree-sitter-python wraps them in decorated_definition. The symbol path is the breadcrumb: a method inside JSONDecoder gets JSONDecoder.raw_decode, so a query can match literally even if the method body never contains the words "JSON decoder."

What if a function is bigger than the budget

Function-level chunking breaks down on oversized definitions. A 300-line class or a 200-line function won't fit in 480 tokens, and handing the whole thing to the embedding model is useless anyway. all-MiniLM-L6-v2's model card says it plainly: "By default, input text longer than 256 word pieces is truncated." Truncated content is invisible to retrieval, and there's no warning.

The chunker handles the size problem in three steps:

Oversized classes: recurse. define calls block on the class body, and each method becomes its own chunk with symbol Class.method. The class signature becomes a head, prepended to chunks of class-level loose statements, so docstrings and class attributes still know which class they belong to.

Oversized functions: pack statements under the signature. pack walks the function body's top-level statements, greedily merging adjacent statements until the next one would exceed the budget. Every chunk starts with the full signature, decorators included. A fragment retrieved from the middle of a long function still starts with def _make_iterencode(...). This is exactly the merge half of cAST's "split large nodes, merge small siblings."

Oversized single statements: fall back to windows. A 150-line if block is one statement. When emit finds it doesn't fit, it runs the line-window chunker on just that statement, with the signature still on top. cAST recurses into the statement's child nodes for cleaner boundaries; the fallback here is the cheap version, and should be replaced with recursion when such statements are common in your codebase.

There are two costs to know about: headers eat into the budget, so a function with a 15-line signature leaves less room for its body per fragment; and pack re-tokenizes growing spans repeatedly, which is quadratic in the worst case. A few files don't matter, but in a monorepo you should count tokens once per statement and then accumulate.

Scoring both on an annotated query set

tree-sitter still produces useful results when source code has syntax errors, so a half-refactored file still yields a syntax tree. py-tree-sitter exposes has_error on every node; check the root node, and when a file is badly broken, fall back to fixed_chunks for that file rather than indexing garbage boundaries.

The evaluation needs a corpus every reader has—the Python standard library is right there on your disk. The queries describe what code does, mostly without using identifiers, like searching a codebase you didn't write. Each annotation is a file plus a symbol path. eval.py reads in six files, builds both chunk lists, and recall checks top-5 results in two ways:

• touched@5: whether any chunk in the top 5 overlaps the target function at all. This metric is loose, and it's where the token-window chunker looks best, because a chunk holding three lines of the target function still counts as a hit.
• whole@5: whether a top-5 chunk contains the entire function. This is what an LLM answer actually needs.

The gap between them is the quantified version of the "function body loses its name" problem. The article deliberately doesn't print results: your Python version changes the standard library files, and 12 queries is a smoke test, not a benchmark. Run it and read your own numbers. Some annotations are longer than the 480-token budget, and the AST chunker splits them too, so neither chunker can achieve whole on those—namedtuple and _make_iterencode are obvious examples. To make a real decision, write 50 to 100 queries against your own repository, drawn from code reviews or questions people actually asked in Slack, and version-control them alongside the chunker.

Recall@K is a common metric in retrieval evaluation: among the top K returned results, a hit is counted if any item is relevant to the ground-truth answer; Recall@5 is the fraction of all queries whose top-5 contains the target. The example splits a hit into two finer granularities: touched@5 requires any chunk in the top 5 to overlap the target function on any line, which is loose and friendly to fixed-window chunkers—a chunk that happens to contain three lines of the target function counts. whole@5 requires at least one chunk in the top 5 to contain the entire function, which is the prerequisite for an LLM to generate an accurate answer from it. The gap between the two directly quantifies the cost of "the function body and its signature/symbol name being torn apart," making it the most decision-relevant metric when choosing a chunking strategy.

Before pointing it at your repo

Going from the standard-library demo to indexing a monorepo takes four changes:

More languages. tree-sitter covers most grammars you'll encounter, but node type names vary by grammar. function_definition is Python's spelling; consult each grammar's node-types.json and give each language its own DEFS set. For maintained versions for Python, Java, C#, and TypeScript, astchunk already does this.

Parse errors. As above, check the root node with has_error and fall back to fixed chunking for damaged files.

Code-aware embedding models. bge-small-en-v1.5 is a general English model, and its WordPiece tokenizer shreds snake_case and camelCase. Fine for a demo; in production you should use the same eval.py to compare code-trained embedding models on the annotated set—that's exactly what the set is for.

Keyword retrieval alongside vectors. Exact identifiers, error codes, and config keys are where dense retrieval is weakest. Index the symbol field into a keyword index too and fuse the results; a query containing raw_decode should find raw_decode directly rather than making the embedding model guess.

A retriever can only return what the chunker keeps together, and code makes this plain. The stages on both ends of retrieval and reranking, and chunking patterns for production RAG, are discussed more fully in the RAG Pocket Guide.