Chunking

After this page your indexing pipeline splits code along syntax boundaries, markdown and text by characters, and any other format with a closure, with chunks sized for your embedder. A chunker takes one node and returns many; you add it with then_chunk.

Source code

ChunkCode needs the tree-sitter feature flag:

cargo add swiftide@0.34 --features tree-sitter

// At most 1500 bytes per chunk, the default
let python = ChunkCode::try_for_language("python")?;

// At most 2000 bytes per chunk; chunks of 200 bytes or less are dropped
let rust = ChunkCode::try_for_language_and_chunk_size(SupportedLanguages::Rust, 200..2000)?
    .with_concurrency(4);

Pipeline::from_loader(FileLoader::new("src").with_extensions(&["rs"]))
    .then_chunk(rust)
    .then_store_with(MemoryStorage::default())
    .run()
    .await?;

Source: snippets/src/indexing/chunking.rs#L17-L28

ChunkCode parses each file with tree-sitter and splits it along syntax boundaries such as functions, impl blocks and classes. It adds top-level items to a chunk. It closes the chunk when the next item would bring it to the maximum or past it. An item that is too large on its own is split along its children.

The size is a ChunkSize, and it counts bytes. A number such as 1500 is a maximum. A range such as 200..2000 adds a minimum: a chunk of 200 bytes or less is dropped. try_for_language uses a maximum of 1500 bytes. The language is a SupportedLanguages value or a name such as "python"; an unknown name returns an error.

with_concurrency sets how many files are chunked at once. Without it the chunker uses the pipeline’s concurrency.

When tree-sitter finds a syntax error, ChunkCode logs a warning and returns the whole file as one chunk, whatever its size.

Several languages

ChunkCode handles one language. The ask-your-codebase showcase indexes several in one pipeline with its own chunker:

/// Splits each file with a tree-sitter `ChunkCode` for the file's language.
///
/// `ChunkCode` is built for a single language; this chunker picks the right one per file from
/// its extension, so one pipeline can index a polyglot codebase.
#[derive(Clone)]
pub struct ChunkByLanguage {
    chunkers: Vec<(SupportedLanguages, ChunkCode)>,
}

impl ChunkByLanguage {
    /// Chunks are between `chunk_size.start` and `chunk_size.end` bytes; smaller ones are dropped.
    pub fn try_new(languages: &SourceLanguages, chunk_size: Range<usize>) -> Result<Self> {
        let chunkers = languages
            .0
            .iter()
            .map(|&language| {
                let chunker =
                    ChunkCode::try_for_language_and_chunk_size(language, chunk_size.clone())?;
                Ok((language, chunker))
            })
            .collect::<Result<_>>()?;

        Ok(Self { chunkers })
    }

    fn chunker_for(&self, path: &Path) -> Option<&ChunkCode> {
        let extension = path.extension()?.to_str()?;

        self.chunkers
            .iter()
            .find(|(language, _)| language.file_extensions().contains(&extension))
            .map(|(_, chunker)| chunker)
    }
}

#[async_trait]
impl ChunkerTransformer for ChunkByLanguage {
    type Input = String;
    type Output = String;

    async fn transform_node(&self, node: TextNode) -> IndexingStream<String> {
        match self.chunker_for(&node.path) {
            Some(chunker) => chunker.transform_node(node).await,
            None => anyhow!("No chunker for {}", node.path.display()).into(),
        }
    }
}

Source: examples/ask-your-codebase/src/languages.rs#L39-L85

The showcase builds it with a range of 10..4096 bytes. The tree-sitter integration explains how it picks a language when extensions overlap.

A file without a matching chunker becomes an error in the stream, and the first error stops run. Load only the extensions you have chunkers for, as the showcase does, or add log_errors() and filter_errors() (Streams and concurrency).

Markdown and text

// Aims for 100 characters, at most 1499; chunks under 100 bytes after trimming are dropped
let markdown = ChunkMarkdown::from_chunk_range(100..1500);

// At most 1000 characters; only blank chunks are dropped
let text = ChunkText::from_max_characters(1000);

Source: snippets/src/indexing/chunking.rs#L41-L45

ChunkMarkdown and ChunkText use the text-splitter crate and count characters, not bytes. ChunkMarkdown prefers to split at headings, then at paragraphs and other blocks. ChunkText splits at blank lines and line breaks, then at sentences and words. Both default to a maximum of 2056 characters.

