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

When a codebase assistant retrieves the wrong context for a query like "where does the retry delay get computed?", the culprit is often a naive chunker. Standard fixed-size token chunkers slice code every 512 tokens just like prose, cutting right through function definitions and leaving the LLM to guess based on orphaned syntax. To fix this, you need to parse code at natural semantic boundaries using ASTs.

High-performance AI engineering requires robust infrastructure. Whether you are building internal developer tools, IDE integrations, or enterprise-grade code search, managing your LLM and embedding pipelines efficiently is critical. Platforms like n1n.ai provide the unified API aggregation layer developers need to route requests across top-tier models with minimal latency.

The Problem with Fixed-Size Code Chunks

Fixed-window chunkers have no concept of syntax. If a function begins at token 470 of a 512-token window, the first chunk gets the decorator and signature, while the next chunk gets the body without any identifiers or names attached. When a retrieval engine matches queries against this fragmented data, it fails.

Research from Carnegie Mellon and Augment Code confirms this issue in the cAST paper: "Existing line-based chunking heuristics often break semantic structures, splitting functions or merging unrelated code." Their AST-based approach yielded a 4.3-point Recall@5 gain on RepoEval and a 2.67-point Pass@1 gain on SWE-bench.

Implementing an AST Chunker with tree-sitter

To respect code structure, we use tree-sitter to break files at function and class definitions, attaching the file path and enclosing symbol as metadata.

First, install the necessary dependencies including py-tree-sitter and your embedder model:

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

Next, implement chunkers.py to handle AST-based parsing and the baseline fixed-window comparison:

# 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