Add AST-aware code chunking for better code understanding

[gabriel-dehan]

Introduces intelligent AST-aware code chunking to LEANN using astchunk, improving code understanding in RAG applications. Instead of arbitrarily splitting code at character boundaries, we now have the option to preserve semantic structure by chunking at function, class, and method boundaries.

This enhancement should make LEANN significantly better for code-related RAG applications while maintaining full backward compatibility.

I am also planning on adding ruby support to astchunk and have opened a pull request there.

What's new?

AST-Aware Code Chunking

• Smart boundary detection: Preserves functions, classes, and methods as complete units
• Multi-language support: Python, Java, C#, TypeScript out of the box (Ruby soon #PR_LINK)
• Graceful fallback: Falls back to traditional chunking for unsupported languages or parsing errors
• Rich metadata: Each chunk includes line numbers, file paths, and AST node information

📱 New application

• Code RAG app (apps/code_rag.py): Specialized for code repositories with optimized parameters
• Enhanced document RAG: New --enable-code-chunking flag for mixed content

🛠️ CLI Integration

• Global --use-ast-chunking flag for the main CLI
• Configurable chunk sizes and overlap for both code and text
• Proper error handling and user feedback

How this helps

Before: Code was chunked arbitrarily, often splitting functions mid-way

# Traditional chunking might split this:
def important_function():
    """This docstring explains the function."""
    # Implementation here...
--- CHUNK BOUNDARY ---
    return result

After: AST chunking preserves semantic boundaries

# AST chunking keeps this together:
def important_function():
    """This docstring explains the function."""
    # Complete implementation
    return result

This should result in a much better code understanding and search quality in RAG applications.

Examples

# Enable AST chunking for mixed content
python -m apps.document_rag --enable-code-chunking --data-dir ./my_project

# Use the specialized code RAG
python -m apps.code_rag --repo-dir ./codebase --query "How does authentication work?"

# Global CLI with AST support
leann build my-index --docs ./src --use-ast-chunking

Files Changed

• apps/chunking_utils.py - New AST-aware chunking utilities
• apps/code_rag.py - Specialized code RAG application
• apps/base_rag_example.py - Enhanced with AST chunking arguments
• apps/document_rag.py - Added code chunking flag
• packages/leann-core/src/leann/cli.py - CLI integration
• pyproject.toml - Added astchunk dependencies
• tests/test_astchunk_integration.py - Comprehensive test suite
• README.md & docs/features.md - Updated documentation

Dependencies

Added optional dependencies for AST parsing:

• astchunk>=0.1.0 - Core AST chunking library
• tree-sitter>=0.20.0 + language parsers - AST parsing backend

PS: Hopefully I didn't break anything during the rebase with main, I haven't tested everything after rebasing but the basic things are still working as expected.

[Copilot]

Pull Request Overview

This pull request introduces AST-aware code chunking to LEANN using the astchunk library, enhancing code understanding in RAG applications by preserving semantic boundaries like functions and classes instead of arbitrary text splits.

Key changes include:

• AST-aware chunking utilities that detect code files and apply semantic chunking
• Enhanced document RAG with --enable-code-chunking flag for mixed content
• New specialized code RAG application optimized for code repositories

Reviewed Changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
tests/test_document_rag.py	Adds test for document RAG with AST chunking enabled
tests/test_astchunk_integration.py	Comprehensive test suite for AST chunking functionality
pyproject.toml	Adds astchunk and tree-sitter dependencies
packages/leann-core/src/leann/cli.py	Integrates AST chunking flags and logic into main CLI
docs/features.md	Documents new AST-aware code chunking feature
apps/document_rag.py	Adds code chunking support to document RAG
apps/code_rag.py	New specialized application for code repository RAG
apps/chunking/utils.py	Core AST chunking utilities and fallback mechanisms
apps/chunking/__init__.py	Package initialization for chunking utilities
apps/base_rag_example.py	Adds AST chunking parameters to base RAG class
README.md	Updates documentation with code chunking examples
ASTCHUNK_INTEGRATION.md	Comprehensive integration documentation

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[yichuan-w]

Thanks for your contribution!! It is a very important feature, we will review and merge asap.

[yichuan-w]

We recently came across astchunk, the pkg you are using and its accompanying paper, too. Have you had a chance to evaluate it or compare it against naive chunking? I’m not entirely convinced about the extent of its performance gains, but it could be worth discussing further.

[andylizf]

Mostly LGTM! Thanks for your hard work @gabriel-dehan ! Left some discussions.

[andylizf]

Suggested change

	from .utils import CODE_EXTENSIONS

[gabriel-dehan]

This wouldn't work as we are not in the same directory no? We've got an init.py in chunking/ so chunking should be correct (or chunking.utils, but not .utils)
I've left it as is

[andylizf]

Suggested change

	if not use_ast:
	else:

[gabriel-dehan]

Actually no because we want to fallback if there is an error (see line right above, we set use_ast = False) so the next if needs to be separated and not an alternative branch of the previous condition!

[andylizf]

Also, in the future, we may want to refactor our document_rag to make it the default implementation of leann cli, which will reduce a lot of redundant code.

[yichuan-w]

https://github.com/yichuan-w/LEANN/blob/da6540decd08bd69b965a85aa99d84dc0e04ba08/docs/CONTRIBUTING.md#setup-pre-commit

@gabriel-dehan it is better to have pre-commit, thanks!

[gabriel-dehan]

We recently came across astchunk, the pkg you are using and its accompanying paper, too. Have you had a chance to evaluate it or compare it against naive chunking? I’m not entirely convinced about the extent of its performance gains, but it could be worth discussing further.

I've been asking myself the same thing. I've got no idea on how to measure that properly so I used Claude Code to generate a quick and dirty evaluation framework on small and big codebases, it does a lot of things and isn't very clean but it clones big repositories like VSCode / Node / Jango and performs complex queries on it amongst other things.

The results are here: gabriel-dehan#2
I am not saying we should trust them 100%, I've skimmed over the evaluation code and it seems coherent overall but there is way too much stuff there and a lot I don't have the time to understand. But I feel it gives an idea, and the results are not that far from what I'd expect (except for the memory and speed difference, I am a bit skeptical about this)

I'll check your other comments and implement the fixes.

[yichuan-w]

Yeah, thanks for the benchmark. I briefly looked at that, and I am confused about the efficiency and memory-saving gain. Maybe I can investigate that later, but at least we can first merge this PR and make it an option, thanks a lot!! Both this PR and the benchmark are great.

Let's fix this PR and merge it first

[gabriel-dehan]

Rebased on main, made the requested changes (for the most part) and installed the pre-commit hook and fixed a bunch of violations. Also commited the uv.lock so the PR seems like a lot more changes than what it actually has 🤭

[yichuan-w]

I guess still some lint error? @gabriel-dehan and thanks for fixing!

[gabriel-dehan]

Yeah I'll try to look at it tonight!

[yichuan-w]

LGTM, let's wait for CI and we will add you to our ack!

[yichuan-w]

Looks good, we can merge first and make it default false, if it works well, we can make it default true!