# libragen: First-Class AST-Aware Code Chunking Support
**Version:** 1.0.0
**Date:** 2024-12-20
**Status:** In Progress
## Summary
Integrate the `code-chunk` library from supermemoryai to provide AST-aware, semantic
code chunking as a first-class feature in libragen. This replaces the current
LangChain-based `RecursiveCharacterTextSplitter` for supported code files with
tree-sitter-powered chunking that respects semantic boundaries (functions, classes,
methods) and provides rich context (scope chains, imports, siblings, entity signatures)
for better embedding quality.
## Objectives & Scope
### In Scope
- **Add `code-chunk` as a dependency** to `@libragen/core`
- **Create a new `CodeChunker` class** that wraps `code-chunk` and implements the same
interface pattern as the existing `Chunker`
- **Extend `Chunk` and `ChunkMetadata` types** to include semantic context from code-chunk
(scope, entities, imports, siblings)
- **Update `Builder`** to use `CodeChunker` for supported code files, falling back to the
existing `Chunker` for unsupported files
- **Store semantic context** in the database for retrieval-time enrichment
- **Store `contextualizedText`** in the database alongside raw content for maximum
flexibility
- **Update CLI and MCP** to expose new chunking options (e.g., `--no-ast-chunking`,
`--context-mode`)
- **Add build option** to enable/disable AST-aware chunking (default: enabled for code
files)
- **Update documentation** across all packages
### Out of Scope
- WASM/Cloudflare Workers support (code-chunk supports this, but libragen is
Node.js-focused)
- Effect.js integration (use Promise-based API)
- Streaming chunking API (batch processing is sufficient for build)
### Future Considerations
- **Custom tree-sitter grammar support** beyond what code-chunk provides - this would
allow users to add support for additional languages by providing their own tree-sitter
grammars
## Assumptions & Open Questions
### Assumptions
- `code-chunk` is stable enough for production use (v0.1.11)
- The `contextualizedText` field from code-chunk is suitable for embedding (this is its
intended use)
- Node.js native tree-sitter bindings will work in libragen's target environments
(Node 24+)
- Users will benefit from richer chunk context even if it increases storage slightly
### Resolved Questions
1. **Should AST chunking be opt-in or opt-out?**
- **Answer:** Opt-out (enabled by default for code files)
2. **Should we store both raw `text` and `contextualizedText`?**
- **Answer:** Yes, store both. Use `contextualizedText` for embeddings, store raw
`content` for display. This provides maximum flexibility and highest quality output.
3. **How should we handle files that code-chunk doesn't support?**
- **Answer:** Fall back to existing `Chunker`
4. **Should chunk context (scope, entities, etc.) be stored in the database?**
- **Answer:** Yes, in the metadata JSON column
5. **What about backward compatibility with old libraries?**
- **Answer:** Old libraries won't have semantic context data, but CLI/MCP should
handle both old and new formats gracefully without errors
## Requirements
### Functional
- **FR-1**: Support AST-aware chunking for TypeScript, JavaScript, Python, Rust, Go, and
Java files
- **FR-2**: Fall back to existing text-based chunking for unsupported file types
- **FR-3**: Store semantic context (scope chain, entities, imports, siblings) in chunk
metadata
- **FR-4**: Store `contextualizedText` in the database for embedding and retrieval
- **FR-5**: Use `contextualizedText` for embedding generation to improve retrieval quality
- **FR-6**: Expose `--no-ast-chunking` flag in CLI to disable AST-aware chunking
- **FR-7**: Expose `--context-mode` option (`none`, `minimal`, `full`) for controlling
context richness (default: `full`)
- **FR-8**: Maintain backward compatibility with existing `.libragen` files
- **FR-9**: CLI and MCP must handle libraries with and without semantic context gracefully
### Non-Functional
- **NFR-1 (Performance)**: AST chunking should not significantly increase build time
(tree-sitter is fast)
- **NFR-2 (Storage)**: Chunk metadata storage increase should be reasonable (<30% file
size increase due to storing contextualizedText)
- **NFR-3 (Compatibility)**: Must work with Node.js 24+ on macOS, Linux, and Windows
## Architecture & Design Overview
### Data Flow
```
Source Files
│
▼
┌─────────────────────────────────────────────────────────┐
│ Builder._chunkSource() │
│ ├─ For each file: │
│ │ ├─ Is code file supported by code-chunk? │
│ │ │ ├─ YES → CodeChunker.chunkText() │
│ │ │ │ └─ Returns Chunk[] with context │
│ │ │ │ + embeddingContent │
│ │ │ └─ NO → Chunker.chunkText() (existing) │
│ │ │ └─ Returns Chunk[] (basic metadata) │
│ │ └─ Collect all chunks │
└─────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────┐
│ Builder._generateEmbeddings() │
│ └─ Use chunk.embeddingContent ?? chunk.content │
│ (contextualizedText for code, raw for others) │
└─────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────┐
│ VectorStore.addChunks() │
│ └─ Store: │
│ - chunk.content (raw text for display) │
│ - chunk.embeddingContent (contextualizedText) │
│ - chunk.metadata (includes codeContext) │
└─────────────────────────────────────────────────────────┘
```
### Key Interfaces
```typescript
// Semantic context from code-chunk
interface CodeContext {
scope: EntityInfo[]; // Scope chain (e.g., class > method)
entities: ChunkEntityInfo[]; // Entities defined in this chunk
siblings: SiblingInfo[]; // Nearby entities for context
imports: ImportInfo[]; // Relevant imports
}
// Extended ChunkMetadata
interface ChunkMetadata {
sourceFile: string;
startLine?: number;
endLine?: number;
language?: string;
// NEW: Semantic context from code-chunk
codeContext?: CodeContext;
}
// Extended Chunk
interface Chunk {
content: string; // Raw code text (for display)
embeddingContent?: string; // contextualizedText (for embedding)
metadata: ChunkMetadata;
}
```
### Database Schema Changes
The `chunks` table already has a `metadata` JSON column. We will:
1. Store `codeContext` in the metadata JSON
2. Add a new column `embedding_content` to store `contextualizedText` separately from
`content`
This allows:
- Displaying raw code to users
- Using enriched context for embeddings
- Retrieving semantic context at search time
### Decisions & Trade-offs
- **Decision**: Store both `content` and `embeddingContent` (contextualizedText)
- **Rationale**: Maximum flexibility - raw content for display, contextualized for
embeddings and potential re-embedding
- **Trade-off**: ~20-30% storage increase, but worth it for quality
- **Decision**: Store semantic context in metadata JSON column
- **Rationale**: Enables rich retrieval-time features without schema migration
- **Trade-off**: Increases storage, but provides valuable context
- **Decision**: Default to AST chunking with `contextMode: 'full'`
- **Rationale**: Best quality out of the box
- **Trade-off**: Users who want smaller files can opt out
- **Decision**: Graceful fallback for unsupported files and old libraries
- **Rationale**: Seamless experience regardless of file type or library age
- **Trade-off**: Some code paths need to handle both cases
## Task Grid
| Status | ID | Task | Priority | Depends On | Acceptance Criteria |
|---|---|---|---|---|---|
| [ ] | T-01 | Add `code-chunk` dependency | H | — | Package installed, types available |
| [ ] | T-02 | Extend `Chunk` and `ChunkMetadata` types | H | T-01 | Types include semantic context fields |
| [ ] | T-03 | Create `CodeChunker` class | H | T-02 | Class wraps code-chunk, matches Chunker pattern |
| [ ] | T-04 | Update `Builder` to use `CodeChunker` | H | T-03 | Builder uses AST chunking for supported files |
| [ ] | T-05 | Update `VectorStore` schema | M | T-02 | `embedding_content` column added |
| [ ] | T-06 | Add build options for AST chunking | M | T-04 | `noAstChunking`, `contextMode` options work |
| [ ] | T-07 | Update CLI with new options | M | T-06 | `--no-ast-chunking`, `--context-mode` flags |
| [ ] | T-08 | Update MCP tools with new options | M | T-06 | MCP build tool accepts new options |
| [ ] | T-09 | Write unit tests | H | T-03, T-04 | ≥90% coverage for new code |
| [ ] | T-10 | Update documentation | M | T-07, T-08 | READMEs, website docs updated |
## Task Details
### T-01 — Add `code-chunk` dependency
**Goal:** Install `code-chunk` package in `@libragen/core`.
**Step-by-step instructions:**
1. Run `npm install --save-exact code-chunk` in `packages/core`
2. Verify TypeScript types are available
3. Test basic import: `import { chunk, detectLanguage } from 'code-chunk'`
---
### T-02 — Extend `Chunk` and `ChunkMetadata` types
**Goal:** Add semantic context fields to chunk types.
**Step-by-step instructions:**
1. Update `packages/core/src/chunker.ts`:
- Add `CodeContext` interface with scope, entities, siblings, imports
- Add `codeContext?: CodeContext` to `ChunkMetadata`
- Add `embeddingContent?: string` to `Chunk` interface
2. Export new types from `packages/core/src/index.ts`
---
### T-03 — Create `CodeChunker` class
**Goal:** Create a wrapper around `code-chunk` that matches the `Chunker` interface
pattern.
**Step-by-step instructions:**
1. Create `packages/core/src/code-chunker.ts`
2. Implement `CodeChunker` class with:
- `static isSupported(filePath: string): boolean` - check if file is supported
- `static detectLanguage(filePath: string): Language | null`
- `async chunkText(content: string, filePath: string): Promise<Chunk[]>`
- `async chunkFile(filePath: string): Promise<Chunk[]>`
- `async chunkSourceFiles(files: SourceFile[]): Promise<Chunk[]>`
3. Map code-chunk's `Chunk` type to libragen's `Chunk` type
4. Handle errors gracefully (fall back to null/empty on parse errors)
---
### T-04 — Update `Builder` to use `CodeChunker`
**Goal:** Integrate `CodeChunker` into the build pipeline.
**Step-by-step instructions:**
1. Update `packages/core/src/builder.ts`:
- Import `CodeChunker`
- Add `noAstChunking?: boolean` and `contextMode?: 'none' | 'minimal' | 'full'` to
`BuildOptions`
- Modify `_chunkSource()` to:
- Use `CodeChunker` for supported files (unless `noAstChunking` is true)
- Fall back to `Chunker` for unsupported files
- Modify `_generateEmbeddings()` to use `chunk.embeddingContent ?? chunk.content`
2. Update `chunking.strategy` in metadata to indicate AST chunking was used
---
### T-05 — Update `VectorStore` schema
**Goal:** Add `embedding_content` column to store contextualizedText.
**Step-by-step instructions:**
1. Update `packages/core/src/store.ts`:
- Add `embedding_content TEXT` column to chunks table
- Update `addChunk()` and `addChunks()` to store `embeddingContent`
- Update `StoredChunk` type to include `embeddingContent`
- Update retrieval methods to return `embeddingContent`
2. Handle backward compatibility: if `embedding_content` is NULL, fall back to `content`
---
### T-06 — Add build options for AST chunking
**Goal:** Allow users to control AST chunking behavior.
**Step-by-step instructions:**
1. Add to `BuildOptions` in `packages/core/src/builder.ts`:
```typescript
/** Disable AST-aware chunking for code files (default: false) */
noAstChunking?: boolean;
/** Context mode for AST chunking (default: 'full') */
contextMode?: 'none' | 'minimal' | 'full';
```
2. Pass options through to `CodeChunker`
---
### T-07 — Update CLI with new options
**Goal:** Expose AST chunking options in the CLI.
**Step-by-step instructions:**
1. Update `packages/cli/src/commands/build.ts`:
- Add `--no-ast-chunking` flag
- Add `--context-mode <mode>` option with choices: `none`, `minimal`, `full`
2. Pass options to `Builder.build()`
3. Update CLI help text
---
### T-08 — Update MCP tools with new options
**Goal:** Expose AST chunking options in MCP build tool.
**Step-by-step instructions:**
1. Update `packages/mcp/src/tools/build.ts`:
- Add `noAstChunking` and `contextMode` to tool parameters
2. Update `packages/mcp/src/tasks/build-worker.ts` to pass options
---
### T-09 — Write unit tests
**Goal:** Ensure new functionality is well-tested.
**Step-by-step instructions:**
1. Create `packages/core/src/__tests__/code-chunker.test.ts`:
- Test `isSupported()` for various file extensions
- Test `chunkText()` with TypeScript, Python, etc.
- Test fallback behavior for unsupported files
- Test error handling for malformed code
2. Update `packages/core/src/__tests__/builder.test.ts`:
- Test build with AST chunking enabled
- Test build with `noAstChunking: true`
- Test `contextMode` options
- Test mixed file types (code + markdown)
---
### T-10 — Update documentation
**Goal:** Document new features for users.
**Step-by-step instructions:**
1. Update `packages/core/README.md`:
- Document `CodeChunker` class
- Document new `BuildOptions`
2. Update `packages/cli/README.md`:
- Document `--no-ast-chunking` and `--context-mode` flags
3. Update `packages/mcp/README.md`:
- Document new build tool parameters
4. Update `packages/website/src/content/docs/building.md`:
- Add section on AST-aware chunking
- Explain benefits and when to use/disable
## New Code
### `packages/core/src/code-chunker.ts` (new file)
Creates a `CodeChunker` class that:
- Wraps the `code-chunk` library
- Provides `isSupported()`, `detectLanguage()`, `chunkText()`, `chunkFile()`,
`chunkSourceFiles()` methods
- Maps code-chunk's `Chunk` type to libragen's `Chunk` type
- Handles errors gracefully
### `packages/core/src/chunker.ts` (modifications)
- Add `CodeContext` interface
- Add `codeContext?: CodeContext` to `ChunkMetadata`
- Add `embeddingContent?: string` to `Chunk`
### `packages/core/src/builder.ts` (modifications)
- Import `CodeChunker`
- Add `noAstChunking`, `contextMode` to `BuildOptions`
- Update `_chunkSource()` to use `CodeChunker` for supported files
- Update `_generateEmbeddings()` to use `embeddingContent`
- Update metadata to reflect chunking strategy
### `packages/core/src/store.ts` (modifications)
- Add `embedding_content` column
- Update insert/retrieve methods
- Handle backward compatibility
### `packages/cli/src/commands/build.ts` (modifications)
- Add `--no-ast-chunking` flag
- Add `--context-mode` option
### `packages/mcp/src/tools/build.ts` (modifications)
- Add `noAstChunking`, `contextMode` parameters
## Tests
1. **`code-chunker.test.ts`**:
- `isSupported()` returns true for `.ts`, `.tsx`, `.js`, `.jsx`, `.py`, `.rs`, `.go`,
`.java`
- `isSupported()` returns false for `.md`, `.json`, `.txt`, etc.
- `chunkText()` produces chunks with semantic context for TypeScript code
- `chunkText()` handles parse errors gracefully
- `chunkSourceFiles()` processes multiple files correctly
2. **`builder.test.ts` (additions)**:
- Build with AST chunking produces chunks with `codeContext`
- Build with `noAstChunking: true` produces chunks without `codeContext`
- `contextMode` option is respected
- Mixed file types (code + markdown) are handled correctly
3. **Integration tests**:
- Build a library from a TypeScript project
- Verify chunks have semantic context
- Verify search quality improvement (manual verification)
## Review Checklist
- [x] Have all outstanding questions been answered?
- [x] Are there any ambiguities that need to be resolved?