With from_chunk_range, text-splitter treats the start of the range as the size to aim for and the end, exclusive, as a hard limit. A chunk grows past the start only when that keeps a larger unit, such as a section, in one piece. Chunks that are empty after trimming, or shorter than the start, are dropped. That check uses the byte length of the trimmed text, which equals the character count only for ASCII.

The docs-qa showcase splits documentation pages with ChunkMarkdown::from_chunk_range(120..1500), because shorter chunks were mostly badges and stray headings.

Closures

let paragraphs = |node: TextNode| {
    let chunks: Vec<String> = node
        .chunk
        .split("\n\n")
        .map(str::trim)
        .filter(|paragraph| !paragraph.is_empty())
        .map(str::to_owned)
        .collect();

    // Each chunk keeps the path, metadata and embed mode of `node`
    IndexingStream::iter(node.into_chunks(chunks).map(Ok))
};

Pipeline::from_loader(FileLoader::new("notes").with_extensions(&["txt"]))
    .then_chunk(paragraphs)
    .then_store_with(MemoryStorage::default())
    .run()
    .await?;

Source: snippets/src/indexing/chunking.rs#L53-L70

A closure Fn(TextNode) -> IndexingStream<String> that is Clone + Send + Sync is a chunker. It works on text nodes only. IndexingStream::iter turns an iterator of Result<TextNode> into a stream.

node.into_chunks builds one node per chunk from the input node. Every chunk keeps the path, metadata, vectors and embed mode of that node. For anything more involved, or for other node types, implement ChunkerTransformer on a struct as the showcase does.

Node expansion

// A JSON Lines file holds one record per line
let records = |node: TextNode| {
    let lines: Vec<String> = node
        .chunk
        .lines()
        .filter(|line| !line.trim().is_empty())
        .map(str::to_owned)
        .collect();

    IndexingStream::iter(
        node.into_chunks(lines)
            .enumerate()
            .map(|(index, mut record)| {
                record.metadata.insert("Record", index);
                Ok(record)
            }),
    )
};

Pipeline::from_loader(FileLoader::new("exports").with_extensions(&["jsonl"]))
    .then_expand(records)
    .then_chunk(ChunkText::from_max_characters(1000))
    .then_store_with(MemoryStorage::default())
    .run()
    .await?;

Source: snippets/src/indexing/chunking.rs#L78-L102

then_expand runs the same ChunkerTransformer as then_chunk; only its tracing span is named differently. Use it when one node holds several records, as in this JSON Lines file, and use then_chunk when you cut content to size.

Metadata and parents

Metadata that a node carries before chunking is copied to every chunk. The showcase stores each file’s relative path as metadata before it chunks, so every chunk can cite its file. The same holds for metadata written by a large language model (LLM). Put a transformer such as MetadataQACode before the chunker and every chunk of a file carries the questions for the whole file. Put it after the chunker for metadata per chunk, at one request per chunk. MetadataQACode also needs the tree-sitter feature flag.

Since 0.33 every chunk also carries parent_id, the id its source node had when it reached the first chunker. The node cache marks that id when the file’s chunks finish. Two consequences matter while you chunk:

A step between filter_cached and the chunker that changes the path or the chunk breaks that match, so the file is never cached. Edit the content before filter_cached, as the showcase does.
A file whose chunks all fall under the minimum is never cached and is processed again on every run.

Cache keys explains the match in full.

Gotchas

Size chunks for the embedder. With the default embed mode the metadata is embedded together with the chunk, so leave room for it under the model’s input limit.
Chunks come out of then_chunk in no particular order, and chunks of different files interleave. ChunkCode sets offset on each chunk, a running byte count over the chunks before it in the same file. The text chunkers keep the offset of the input node.
Only ChunkCode sees syntax. Running ChunkText over code cuts through functions at line or sentence boundaries.
ChunkCode cannot split a syntax node that has no children. A single token larger than the maximum, such as a long comment or string, is silently left out of the chunks.

Reference

ChunkCode
ChunkSize
SupportedLanguages
ChunkMarkdown
ChunkText
ChunkerTransformer
Pipeline::then_chunk and Pipeline::then_expand
Pipeline::filter_errors
Node::into_chunks
IndexingStream
The ask-your-codebase and docs-qa showcases, the tree-sitter integration and feature flags

Transformers and metadata