- Book: RAG Pocket Guide
- Also by me: Thinking in Go (2-book series) — Complete Guide to Go Programming + Hexagonal Architecture in Go
- My project: Hermes IDE | GitHub — an IDE for developers who ship with Claude Code and other AI coding tools
- Me: xgabriel.com | GitHub
The chunk the codebase assistant retrieved for "where does the retry delay get computed?" starts halfway down a function, at an except clause, and has no def line, no class name, nothing but a file path to say where it came from. The model writes a confident answer about a function it never saw the name of.
The bug is in the chunker. It cut the file every 512 tokens, the way it cuts prose, and the cut landed in the middle of a function. Code has natural seams at def and class; a token counter ignores them.
The fix fits in two small files. Parse Python with tree-sitter, cut at function and class boundaries, carry the file path and the enclosing symbol as metadata, and split the definitions that are too big on statement boundaries. Then score both chunkers on a labelled query set, so you see the difference on your own machine instead of taking my word for it.
What a token window does to a function
A fixed window has no idea where a function starts. If a function begins at token 470 of a 512-token window, the chunk gets the decorator, the signature and four lines of body. The next chunk gets the rest of the body with no name attached. Ask for "the function that parses the config file" and the retriever has to match either a signature with no logic or logic with no signature.
The cAST paper from Carnegie Mellon and Augment Code found the same problem: "Existing line-based chunking heuristics often break semantic structures, splitting functions or merging unrelated code." Their AST-based chunker, which splits large syntax nodes and merges small siblings under a size limit, reported a 4.3-point Recall@5 gain on RepoEval retrieval and a 2.67-point Pass@1 gain on SWE-bench generation over the fixed-size baseline.
Frameworks have moved the same way. LlamaIndex ships a CodeSplitter that parses with tree-sitter and, in its source, caps chunks at 1,500 characters by default. The cAST authors published astchunk, which measures chunk size in non-whitespace characters and supports Python, Java, C# and TypeScript. Both are fine to use. Writing your own once is how you find out what they are doing for you, and how you tune it when your codebase has 600-line functions.
Token windows cut wherever the count runs out. AST chunks cut where the code does.
Two chunkers, same token budget
Install the pieces. py-tree-sitter 0.26.0 needs Python 3.10 or newer, and the Python grammar ships as its own wheel.
pip install tree-sitter tree-sitter-python \
sentence-transformers numpy
The embedder is BAAI/bge-small-en-v1.5: a small English model with 384 dimensions and a 512-token max sequence length. Both chunkers get the same budget of 480 tokens, counted with the embedder's own tokenizer, so a short header and the special tokens still fit under 512 and nothing gets silently truncated. That is the fair version of "512-token chunks".
First file, chunkers.py, starting with the Chunk type:
# 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}"
The baseline comes next. It adds lines while their token count stays under the budget, emits a chunk, then backs up about 64 tokens for the overlap. The BGE tokenizer emits no token for a newline, so summing per line counts the same tokens the AST chunker counts on joined text:
# chunkers.py (continued)
def fixed_chunks(path: str, src: str, ntok: Count,
size: int = 480, overlap: int = 64):
lines = src.split("\n")
out, i = [], 0
while i < len(lines):
j, used = i, 0
while j < len(lines):
cost = ntok(lines[j])
if used + cost > size and j > i:
break
used += cost
j += 1
text = "\n".join(lines[i:j])
out.append(Chunk(path, "", i + 1, j, text))
if j >= len(lines):
break
back, k = 0, j
while k > i + 1 and back < overlap:
k -= 1
back += ntok(lines[k])
i = k
return out
The baseline is slightly generous: it snaps to line ends and overlaps about 64 tokens, where a plain token splitter would cut mid-line. Its chunks still get the file path in the header. The only things it lacks are syntax-aware boundaries and the symbol name, which keeps the comparison about exactly those two things.
src.split("\n") rather than splitlines() is deliberate. tree-sitter counts rows by \n only, and splitlines() also splits on form feeds, which some older Python files contain. Mixing the two puts every line number after the first form feed off by one.
The AST chunker goes in the same file. First the helpers that write chunks:
# chunkers.py (continued)
def unwrap(node):
if node.type == "decorated_definition":
return node.child_by_field_name("definition")
return node
def name_of(defn) -> str:
return defn.child_by_field_name("name").text.decode()
def ast_chunks(path: str, src: str, ntok: Count,
budget: int = 480):
lines = src.split("\n")
tree = PARSER.parse(src.encode())
out: list[Chunk] = []
def span(a, b): # 0-based rows, inclusive
return "\n".join(lines[a:b + 1])
def emit(sym, a, b, head=""):
body = span(a, b)
if ntok(head + body) <= budget:
out.append(Chunk(path, sym, a + 1, b + 1, head + body))
return
# one statement is bigger than the budget: window it
room = max(budget - ntok(head), 64)
for c in fixed_chunks(path, body, ntok, size=room):
out.append(Chunk(path, sym, a + c.start, a + c.end,
head + c.text))
emit writes one chunk when the span fits and falls back to line windows only when a single statement is bigger than the budget. pack sits on top of it and greedily groups neighbouring statements. It is still inside ast_chunks, so keep the indentation:
# ast_chunks, continued
def pack(nodes, sym, head=""):
group = []
for node in nodes:
if group:
a = group[0].start_point[0]
b = node.end_point[0]
if ntok(head + span(a, b)) > budget:
emit(sym, a, group[-1].end_point[0], head)
group = []
group.append(node)
if group:
emit(sym, group[0].start_point[0],
group[-1].end_point[0], head)
The last two helpers decide what a unit is, and the final line starts the walk at the module root:
# ast_chunks, continued
def define(node, defn, scope):
name = name_of(defn)
sym = f"{scope}.{name}" if scope else name
a, b = node.start_point[0], node.end_point[0]
if ntok(span(a, b)) <= budget:
emit(sym, a, b)
return
body = defn.child_by_field_name("body")
head = span(a, body.start_point[0] - 1) + "\n"
if defn.type == "class_definition":
block(body.named_children, sym, head)
else:
pack(body.named_children, sym, head)
def block(nodes, scope, head=""):
loose = []
for node in nodes:
defn = unwrap(node)
if defn.type in DEFS:
if loose:
pack(loose, scope or "<module>", head)
loose = []
define(node, defn, scope)
else:
loose.append(node)
if loose:
pack(loose, scope or "<module>", head)
block(tree.root_node.named_children, "")
return out
Walk it top-down. block goes through the statements of a module or a class body. Each definition goes to define, and the statements between definitions (imports, constants, a class docstring) collect into loose and get packed into <module> chunks or chunks named after the class. define emits a definition whole if it fits. Decorators count as part of the definition, because tree-sitter-python wraps them in a decorated_definition whose definition field holds the function or class.
The symbol path is the breadcrumb. A method inside JSONDecoder gets JSONDecoder.raw_decode, and embed_text() puts # json/decoder.py :: JSONDecoder.raw_decode on the first line of what gets embedded. A query that says "JSON decoder" now has a literal match even when the method body never says it.
When one function is bigger than the budget
Function-level chunks break down on the big ones. A 300-line class or a 200-line function does not fit in 480 tokens, and dumping it whole into the embedder does not help. The model card for all-MiniLM-L6-v2 says so directly: "By default, input text longer than 256 word pieces is truncated." Everything past the cutoff is invisible to retrieval, and nothing warns you.
The chunker above handles size in three steps.
Too-big class: recurse. define calls block on the class body, so every method becomes its own chunk with a Class.method symbol. The class signature becomes a head that is prepended to the chunks of loose class-level statements, so the docstring and class attributes still say which class they belong to.
Too-big function: pack statements under the signature. pack walks the body's top-level statements and greedily merges neighbours until the next one would blow the budget. Every piece starts with the full signature, decorators included. A retrieved piece from the middle of a long function still opens with def _make_iterencode(...). This is the merge half of what cAST does: split big nodes, merge small siblings.
Too-big single statement: fall back to windows. A 150-line if block is one statement. emit notices it does not fit and runs the line-window chunker over just that statement, with the signature still on top. cAST recurses into the statement's children instead, which keeps boundaries cleaner. The fallback is the cheap version; swap it for recursion when your codebase has a lot of these.
Small siblings merge, big nodes split, and every split piece keeps the signature on top.
Two costs to know about. The header eats budget, so a function with a 15-line signature leaves less room for the body in each piece. And pack re-tokenizes growing spans, which is quadratic in the worst case. On a few files that is nothing. On a monorepo, count tokens per statement once and add them up instead.
Score both on a labelled query set
A chunker is only better if retrieval gets better, so measure it. You need a corpus every reader already has, and the Python standard library is on your disk. The queries below describe what code does, mostly without using its identifiers, the way you search a codebase you did not write. Each label is a file plus a symbol path.
Before the eval, append one helper to chunkers.py, so the gold line ranges come from the parser instead of from hand-typed line numbers:
# chunkers.py (append)
def symbols(src: str) -> dict[str, tuple[int, int]]:
found = {}
def walk(node, scope):
for child in node.named_children:
defn = unwrap(child)
if defn.type not in DEFS:
continue
name = name_of(defn)
sym = f"{scope}.{name}" if scope else name
found[sym] = (child.start_point[0] + 1,
child.end_point[0] + 1)
if defn.type == "class_definition":
walk(defn.child_by_field_name("body"), sym)
walk(PARSER.parse(src.encode()).root_node, "")
return found
Second file, eval.py. The top holds the model, the files and the labels:
# eval.py
import sysconfig
from pathlib import Path
import numpy as np
from sentence_transformers import SentenceTransformer
from chunkers import ast_chunks, fixed_chunks, symbols
MODEL = "BAAI/bge-small-en-v1.5"
QUERY = "Represent this sentence for searching relevant passages: "
FILES = [
"json/__init__.py", "json/decoder.py", "json/encoder.py",
"shlex.py", "textwrap.py", "collections/__init__.py",
]
LABELS = [
("quote a string so a POSIX shell sees one argument",
"shlex.py", "quote"),
("split a command line into tokens like a shell",
"shlex.py", "split"),
("strip the indentation shared by every line",
"textwrap.py", "dedent"),
("break a paragraph into lines of a maximum width",
"textwrap.py", "wrap"),
("add a prefix to the start of selected lines",
"textwrap.py", "indent"),
("decode JSON at the start of a string, return end",
"json/decoder.py", "JSONDecoder.raw_decode"),
("scan a JSON string literal handling escapes",
"json/decoder.py", "py_scanstring"),
("serialize a Python object to a JSON str",
"json/__init__.py", "dumps"),
("walk nested lists and dicts yielding JSON text",
"json/encoder.py", "_make_iterencode"),
("the n elements with the highest counts",
"collections/__init__.py", "Counter.most_common"),
("move a key to either end of an ordered dict",
"collections/__init__.py", "OrderedDict.move_to_end"),
("make a tuple subclass with fields accessed by name",
"collections/__init__.py", "namedtuple"),
]
recall embeds every chunk once, then checks the top five results for each query in two ways, explained after the code:
# eval.py (continued)
def recall(chunks, gold, model, k=5):
vecs = model.encode([c.embed_text() for c in chunks],
normalize_embeddings=True)
touched = whole = 0
for query, path, (s, e) in gold:
qv = model.encode(QUERY + query, normalize_embeddings=True)
top = np.argsort(-(vecs @ qv))[:k]
hits = [chunks[i] for i in top if chunks[i].path == path]
touched += any(c.start <= e and c.end >= s for c in hits)
whole += any(c.start <= s and c.end >= e for c in hits)
return touched / len(gold), whole / len(gold)
main reads the six files from the standard library of whichever Python runs it and builds both chunk lists:
# eval.py (continued)
def main():
root = Path(sysconfig.get_paths()["stdlib"])
model = SentenceTransformer(MODEL)
def ntok(text: str) -> int:
ids = model.tokenizer.encode(text, add_special_tokens=False)
return len(ids)
fixed, by_ast, spans = [], [], {}
for rel in FILES:
src = (root / rel).read_text(encoding="utf-8")
fixed += fixed_chunks(rel, src, ntok)
by_ast += ast_chunks(rel, src, ntok)
spans[rel] = symbols(src)
gold = [(q, f, spans[f][sym]) for q, f, sym in LABELS]
for name, chunks in (("fixed", fixed), ("ast", by_ast)):
t, w = recall(chunks, gold, model)
print(f"{name:5} chunks={len(chunks):4} "
f"touched@5={t:.2f} whole@5={w:.2f}")
if __name__ == "__main__":
main()
The query prefix is the retrieval instruction from the BGE model card, which says passages need no prefix. Run python eval.py. The tokenizer may log a "sequence length is longer than" warning while counting tokens for whole classes; that is the counter measuring, not the embedder truncating.
Read the two numbers separately. touched@5 asks whether any top-5 chunk overlaps the target function at all. That is lenient, and it is the metric a token-window chunker looks best on, because a chunk holding three lines of the right function counts. whole@5 asks whether one top-5 chunk contains the entire function, which is what the LLM needs in order to answer. The gap between them is the "body without its name" problem, measured.
There are no results printed in this post, on purpose: your Python version changes the stdlib files, and 12 queries is a smoke test, not a benchmark. Run it and read your own numbers. Some labels are longer than the 480-token budget, so the AST chunker splits them too and neither chunker can score whole on them. namedtuple and _make_iterencode are the obvious ones; dumps (its docstring alone is long) and py_scanstring may be over as well, depending on your Python version. Expect that, and look at which chunk comes back instead. For a real decision, write 50 to 100 queries against your own repo, from real questions people asked in code review or on Slack, and keep the file in version control next to the chunker.
Before you point it at your repo
Four changes take this from a stdlib demo to something you would index a monorepo with.
More languages. tree-sitter has grammars for most languages you will meet, but node type names differ per grammar. function_definition is the Python name; check each grammar's node-types.json and give every language its own DEFS set. If you want a maintained version for Python, Java, C# and TypeScript, astchunk already does it.
Parse errors. tree-sitter is built to keep producing useful results even when the source has syntax errors, so a half-written file mid-refactor still produces a tree. py-tree-sitter exposes has_error on every node. Check it on the root, and when a file is badly broken, fall back to fixed_chunks for that file rather than indexing garbage boundaries.
A code-aware embedder. bge-small-en-v1.5 is a general English model with a WordPiece tokenizer that shreds snake_case and camelCase into fragments. It works for the demo. For production, compare a code-trained embedding model on your labelled set using the same eval.py; that is what the set is for.
Keyword search next to vectors. Exact identifiers, error codes and config keys are where dense retrieval is weakest. Index the symbol field in a keyword index too and fuse the results. A query that contains raw_decode should find raw_decode without asking an embedding model to guess.
If this was useful
A retriever can only return what the chunker kept together, and code makes that easy to see. The RAG Pocket Guide covers the stages on either side of this one, retrieval and reranking, plus the chunking patterns for production RAG. If you are building search over a codebase or anything else with structure, it is the short version of the decisions this post only touches.
Top comments (0)