Code-aware splitting (CodeSplitter)

Worked

This is a worked example — study the fully solved case below and understand why splitting code by sentences breaks identifier grounding.

Goal

Implement a function split_source(file_path, text) that reads a source file and splits it into LlamaIndex nodes using the correct node parser:

For .py, .ts, .tsx files: use a tree-sitter-backed CodeSplitter (function/class-aware).
For all other files (e.g., .md, .txt): use SentenceSplitter.

The splitting must match how the real code (_load_nodes, _code_splitter, _lang_for) processes sources in the repository.

Why does sentence-splitting code break grounding?

A SentenceSplitter splits text at sentence boundaries (periods, newlines). For code, this can:

Split a class definition across chunks (e.g., class Foo: on one chunk, the body on another).
Break long identifiers like NativeD1Saver in the middle (e.g., NativeD in one chunk, 1Saver in another).
Separate a function from its docstring or parameter list.

When the retriever later searches for exact identifiers (e.g., route_for), a sentence-split chunk may contain only part of the identifier, causing the token to be missing from the candidate set. The code-aware CodeSplitter uses tree-sitter to keep functions, classes, and code blocks intact, preserving identifier and context boundaries.

Worked Solution

Read the complete solution below. Notice:

_lang_for maps file suffix to language string or None.
_code_splitter(language) builds a CodeSplitter with an explicit tree-sitter parser.
For code files we try the code splitter; if it raises an exception (e.g., tree-sitter parse hiccup), we fall back to SentenceSplitter — just as in the real _load_nodes.
Nodes are tagged with metadata ({"file": file_path.name}).

python

from pathlib import Path
from typing import List, Optional

from llama_index.core import Document
from llama_index.core.node_parser import SentenceSplitter, CodeSplitter
from tree_sitter import Parser
from tree_sitter_language_pack import get_language

import os
import re

def _lang_for(path: Path) -> Optional[str]:
    """Determine tree-sitter language from file extension."""
    return {".py": "python", ".ts": "typescript", ".tsx": "typescript"}.get(path.suffix)

def _code_splitter(language: str) -> CodeSplitter:
    """Build a CodeSplitter with an explicit tree-sitter parser."""
    from tree_sitter import Parser
    from tree_sitter_language_pack import get_language
    return CodeSplitter(
        language=language,
        parser=Parser(get_language(language)),
        chunk_lines=int(os.environ.get("CODE_CHUNK_LINES", "60")),
        chunk_lines_overlap=int(os.environ.get("CODE_CHUNK_OVERLAP", "12")),
    )

def split_source(file_path: Path, text: str) -> List[Document]:
    """Split source text into LlamaIndex nodes using the appropriate splitter.
    
    Args:
        file_path: Path object (used to detect language and tag metadata).
        text: The full source text to split.
    
    Returns:
        List of LlamaIndex Document (or BaseNode) objects.
    """
    lang = _lang_for(file_path)
    if lang is not None:
        # Code file – try CodeSplitter, fall back to SentenceSplitter on error
        try:
            code_splitter = _code_splitter(lang)
            doc = Document(text=text, metadata={"file": file_path.name})
            nodes = code_splitter.get_nodes_from_documents([doc])
            return nodes
        except Exception:
            # tree-sitter parse failure – degrade gracefully
            pass
    
    # Prose or fallback: use SentenceSplitter with env-configured size and overlap
    splitter = SentenceSplitter(
        chunk_size=int(os.environ.get("CHUNK_SIZE", "512")),
        chunk_overlap=int(os.environ.get("CHUNK_OVERLAP", "64")),
    )
    doc = Document(text=text, metadata={"file": file_path.name})
    return splitter.get_nodes_from_documents([doc])

Your Skeleton

Start from the skeleton below. Fill in the missing logic to produce the same behaviour as the worked solution. The function must use the same environment variables, the same fallback logic, and the same _lang_for helper. You may reuse the imports from the skeleton.

Your code

Sources

roadmap-kg/kg/rerank.py:151-186
roadmap-kg/kg/glossary_llamaindex.py:215-223
roadmap-kg/kg/ground_content.py:218-261