code-chunk

AST-aware code chunking for semantic search and RAG pipelines.

Uses tree-sitter to split source code at semantic boundaries (functions, classes, methods) rather than arbitrary character limits. Each chunk includes rich context: scope chain, imports, siblings, and entity signatures.

Table of Contents

• Features
• How It Works
• Installation
• Quickstart
• API Reference
• License

Features

• AST-aware: Splits at semantic boundaries, never mid-function
• Rich context: Scope chain, imports, siblings, entity signatures
• Contextualized text: Pre-formatted for embedding models
• Multi-language: TypeScript, JavaScript, Python, Rust, Go, Java
• Batch processing: Process entire codebases with controlled concurrency
• Streaming: Process large files incrementally
• Effect support: First-class Effect integration

How It Works

Traditional text splitters chunk code by character count or line breaks, often cutting functions in half or separating related code. code-chunk takes a different approach:

1. Parse

Source code is parsed into an Abstract Syntax Tree (AST) using tree-sitter. This gives us a structured representation of the code that understands language grammar.

2. Extract

We traverse the AST to extract semantic entities: functions, methods, classes, interfaces, types, and imports. For each entity, we capture:

• Name and type
• Full signature (e.g., async getUser(id: string): Promise<User>)
• Docstring/comments if present
• Byte and line ranges

3. Build Scope Tree

Entities are organized into a hierarchical scope tree that captures nesting relationships. A method inside a class knows its parent; a nested function knows its containing function. This enables us to provide scope context like UserService > getUser.

4. Chunk

Code is split at semantic boundaries while respecting the maxChunkSize limit. The chunker:

• Prefers to keep complete entities together
• Splits oversized entities at logical points (statement boundaries)
• Never cuts mid-expression or mid-statement
• Merges small adjacent chunks to reduce fragmentation

5. Enrich with Context

Each chunk is enriched with contextual metadata:

• Scope chain: Where this code lives (e.g., inside which class/function)
• Entities: What's defined in this chunk
• Siblings: What comes before/after (for continuity)
• Imports: What dependencies are used

This context is formatted into contextualizedText, optimized for embedding models to understand semantic relationships.

Installation

bun add code-chunk
# or
npm install code-chunk

Quickstart

Basic Usage

import { chunk } from 'code-chunk'

const chunks = await chunk('src/user.ts', sourceCode)

for (const c of chunks) {
  console.log(c.text)
  console.log(c.context.scope)    // [{ name: 'UserService', type: 'class' }]
  console.log(c.context.entities) // [{ name: 'getUser', type: 'method', ... }]
}

Using Contextualized Text for Embeddings

Use contextualizedText for better embedding quality in RAG systems:

for (const c of chunks) {
  const embedding = await embed(c.contextualizedText)
  await vectorDB.upsert({
    id: `${filepath}:${c.index}`,
    embedding,
    metadata: { filepath, lines: c.lineRange }
  })
}

The contextualizedText prepends semantic context to the raw code:

# src/services/user.ts
# Scope: UserService
# Defines: async getUser(id: string): Promise<User>
# Uses: Database
# After: constructor

  async getUser(id: string): Promise<User> {
    return this.db.query('SELECT * FROM users WHERE id = ?', [id])
  }

Streaming Large Files

Process chunks incrementally without loading everything into memory:

import { chunkStream } from 'code-chunk'

for await (const c of chunkStream('src/large.ts', code)) {
  await process(c)
}

Reusable Chunker

Create a chunker instance when processing multiple files with the same config:

import { createChunker } from 'code-chunk'

const chunker = createChunker({
  maxChunkSize: 2048,
  contextMode: 'full',
  siblingDetail: 'signatures',
})

for (const file of files) {
  const chunks = await chunker.chunk(file.path, file.content)
}

Batch Processing

Process multiple files concurrently with error handling per file:

import { chunkBatch } from 'code-chunk'

const files = [
  { filepath: 'src/user.ts', code: userCode },
  { filepath: 'src/auth.ts', code: authCode },
  { filepath: 'lib/utils.py', code: utilsCode },
]

const results = await chunkBatch(files, {
  maxChunkSize: 1500,
  concurrency: 10,
  onProgress: (done, total, path, success) => {
    console.log(`[${done}/${total}] ${path}: ${success ? 'ok' : 'failed'}`)
  }
})

for (const result of results) {
  if (result.error) {
    console.error(`Failed: ${result.filepath}`, result.error)
  } else {
    await indexChunks(result.filepath, result.chunks)
  }
}

Stream results as they complete:

import { chunkBatchStream } from 'code-chunk'

for await (const result of chunkBatchStream(files, { concurrency: 5 })) {
  if (result.chunks) {
    await indexChunks(result.filepath, result.chunks)
  }
}

Effect Integration

For Effect-based pipelines:

import { chunkStreamEffect } from 'code-chunk'
import { Effect, Stream } from 'effect'

const program = Stream.runForEach(
  chunkStreamEffect('src/utils.ts', code),
  (chunk) => Effect.log(chunk.text)
)

await Effect.runPromise(program)

API Reference

chunk(filepath, code, options?)

Chunk source code into semantic pieces with context.

Parameters:

• filepath: File path (used for language detection)
• code: Source code string
• options: Optional configuration

Returns: Promise<Chunk[]>

Throws: ChunkingError, UnsupportedLanguageError

chunkStream(filepath, code, options?)

Stream chunks as they're generated. Useful for large files.

Returns: AsyncGenerator<Chunk>

Note: chunk.totalChunks is -1 in streaming mode (unknown upfront).

chunkStreamEffect(filepath, code, options?)

Effect-native streaming API for composable pipelines.

Returns: Stream.Stream<Chunk, ChunkingError | UnsupportedLanguageError>

createChunker(options?)

Create a reusable chunker instance with default options.

Returns: Chunker with chunk(), stream(), chunkBatch(), and chunkBatchStream() methods

chunkBatch(files, options?)

Process multiple files concurrently with per-file error handling.

Parameters:

• files: Array of { filepath, code, options? }
• options: Batch options (extends ChunkOptions with concurrency and onProgress)

Returns: Promise<BatchResult[]> where each result has { filepath, chunks, error }

chunkBatchStream(files, options?)

Stream batch results as files complete processing.

Returns: AsyncGenerator<BatchResult>

chunkBatchEffect(files, options?)

Effect-native batch processing.

Returns: Effect.Effect<BatchResult[], never>

chunkBatchStreamEffect(files, options?)

Effect-native streaming batch processing.

Returns: Stream.Stream<BatchResult, never>

formatChunkWithContext(text, context, overlapText?)

Format chunk text with semantic context prepended. Useful for custom embedding pipelines.

Returns: string

detectLanguage(filepath)

Detect programming language from file extension.

Returns: Language | null

ChunkOptions

| Option | Type | Default | Description | |--------|------|---------|-------------| | maxChunkSize | number | 1500 | Maximum chunk size in bytes | | contextMode | 'none' \| 'minimal' \| 'full' | 'full' | How much context to include | | siblingDetail | 'none' \| 'names' \| 'signatures' | 'signatures' | Level of sibling detail | | filterImports | boolean | false | Filter out import statements | | language | Language | auto | Override language detection | | overlapLines | number | 10 | Lines from previous chunk to include in contextualizedText |

BatchOptions

Extends ChunkOptions with:

| Option | Type | Default | Description | |--------|------|---------|-------------| | concurrency | number | 10 | Maximum files to process concurrently | | onProgress | function | - | Callback (completed, total, filepath, success) => void |

Supported Languages

| Language | Extensions | |----------|------------| | TypeScript | .ts, .tsx, .mts, .cts | | JavaScript | .js, .jsx, .mjs, .cjs | | Python | .py, .pyi | | Rust | .rs | | Go | .go | | Java | .java |

Errors

ChunkingError: Thrown when chunking fails (parsing error, extraction error, etc.)

UnsupportedLanguageError: Thrown when the file extension is not supported

Both errors have a _tag property for Effect-style error handling.

License

MIT