chore: import upstream snapshot with attribution
Publish CLI Package / publish-npm (push) Waiting to run
Publish Python SDK / publish-pypi (push) Waiting to run
Publish TypeScript SDK / publish-npm (push) Waiting to run
CI / Migrate Dev DB (push) Has been skipped
CI / Detect Version (push) Has been cancelled
CI / Migrate DB (push) Has been cancelled
CI / Build Dev ECR (./docker/app.Dockerfile, ECR_APP) (push) Has been cancelled
CI / Build Dev ECR (./docker/db.Dockerfile, ECR_MIGRATIONS) (push) Has been cancelled
CI / Build Dev ECR (./docker/pii.Dockerfile, ECR_PII) (push) Has been cancelled
CI / Build Dev ECR (./docker/realtime.Dockerfile, ECR_REALTIME) (push) Has been cancelled
CI / Deploy Trigger.dev (Dev) (push) Has been cancelled
CI / Build AMD64 (./docker/app.Dockerfile, ECR_APP, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build AMD64 (./docker/db.Dockerfile, ECR_MIGRATIONS, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build AMD64 (./docker/pii.Dockerfile, ECR_PII, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build AMD64 (./docker/realtime.Dockerfile, ECR_REALTIME, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/app.Dockerfile, ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/db.Dockerfile, ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/pii.Dockerfile, ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Build ARM64 (GHCR Only) (./docker/realtime.Dockerfile, ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/migrations) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/pii) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/realtime) (push) Has been cancelled
CI / Create GHCR Manifests (ghcr.io/simstudioai/simstudio) (push) Has been cancelled
CI / Check Docs Changes (push) Has been cancelled
CI / Process Docs (push) Has been cancelled
CI / Create GitHub Release (push) Has been cancelled
CI / Test and Build (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:20:55 +08:00
commit d25d482dc2
13754 changed files with 4996608 additions and 0 deletions
+387
View File
@@ -0,0 +1,387 @@
import fs from 'fs/promises'
import path from 'path'
import { createLogger } from '@sim/logger'
import { TextChunker } from '@/lib/chunkers/text-chunker'
import type { DocChunk, DocsChunkerOptions } from '@/lib/chunkers/types'
import { estimateTokens } from '@/lib/chunkers/utils'
import { generateEmbeddings, getConfiguredEmbeddingModel } from '@/lib/knowledge/embeddings'
interface HeaderInfo {
level: number
text: string
anchor?: string
position?: number
}
interface Frontmatter {
title?: string
description?: string
[key: string]: unknown
}
const logger = createLogger('DocsChunker')
export class DocsChunker {
private readonly textChunker: TextChunker
private readonly baseUrl: string
private readonly chunkSize: number
constructor(options: DocsChunkerOptions = {}) {
this.chunkSize = options.chunkSize ?? 300
this.textChunker = new TextChunker({
chunkSize: this.chunkSize,
minCharactersPerChunk: options.minCharactersPerChunk ?? 1,
chunkOverlap: options.chunkOverlap ?? 50,
})
this.baseUrl = options.baseUrl ?? 'https://docs.sim.ai'
}
async chunkAllDocs(docsPath: string): Promise<DocChunk[]> {
const allChunks: DocChunk[] = []
try {
const mdxFiles = await this.findMdxFiles(docsPath)
logger.info(`Found ${mdxFiles.length} .mdx files to process`)
for (const filePath of mdxFiles) {
try {
const chunks = await this.chunkMdxFile(filePath, docsPath)
allChunks.push(...chunks)
logger.info(`Processed ${filePath}: ${chunks.length} chunks`)
} catch (error) {
logger.error(`Error processing ${filePath}:`, error)
}
}
logger.info(`Total chunks generated: ${allChunks.length}`)
return allChunks
} catch (error) {
logger.error('Error processing docs:', error)
throw error
}
}
async chunkMdxFile(filePath: string, basePath: string): Promise<DocChunk[]> {
const content = await fs.readFile(filePath, 'utf-8')
const relativePath = path.relative(basePath, filePath)
const { data: frontmatter, content: markdownContent } = this.parseFrontmatter(content)
const documentUrl = this.generateDocumentUrl(relativePath)
const { chunks: textChunks, cleanedContent } = await this.splitContent(markdownContent)
const headers = this.extractHeaders(cleanedContent)
logger.info(`Generating embeddings for ${textChunks.length} chunks in ${relativePath}`)
const embeddingModel = getConfiguredEmbeddingModel()
const embeddings: number[][] =
textChunks.length > 0 ? (await generateEmbeddings(textChunks, embeddingModel)).embeddings : []
const chunks: DocChunk[] = []
let currentPosition = 0
for (let i = 0; i < textChunks.length; i++) {
const chunkText = textChunks[i]
const chunkStart = currentPosition
const chunkEnd = currentPosition + chunkText.length
const relevantHeader = this.findRelevantHeader(headers, chunkStart)
const chunk: DocChunk = {
text: chunkText,
tokenCount: estimateTokens(chunkText),
sourceDocument: relativePath,
headerLink: relevantHeader ? `${documentUrl}#${relevantHeader.anchor}` : documentUrl,
headerText: relevantHeader?.text || frontmatter.title || 'Document Root',
headerLevel: relevantHeader?.level || 1,
embedding: embeddings[i] || [],
embeddingModel,
metadata: {
startIndex: chunkStart,
endIndex: chunkEnd,
title: frontmatter.title,
},
}
chunks.push(chunk)
currentPosition = chunkEnd
}
return chunks
}
private async findMdxFiles(dirPath: string): Promise<string[]> {
const files: string[] = []
const entries = await fs.readdir(dirPath, { withFileTypes: true })
for (const entry of entries) {
const fullPath = path.join(dirPath, entry.name)
if (entry.isDirectory()) {
const subFiles = await this.findMdxFiles(fullPath)
files.push(...subFiles)
} else if (entry.isFile() && entry.name.endsWith('.mdx')) {
files.push(fullPath)
}
}
return files
}
private extractHeaders(content: string): HeaderInfo[] {
const headers: HeaderInfo[] = []
const headerRegex = /^(#{1,6})\s+(.+)$/gm
let match
while ((match = headerRegex.exec(content)) !== null) {
const level = match[1].length
const text = match[2].trim()
const anchor = this.generateAnchor(text)
headers.push({
text,
level,
anchor,
position: match.index,
})
}
return headers
}
private generateAnchor(headerText: string): string {
return headerText
.toLowerCase()
.replace(/[^\w\s-]/g, '')
.replace(/\s+/g, '-')
.replace(/-+/g, '-')
.replace(/^-|-$/g, '')
}
/** index.mdx files are served at the parent directory path */
private generateDocumentUrl(relativePath: string): string {
let urlPath = relativePath.replace(/\.mdx$/, '').replace(/\\/g, '/')
if (urlPath.endsWith('/index')) {
urlPath = urlPath.slice(0, -6)
} else if (urlPath === 'index') {
urlPath = ''
}
return `${this.baseUrl}/${urlPath}`
}
private findRelevantHeader(headers: HeaderInfo[], position: number): HeaderInfo | null {
if (headers.length === 0) return null
let relevantHeader: HeaderInfo | null = null
for (const header of headers) {
if (header.position !== undefined && header.position <= position) {
relevantHeader = header
} else {
break
}
}
return relevantHeader
}
/** Returns both chunks and cleaned content so header extraction uses aligned positions. */
private async splitContent(
content: string
): Promise<{ chunks: string[]; cleanedContent: string }> {
const cleanedContent = this.cleanContent(content)
const tableBoundaries = this.detectTableBoundaries(cleanedContent)
const chunks = await this.textChunker.chunk(cleanedContent)
const processedChunks = this.mergeTableChunks(
chunks.map((chunk) => chunk.text),
tableBoundaries,
cleanedContent
)
const finalChunks = this.enforceSizeLimit(processedChunks)
return { chunks: finalChunks, cleanedContent }
}
private cleanContent(content: string): string {
return content
.replace(/\r\n/g, '\n')
.replace(/\r/g, '\n')
.replace(/^import\s+.*$/gm, '')
.replace(/^export\s+.*$/gm, '')
.replace(/<\/?[a-zA-Z][^>]*>/g, ' ')
.replace(/\{\/\*[\s\S]*?\*\/\}/g, ' ')
.replace(/\{[^{}]*\}/g, ' ')
.replace(/\n{3,}/g, '\n\n')
.replace(/[ \t]{2,}/g, ' ')
.trim()
}
private parseFrontmatter(content: string): { data: Frontmatter; content: string } {
const frontmatterRegex = /^---\r?\n([\s\S]*?)\r?\n---\r?\n([\s\S]*)$/
const match = content.match(frontmatterRegex)
if (!match) {
return { data: {}, content }
}
const [, frontmatterText, markdownContent] = match
const data: Frontmatter = {}
const lines = frontmatterText.split('\n')
for (const line of lines) {
const colonIndex = line.indexOf(':')
if (colonIndex > 0) {
const key = line.slice(0, colonIndex).trim()
const value = line
.slice(colonIndex + 1)
.trim()
.replace(/^['"]|['"]$/g, '')
data[key] = value
}
}
return { data, content: markdownContent }
}
/** Detects table boundaries to avoid splitting tables across chunks. */
private detectTableBoundaries(content: string): { start: number; end: number }[] {
const tables: { start: number; end: number }[] = []
const lines = content.split('\n')
let inTable = false
let inCodeBlock = false
let tableStart = -1
for (let i = 0; i < lines.length; i++) {
const line = lines[i].trim()
if (line.startsWith('```')) {
inCodeBlock = !inCodeBlock
continue
}
if (inCodeBlock) continue
if (line.includes('|') && line.split('|').length >= 3 && !inTable) {
const nextLine = lines[i + 1]?.trim()
if (nextLine?.includes('|') && nextLine.includes('-')) {
inTable = true
tableStart = i
}
} else if (inTable && (!line.includes('|') || line === '' || line.startsWith('#'))) {
tables.push({
start: this.getCharacterPosition(lines, tableStart),
end: this.getCharacterPosition(lines, i - 1) + (lines[i - 1]?.length ?? 0),
})
inTable = false
}
}
if (inTable && tableStart >= 0) {
tables.push({
start: this.getCharacterPosition(lines, tableStart),
end: content.length,
})
}
return tables
}
private getCharacterPosition(lines: string[], lineIndex: number): number {
return lines.slice(0, lineIndex).reduce((acc, line) => acc + line.length + 1, 0)
}
private mergeTableChunks(
chunks: string[],
tableBoundaries: { start: number; end: number }[],
originalContent: string
): string[] {
if (tableBoundaries.length === 0) {
return chunks
}
const mergedChunks: string[] = []
let currentPosition = 0
for (const chunk of chunks) {
const chunkStart = originalContent.indexOf(chunk, currentPosition)
if (chunkStart === -1) {
mergedChunks.push(chunk)
continue
}
const chunkEnd = chunkStart + chunk.length
const intersectsTable = tableBoundaries.some(
(table) =>
(chunkStart >= table.start && chunkStart <= table.end) ||
(chunkEnd >= table.start && chunkEnd <= table.end) ||
(chunkStart <= table.start && chunkEnd >= table.end)
)
if (intersectsTable) {
const affectedTables = tableBoundaries.filter(
(table) =>
(chunkStart >= table.start && chunkStart <= table.end) ||
(chunkEnd >= table.start && chunkEnd <= table.end) ||
(chunkStart <= table.start && chunkEnd >= table.end)
)
const minStart = Math.min(chunkStart, ...affectedTables.map((t) => t.start))
const maxEnd = Math.max(chunkEnd, ...affectedTables.map((t) => t.end))
const completeChunk = originalContent.slice(minStart, maxEnd).trim()
if (completeChunk && !mergedChunks.some((existing) => existing === completeChunk)) {
mergedChunks.push(completeChunk)
}
} else {
mergedChunks.push(chunk)
}
currentPosition = chunkEnd
}
return mergedChunks.filter((chunk) => chunk.length > 50)
}
private enforceSizeLimit(chunks: string[]): string[] {
const finalChunks: string[] = []
for (const chunk of chunks) {
const tokens = estimateTokens(chunk)
if (tokens <= this.chunkSize) {
finalChunks.push(chunk)
} else {
const lines = chunk.split('\n')
let currentChunk = ''
for (const line of lines) {
const testChunk = currentChunk ? `${currentChunk}\n${line}` : line
if (estimateTokens(testChunk) <= this.chunkSize) {
currentChunk = testChunk
} else {
if (currentChunk.trim()) {
finalChunks.push(currentChunk.trim())
}
currentChunk = line
}
}
if (currentChunk.trim()) {
finalChunks.push(currentChunk.trim())
}
}
}
return finalChunks.filter((chunk) => chunk.trim().length > 100)
}
}
+9
View File
@@ -0,0 +1,9 @@
export { DocsChunker } from './docs-chunker'
export { JsonYamlChunker } from './json-yaml-chunker'
export { RecursiveChunker } from './recursive-chunker'
export { RegexChunker } from './regex-chunker'
export { SentenceChunker } from './sentence-chunker'
export { StructuredDataChunker } from './structured-data-chunker'
export { TextChunker } from './text-chunker'
export { TokenChunker } from './token-chunker'
export * from './types'
@@ -0,0 +1,375 @@
/**
* @vitest-environment node
*/
import { describe, expect, it, vi } from 'vitest'
import { JsonYamlChunker } from './json-yaml-chunker'
vi.mock('@/lib/tokenization', () => ({
getAccurateTokenCount: (text: string) => Math.ceil(text.length / 4),
}))
vi.mock('@/lib/tokenization/estimators', () => ({
estimateTokenCount: (text: string) => ({ count: Math.ceil(text.length / 4) }),
}))
describe('JsonYamlChunker', () => {
describe('isStructuredData', () => {
it('should detect valid JSON', () => {
expect(JsonYamlChunker.isStructuredData('{"key": "value"}')).toBe(true)
})
it('should detect valid JSON array', () => {
expect(JsonYamlChunker.isStructuredData('[1, 2, 3]')).toBe(true)
})
it('should detect valid YAML', () => {
expect(JsonYamlChunker.isStructuredData('key: value\nother: data')).toBe(true)
})
it('should return false for plain text parsed as YAML scalar', () => {
expect(JsonYamlChunker.isStructuredData('Hello, this is plain text.')).toBe(false)
})
it('should return false for invalid JSON/YAML with unbalanced braces', () => {
expect(JsonYamlChunker.isStructuredData('{invalid: json: content: {{')).toBe(false)
})
it('should detect nested JSON objects', () => {
const nested = JSON.stringify({ level1: { level2: { level3: 'value' } } })
expect(JsonYamlChunker.isStructuredData(nested)).toBe(true)
})
})
describe('basic chunking', () => {
it.concurrent('should return single chunk for small JSON', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 1000 })
const json = JSON.stringify({ name: 'test', value: 123 })
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should return empty array for empty object', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = '{}'
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThanOrEqual(0)
})
it.concurrent('should chunk large JSON object', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 50 })
const largeObject: Record<string, string> = {}
for (let i = 0; i < 100; i++) {
largeObject[`key${i}`] = `value${i}`.repeat(10)
}
const json = JSON.stringify(largeObject)
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should chunk large JSON array', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 50 })
const largeArray = Array.from({ length: 100 }, (_, i) => ({
id: i,
name: `Item ${i}`,
description: 'A description that takes some space',
}))
const json = JSON.stringify(largeArray)
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should include token count in chunk metadata', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 1000 })
const json = JSON.stringify({ hello: 'world' })
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].tokenCount).toBeGreaterThan(0)
})
})
describe('YAML chunking', () => {
it.concurrent('should chunk valid YAML', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const yaml = `
name: test
version: 1.0.0
config:
debug: true
port: 8080
`.trim()
const chunks = await chunker.chunk(yaml)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle YAML with arrays', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const yaml = `
items:
- name: first
value: 1
- name: second
value: 2
- name: third
value: 3
`.trim()
const chunks = await chunker.chunk(yaml)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle YAML with nested structures', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 50 })
const yaml = `
database:
host: localhost
port: 5432
credentials:
username: admin
password: secret
server:
host: 0.0.0.0
port: 3000
`.trim()
const chunks = await chunker.chunk(yaml)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('structured data handling', () => {
it.concurrent('should preserve context path for nested objects', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 30 })
const data = {
users: [
{ id: 1, name: 'Alice', email: 'alice@example.com' },
{ id: 2, name: 'Bob', email: 'bob@example.com' },
],
}
const json = JSON.stringify(data)
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle deeply nested structures', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 50 })
const deepObject = {
l1: {
l2: {
l3: {
l4: {
l5: 'deep value',
},
},
},
},
}
const json = JSON.stringify(deepObject)
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle mixed arrays and objects', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const mixed = {
settings: { theme: 'dark', language: 'en' },
items: [1, 2, 3],
users: [{ name: 'Alice' }, { name: 'Bob' }],
}
const json = JSON.stringify(mixed)
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('edge cases', () => {
it.concurrent('should handle empty array', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = '[]'
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThanOrEqual(0)
})
it.concurrent('should handle JSON with unicode keys and values', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = JSON.stringify({
: '田中太郎',
: '東京都渋谷区',
})
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('名前')
})
it.concurrent('should handle JSON with special characters in strings', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = JSON.stringify({
text: 'Line 1\nLine 2\tTabbed',
special: '!@#$%^&*()',
quotes: '"double" and \'single\'',
})
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle JSON with null values', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = JSON.stringify({
valid: 'value',
empty: null,
another: 'value',
})
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('null')
})
it.concurrent('should handle JSON with boolean values', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = JSON.stringify({
active: true,
deleted: false,
name: 'test',
})
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle JSON with numeric values', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = JSON.stringify({
integer: 42,
float: Math.PI,
negative: -100,
scientific: 1.5e10,
})
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should fall back to text chunking for invalid JSON', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100, minCharactersPerChunk: 10 })
const invalidJson = `{this is not valid json: content: {{${' more content here '.repeat(10)}`
const chunks = await chunker.chunk(invalidJson)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('large inputs', () => {
it.concurrent('should handle JSON with 1000 array items', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 200 })
const largeArray = Array.from({ length: 1000 }, (_, i) => ({
id: i,
name: `Item ${i}`,
}))
const json = JSON.stringify(largeArray)
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle JSON with long string values', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = JSON.stringify({
content: 'A'.repeat(5000),
description: 'B'.repeat(3000),
})
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle deeply nested structure up to depth limit', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 50 })
let nested: Record<string, unknown> = { value: 'deep' }
for (let i = 0; i < 10; i++) {
nested = { [`level${i}`]: nested }
}
const json = JSON.stringify(nested)
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('static chunkJsonYaml method', () => {
it.concurrent('should work with default options', async () => {
const json = JSON.stringify({ test: 'value' })
const chunks = await JsonYamlChunker.chunkJsonYaml(json)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should accept custom options', async () => {
const largeObject: Record<string, string> = {}
for (let i = 0; i < 50; i++) {
largeObject[`key${i}`] = `value${i}`.repeat(20)
}
const json = JSON.stringify(largeObject)
const chunksSmall = await JsonYamlChunker.chunkJsonYaml(json, { chunkSize: 50 })
const chunksLarge = await JsonYamlChunker.chunkJsonYaml(json, { chunkSize: 500 })
expect(chunksSmall.length).toBeGreaterThan(chunksLarge.length)
})
})
describe('chunk metadata', () => {
it.concurrent('should include startIndex and endIndex in metadata', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100 })
const json = JSON.stringify({ key: 'value' })
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].metadata.startIndex).toBeDefined()
expect(chunks[0].metadata.endIndex).toBeDefined()
})
it.concurrent('should have valid metadata indices for array chunking', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 50 })
const largeArray = Array.from({ length: 50 }, (_, i) => ({ id: i, data: 'x'.repeat(20) }))
const json = JSON.stringify(largeArray)
const chunks = await chunker.chunk(json)
for (const chunk of chunks) {
expect(chunk.metadata.startIndex).toBeDefined()
expect(chunk.metadata.endIndex).toBeDefined()
}
})
})
describe('constructor options', () => {
it.concurrent('should use default chunkSize when not provided', async () => {
const chunker = new JsonYamlChunker({})
const json = JSON.stringify({ test: 'value' })
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should respect custom minCharactersPerChunk', async () => {
const chunker = new JsonYamlChunker({ chunkSize: 100, minCharactersPerChunk: 20 })
const json = JSON.stringify({ a: 1, b: 2, c: 3 })
const chunks = await chunker.chunk(json)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text.length).toBeGreaterThan(0)
})
})
})
+278
View File
@@ -0,0 +1,278 @@
import { createLogger } from '@sim/logger'
import * as yaml from 'js-yaml'
import type { Chunk, ChunkerOptions } from '@/lib/chunkers/types'
import { estimateTokens } from '@/lib/chunkers/utils'
const logger = createLogger('JsonYamlChunker')
type JsonPrimitive = string | number | boolean | null
type JsonValue = JsonPrimitive | JsonObject | JsonArray
type JsonObject = { [key: string]: JsonValue }
type JsonArray = JsonValue[]
const MAX_DEPTH = 5
export class JsonYamlChunker {
private chunkSize: number
private minCharactersPerChunk: number
constructor(options: ChunkerOptions = {}) {
this.chunkSize = options.chunkSize ?? 1024
this.minCharactersPerChunk = options.minCharactersPerChunk ?? 100
}
static isStructuredData(content: string): boolean {
try {
const parsed = JSON.parse(content)
return typeof parsed === 'object' && parsed !== null
} catch {
try {
const parsed = yaml.load(content)
return typeof parsed === 'object' && parsed !== null
} catch {
return false
}
}
}
async chunk(content: string): Promise<Chunk[]> {
try {
let data: JsonValue
try {
data = JSON.parse(content) as JsonValue
} catch {
data = yaml.load(content) as JsonValue
}
const chunks = this.chunkStructuredData(data, [], 0)
const totalTokens = chunks.reduce((sum, c) => sum + c.tokenCount, 0)
logger.info(`JSON chunking complete: ${chunks.length} chunks, ${totalTokens} total tokens`)
return chunks
} catch (error) {
logger.info('JSON parsing failed, falling back to text chunking')
return this.chunkAsText(content)
}
}
private chunkStructuredData(data: JsonValue, path: string[], depth: number): Chunk[] {
if (Array.isArray(data)) {
return this.chunkArray(data, path, depth)
}
if (typeof data === 'object' && data !== null) {
return this.chunkObject(data as JsonObject, path, depth)
}
const content = JSON.stringify(data, null, 2)
const contextHeader = path.length > 0 ? `// ${path.join('.')}\n` : ''
const contentTokens = estimateTokens(content)
if (contentTokens > this.chunkSize) {
return this.chunkAsText(contextHeader + content)
}
if (content.length < this.minCharactersPerChunk) {
return []
}
const text = contextHeader + content
return [
{
text,
tokenCount: estimateTokens(text),
metadata: { startIndex: 0, endIndex: text.length },
},
]
}
private chunkArray(arr: JsonArray, path: string[], depth: number): Chunk[] {
const chunks: Chunk[] = []
let currentBatch: JsonValue[] = []
let currentTokens = 0
const contextHeader = path.length > 0 ? `// ${path.join('.')}\n` : ''
for (let i = 0; i < arr.length; i++) {
const item = arr[i]
const itemStr = JSON.stringify(item, null, 2)
const itemTokens = estimateTokens(itemStr)
if (itemTokens > this.chunkSize) {
if (currentBatch.length > 0) {
chunks.push(
this.buildBatchChunk(contextHeader, currentBatch, i - currentBatch.length, i - 1)
)
currentBatch = []
currentTokens = 0
}
if (depth < MAX_DEPTH && typeof item === 'object' && item !== null) {
chunks.push(...this.chunkStructuredData(item, [...path, `[${i}]`], depth + 1))
} else {
chunks.push({
text: contextHeader + itemStr,
tokenCount: itemTokens,
metadata: { startIndex: i, endIndex: i },
})
}
} else if (currentTokens + itemTokens > this.chunkSize && currentBatch.length > 0) {
chunks.push(
this.buildBatchChunk(contextHeader, currentBatch, i - currentBatch.length, i - 1)
)
currentBatch = [item]
currentTokens = itemTokens
} else {
currentBatch.push(item)
currentTokens += itemTokens
}
}
if (currentBatch.length > 0) {
chunks.push(
this.buildBatchChunk(
contextHeader,
currentBatch,
arr.length - currentBatch.length,
arr.length - 1
)
)
}
return chunks
}
private chunkObject(obj: JsonObject, path: string[], depth: number): Chunk[] {
const chunks: Chunk[] = []
const entries = Object.entries(obj)
const fullContent = JSON.stringify(obj, null, 2)
const fullTokens = estimateTokens(fullContent)
if (fullTokens <= this.chunkSize) {
const contextHeader = path.length > 0 ? `// ${path.join('.')}\n` : ''
const text = contextHeader + fullContent
return [
{
text,
tokenCount: estimateTokens(text),
metadata: { startIndex: 0, endIndex: text.length },
},
]
}
const contextHeader = path.length > 0 ? `// ${path.join('.')}\n` : ''
let currentObj: JsonObject = {}
let currentTokens = 0
for (const [key, value] of entries) {
const valueStr = JSON.stringify({ [key]: value }, null, 2)
const valueTokens = estimateTokens(valueStr)
if (valueTokens > this.chunkSize) {
if (Object.keys(currentObj).length > 0) {
const objContent = contextHeader + JSON.stringify(currentObj, null, 2)
chunks.push({
text: objContent,
tokenCount: estimateTokens(objContent),
metadata: { startIndex: 0, endIndex: objContent.length },
})
currentObj = {}
currentTokens = 0
}
if (depth < MAX_DEPTH && typeof value === 'object' && value !== null) {
chunks.push(...this.chunkStructuredData(value, [...path, key], depth + 1))
} else {
chunks.push({
text: contextHeader + valueStr,
tokenCount: valueTokens,
metadata: { startIndex: 0, endIndex: valueStr.length },
})
}
} else if (
currentTokens + valueTokens > this.chunkSize &&
Object.keys(currentObj).length > 0
) {
const objContent = contextHeader + JSON.stringify(currentObj, null, 2)
chunks.push({
text: objContent,
tokenCount: estimateTokens(objContent),
metadata: { startIndex: 0, endIndex: objContent.length },
})
currentObj = { [key]: value }
currentTokens = valueTokens
} else {
currentObj[key] = value
currentTokens += valueTokens
}
}
if (Object.keys(currentObj).length > 0) {
const objContent = contextHeader + JSON.stringify(currentObj, null, 2)
chunks.push({
text: objContent,
tokenCount: estimateTokens(objContent),
metadata: { startIndex: 0, endIndex: objContent.length },
})
}
return chunks
}
private buildBatchChunk(
contextHeader: string,
batch: JsonValue[],
startIdx: number,
endIdx: number
): Chunk {
const batchContent = contextHeader + JSON.stringify(batch, null, 2)
return {
text: batchContent,
tokenCount: estimateTokens(batchContent),
metadata: { startIndex: startIdx, endIndex: endIdx },
}
}
private chunkAsText(content: string): Chunk[] {
const chunks: Chunk[] = []
const lines = content.split('\n')
let currentChunk = ''
let currentTokens = 0
let startIndex = 0
for (const line of lines) {
const lineTokens = estimateTokens(line)
if (currentTokens + lineTokens > this.chunkSize && currentChunk) {
chunks.push({
text: currentChunk,
tokenCount: currentTokens,
metadata: { startIndex, endIndex: startIndex + currentChunk.length },
})
startIndex += currentChunk.length + 1
currentChunk = line
currentTokens = lineTokens
} else {
currentChunk = currentChunk ? `${currentChunk}\n${line}` : line
currentTokens += lineTokens
}
}
if (currentChunk && currentChunk.length >= this.minCharactersPerChunk) {
chunks.push({
text: currentChunk,
tokenCount: currentTokens,
metadata: { startIndex, endIndex: startIndex + currentChunk.length },
})
}
return chunks
}
static async chunkJsonYaml(content: string, options: ChunkerOptions = {}): Promise<Chunk[]> {
const chunker = new JsonYamlChunker(options)
return chunker.chunk(content)
}
}
@@ -0,0 +1,272 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import { RecursiveChunker } from './recursive-chunker'
describe('RecursiveChunker', () => {
describe('empty and whitespace input', () => {
it.concurrent('should return empty array for empty string', async () => {
const chunker = new RecursiveChunker({ chunkSize: 100 })
const chunks = await chunker.chunk('')
expect(chunks).toEqual([])
})
it.concurrent('should return empty array for whitespace-only input', async () => {
const chunker = new RecursiveChunker({ chunkSize: 100 })
const chunks = await chunker.chunk(' \n\n\t ')
expect(chunks).toEqual([])
})
})
describe('small content', () => {
it.concurrent('should return single chunk when content fits in one chunk', async () => {
const chunker = new RecursiveChunker({ chunkSize: 100 })
const text = 'This is a short text.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
})
})
describe('paragraph splitting', () => {
it.concurrent('should split at paragraph boundaries first', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20 })
const text =
'First paragraph with enough content to matter.\n\nSecond paragraph with enough content to matter.\n\nThird paragraph with enough content here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('line splitting fallback', () => {
it.concurrent('should split at newlines when paragraphs are too large', async () => {
const chunker = new RecursiveChunker({ chunkSize: 15 })
const text =
'Line one with content here.\nLine two with content here.\nLine three with content here.\nLine four with content here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('sentence splitting fallback', () => {
it.concurrent('should split at sentence boundaries when lines are too large', async () => {
const chunker = new RecursiveChunker({ chunkSize: 10 })
const text =
'First sentence here. Second sentence here. Third sentence here. Fourth sentence here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('word splitting fallback', () => {
it.concurrent('should split at spaces when sentences are too large', async () => {
const chunker = new RecursiveChunker({ chunkSize: 5 })
const text = 'word1 word2 word3 word4 word5 word6 word7 word8 word9 word10'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('keep_separator behavior', () => {
it.concurrent('should prepend separator to subsequent chunks', async () => {
const chunker = new RecursiveChunker({ chunkSize: 15 })
const text =
'First paragraph content here.\n\nSecond paragraph content here.\n\nThird paragraph content here.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
expect(chunks[1].text.startsWith('\n\n') || chunks[1].text.length > 0).toBe(true)
}
})
})
describe('custom separators', () => {
it.concurrent('should use custom separators instead of default recipe', async () => {
const chunker = new RecursiveChunker({
chunkSize: 15,
separators: ['---', '\n'],
})
const text =
'Section one content here with words.---Section two content here with words.---Section three content here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('recipe: plain', () => {
it.concurrent('should use plain recipe by default', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20 })
const text =
'First paragraph with enough words to exceed the chunk size limit.\n\nSecond paragraph with enough words to exceed the chunk size limit.\n\nThird paragraph with enough words to exceed the chunk size limit.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('recipe: markdown', () => {
it.concurrent('should split at heading boundaries for markdown content', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20, recipe: 'markdown' })
const text =
'\n# Title\n\nParagraph content under the title goes here.\n\n## Subtitle\n\nMore text content under the subtitle goes here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle markdown horizontal rules', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20, recipe: 'markdown' })
const text =
'Section one content here.\n---\nSection two content here.\n---\nSection three content here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('recipe: code', () => {
it.concurrent('should split on function and class boundaries', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20, recipe: 'code' })
const text = [
'const x = 1;',
'function hello() {',
' return "hello";',
'}',
'function world() {',
' return "world";',
'}',
'class MyClass {',
' constructor() {}',
' method() { return true; }',
'}',
].join('\n')
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('chunk size respected', () => {
it.concurrent('should not exceed chunk size in tokens', async () => {
const chunkSize = 30
const chunker = new RecursiveChunker({ chunkSize })
const text = 'This is a test sentence with content. '.repeat(30)
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.tokenCount).toBeLessThanOrEqual(chunkSize + 5)
}
})
})
describe('overlap', () => {
it.concurrent('should share text between consecutive chunks when overlap is set', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20, chunkOverlap: 5 })
const text =
'First paragraph with some content here.\n\nSecond paragraph with different content here.\n\nThird paragraph with more content here.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
expect(chunks[1].text.length).toBeGreaterThan(0)
}
})
it.concurrent('should not add overlap when overlap is 0', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20, chunkOverlap: 0 })
const text =
'First sentence content here. Second sentence content here. Third sentence content here.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
const firstChunkEnd = chunks[0].text.slice(-10)
expect(chunks[1].text.startsWith(firstChunkEnd)).toBe(false)
}
})
})
describe('chunk metadata', () => {
it.concurrent('should include text, tokenCount, and metadata fields', async () => {
const chunker = new RecursiveChunker({ chunkSize: 100 })
const text = 'This is test content for metadata.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
expect(chunks[0].tokenCount).toBe(Math.ceil(text.length / 4))
expect(chunks[0].metadata.startIndex).toBeDefined()
expect(chunks[0].metadata.endIndex).toBeDefined()
})
it.concurrent('should have startIndex of 0 for the first chunk', async () => {
const chunker = new RecursiveChunker({ chunkSize: 100 })
const text = 'Some content here.'
const chunks = await chunker.chunk(text)
expect(chunks[0].metadata.startIndex).toBe(0)
})
it.concurrent('should have non-negative indices for all chunks', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20, chunkOverlap: 5 })
const text = 'First part. Second part. Third part. Fourth part. Fifth part.'
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.metadata.startIndex).toBeGreaterThanOrEqual(0)
expect(chunk.metadata.endIndex).toBeGreaterThanOrEqual(chunk.metadata.startIndex)
}
})
it.concurrent('should have endIndex greater than startIndex for non-empty chunks', async () => {
const chunker = new RecursiveChunker({ chunkSize: 20 })
const text = 'Multiple sentences here. Another one here. And another. And more content.'
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.metadata.endIndex).toBeGreaterThan(chunk.metadata.startIndex)
}
})
})
describe('edge cases', () => {
it.concurrent('should handle very long text', async () => {
const chunker = new RecursiveChunker({ chunkSize: 100 })
const text = 'This is a sentence. '.repeat(1000)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle text with no natural separators', async () => {
const chunker = new RecursiveChunker({ chunkSize: 5 })
const text = 'abcdefghijklmnopqrstuvwxyz'.repeat(5)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle unicode text', async () => {
const chunker = new RecursiveChunker({ chunkSize: 100 })
const text = '这是中文测试。日本語テスト。한국어 테스트.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('中文')
})
it.concurrent('should use default chunkSize of 1024 tokens', async () => {
const chunker = new RecursiveChunker({})
const text = 'Word '.repeat(400)
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
})
})
})
+145
View File
@@ -0,0 +1,145 @@
import { createLogger } from '@sim/logger'
import type { Chunk, RecursiveChunkerOptions } from '@/lib/chunkers/types'
import {
addOverlap,
buildChunks,
cleanText,
estimateTokens,
resolveChunkerOptions,
splitAtWordBoundaries,
tokensToChars,
} from '@/lib/chunkers/utils'
const logger = createLogger('RecursiveChunker')
const RECIPES = {
plain: ['\n\n', '\n', '. ', ' ', ''],
markdown: [
'\n---\n',
'\n***\n',
'\n___\n',
'\n# ',
'\n## ',
'\n### ',
'\n#### ',
'\n##### ',
'\n###### ',
'\n```\n',
'\n> ',
'\n\n',
'\n',
'. ',
' ',
'',
],
code: [
'\nfunction ',
'\nclass ',
'\nexport ',
'\nconst ',
'\nlet ',
'\nvar ',
'\nif ',
'\nfor ',
'\nwhile ',
'\nswitch ',
'\nreturn ',
'\n\n',
'\n',
'; ',
' ',
'',
],
} as const
export class RecursiveChunker {
private readonly chunkSize: number
private readonly chunkOverlap: number
private readonly separators: string[]
constructor(options: RecursiveChunkerOptions = {}) {
const resolved = resolveChunkerOptions(options)
this.chunkSize = resolved.chunkSize
this.chunkOverlap = resolved.chunkOverlap
if (options.separators && options.separators.length > 0) {
this.separators = options.separators
} else {
const recipe = options.recipe ?? 'plain'
this.separators = [...RECIPES[recipe]]
}
}
private splitRecursively(text: string, separatorIndex = 0): string[] {
const tokenCount = estimateTokens(text)
if (tokenCount <= this.chunkSize) {
return text.trim() ? [text] : []
}
if (separatorIndex >= this.separators.length) {
const chunkSizeChars = tokensToChars(this.chunkSize)
return splitAtWordBoundaries(text, chunkSizeChars)
}
const separator = this.separators[separatorIndex]
if (separator === '') {
return this.splitRecursively(text, this.separators.length)
}
const parts = text.split(separator).filter((part) => part.trim())
if (parts.length <= 1) {
return this.splitRecursively(text, separatorIndex + 1)
}
const chunks: string[] = []
let currentChunk = ''
for (const part of parts) {
const testChunk = currentChunk + (currentChunk ? separator : '') + part
if (estimateTokens(testChunk) <= this.chunkSize) {
currentChunk = testChunk
} else {
if (currentChunk.trim()) {
chunks.push(currentChunk.trim())
}
if (estimateTokens(part) > this.chunkSize) {
const subChunks = this.splitRecursively(part, separatorIndex + 1)
for (const subChunk of subChunks) {
chunks.push(subChunk)
}
currentChunk = ''
} else {
currentChunk = part
}
}
}
if (currentChunk.trim()) {
chunks.push(currentChunk.trim())
}
return chunks
}
async chunk(content: string): Promise<Chunk[]> {
if (!content?.trim()) {
return []
}
const cleaned = cleanText(content)
let chunks = this.splitRecursively(cleaned)
if (this.chunkOverlap > 0) {
const overlapChars = tokensToChars(this.chunkOverlap)
chunks = addOverlap(chunks, overlapChars)
}
logger.info(`Chunked into ${chunks.length} recursive chunks`)
return buildChunks(chunks, this.chunkOverlap)
}
}
+391
View File
@@ -0,0 +1,391 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import { RegexChunker } from './regex-chunker'
describe('RegexChunker', () => {
describe('empty and whitespace input', () => {
it.concurrent('should return empty array for empty string', async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n' })
const chunks = await chunker.chunk('')
expect(chunks).toEqual([])
})
it.concurrent('should return empty array for whitespace-only input', async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n' })
const chunks = await chunker.chunk(' \n\n ')
expect(chunks).toEqual([])
})
})
describe('small content', () => {
it.concurrent('should return single chunk when content fits in chunkSize', async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n', chunkSize: 100 })
const text = 'This is a short text.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
})
})
describe('basic regex splitting', () => {
it.concurrent('should split on double newlines with pattern \\n\\n', async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n', chunkSize: 20 })
const text =
'First paragraph content here.\n\nSecond paragraph content here.\n\nThird paragraph content here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('custom pattern splitting', () => {
it.concurrent('should split text at --- delimiters', async () => {
const chunker = new RegexChunker({ pattern: '---', chunkSize: 20 })
const text =
'Section one has enough content to fill a chunk on its own here.---Section two also has enough content to fill another chunk here.---Section three needs content too for splitting.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('segment merging', () => {
it.concurrent('should merge small adjacent segments up to chunkSize', async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n', chunkSize: 100 })
const text = 'Short.\n\nAlso short.\n\nTiny.\n\nSmall too.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toContain('Short.')
expect(chunks[0].text).toContain('Also short.')
})
})
describe('oversized segment fallback', () => {
it.concurrent(
'should sub-chunk segments larger than chunkSize via word boundaries',
async () => {
const chunker = new RegexChunker({ pattern: '---', chunkSize: 10 })
const longSegment =
'This is a very long segment with many words that exceeds the chunk size limit significantly. '
const text = `${longSegment}---${longSegment}`
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(2)
}
)
})
describe('no-match fallback', () => {
it.concurrent(
'should fall back to word-boundary splitting when regex matches nothing',
async () => {
const chunker = new RegexChunker({ pattern: '###SPLIT###', chunkSize: 10 })
const text = 'This is a text with no matching delimiter anywhere in the content at all.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
}
)
})
describe('chunk size respected', () => {
it.concurrent('should not exceed chunkSize tokens approximately', async () => {
const chunkSize = 30
const chunker = new RegexChunker({ pattern: '\\n\\n', chunkSize })
const text =
'Paragraph one with some words. '.repeat(5) +
'\n\n' +
'Paragraph two with more words. '.repeat(5) +
'\n\n' +
'Paragraph three continues here. '.repeat(5)
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.tokenCount).toBeLessThanOrEqual(chunkSize + 10)
}
})
})
describe('overlap', () => {
it.concurrent('should share content between chunks when chunkOverlap > 0', async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n', chunkSize: 20, chunkOverlap: 5 })
const text =
'First paragraph with enough content.\n\nSecond paragraph with more content.\n\nThird paragraph with even more.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
const firstChunkEnd = chunks[0].text.slice(-10)
const secondChunkStart = chunks[1].text.slice(0, 20)
expect(secondChunkStart.length).toBeGreaterThan(0)
expect(chunks[1].text.length).toBeGreaterThan(0)
}
})
})
describe('chunk metadata', () => {
it.concurrent('should include text, tokenCount, and metadata with indices', async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n', chunkSize: 100 })
const text = 'Hello world test content.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
expect(chunks[0].tokenCount).toBe(Math.ceil(text.length / 4))
expect(chunks[0].metadata.startIndex).toBeDefined()
expect(chunks[0].metadata.endIndex).toBeDefined()
expect(chunks[0].metadata.startIndex).toBe(0)
})
it.concurrent('should have non-negative indices across multiple chunks', async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n', chunkSize: 20, chunkOverlap: 0 })
const text = 'First paragraph here.\n\nSecond paragraph here.\n\nThird paragraph here.'
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.metadata.startIndex).toBeGreaterThanOrEqual(0)
expect(chunk.metadata.endIndex).toBeGreaterThanOrEqual(chunk.metadata.startIndex)
}
})
})
describe('invalid regex', () => {
it.concurrent('should throw error for invalid regex pattern', async () => {
expect(() => new RegexChunker({ pattern: '[invalid' })).toThrow()
})
})
describe('empty pattern', () => {
it.concurrent('should throw error for empty pattern', async () => {
expect(() => new RegexChunker({ pattern: '' })).toThrow('Regex pattern is required')
})
})
describe('pattern too long', () => {
it.concurrent('should throw error for pattern exceeding 500 characters', async () => {
const longPattern = 'a'.repeat(501)
expect(() => new RegexChunker({ pattern: longPattern })).toThrow(
'Regex pattern exceeds maximum length of 500 characters'
)
})
})
describe('ReDoS protection', () => {
it.concurrent('should accept safe pattern \\n+', async () => {
expect(() => new RegexChunker({ pattern: '\\n+' })).not.toThrow()
})
it.concurrent('should accept safe pattern [,;]', async () => {
expect(() => new RegexChunker({ pattern: '[,;]' })).not.toThrow()
})
})
describe('capturing groups', () => {
it.concurrent(
'should not include delimiter text as a chunk when pattern has capturing groups',
async () => {
const chunker = new RegexChunker({
pattern: '(---)',
chunkSize: 1024,
strictBoundaries: true,
})
const text = 'Section one content.---Section two content.---Section three content.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(3)
expect(chunks[0].text).toBe('Section one content.')
expect(chunks[1].text).toBe('Section two content.')
expect(chunks[2].text).toBe('Section three content.')
for (const chunk of chunks) {
expect(chunk.text).not.toBe('---')
}
}
)
it.concurrent(
'should not include delimiter text when pattern uses named capture groups',
async () => {
const chunker = new RegexChunker({
pattern: '(?<sep>---)',
chunkSize: 1024,
strictBoundaries: true,
})
const text = 'Section one content.---Section two content.---Section three content.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(3)
expect(chunks[0].text).toBe('Section one content.')
expect(chunks[1].text).toBe('Section two content.')
expect(chunks[2].text).toBe('Section three content.')
for (const chunk of chunks) {
expect(chunk.text).not.toBe('---')
}
}
)
it.concurrent('should preserve lookbehind whose body contains a > character', async () => {
const chunker = new RegexChunker({
pattern: '(?<=</section>)',
chunkSize: 1024,
strictBoundaries: true,
})
const text = '<section>one</section><section>two</section><section>three</section>'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(3)
expect(chunks[0].text).toBe('<section>one</section>')
expect(chunks[1].text).toBe('<section>two</section>')
expect(chunks[2].text).toBe('<section>three</section>')
})
it.concurrent('should leave non-capturing groups and lookarounds intact', async () => {
const chunker = new RegexChunker({
pattern: '(?=\\n\\s*\\{\\s*"id"\\s*:)',
chunkSize: 1024,
strictBoundaries: true,
})
const text = '{"id": 1, "v": "a"}\n{"id": 2, "v": "b"}\n{"id": 3, "v": "c"}'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(3)
})
})
describe('strictBoundaries mode', () => {
it.concurrent(
'should produce one chunk per match without merging small adjacent segments',
async () => {
const chunker = new RegexChunker({
pattern: '\\n\\n',
chunkSize: 1024,
strictBoundaries: true,
})
const text = 'Short.\n\nAlso short.\n\nTiny.\n\nSmall too.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(4)
expect(chunks[0].text).toBe('Short.')
expect(chunks[1].text).toBe('Also short.')
expect(chunks[2].text).toBe('Tiny.')
expect(chunks[3].text).toBe('Small too.')
}
)
it.concurrent('should produce one chunk per QA record using lookahead pattern', async () => {
const chunker = new RegexChunker({
pattern: '(?=\\n\\s*\\{\\s*"id"\\s*:)',
chunkSize: 1024,
strictBoundaries: true,
})
const text =
'{"id": 1, "q": "first?", "a": "yes"}\n{"id": 2, "q": "second?", "a": "no"}\n{"id": 3, "q": "third?", "a": "maybe"}'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(3)
expect(chunks[0].text).toContain('"id": 1')
expect(chunks[0].text).not.toContain('"id": 2')
expect(chunks[1].text).toContain('"id": 2')
expect(chunks[1].text).not.toContain('"id": 3')
expect(chunks[2].text).toContain('"id": 3')
})
it.concurrent('should not apply overlap even when chunkOverlap is set', async () => {
const chunker = new RegexChunker({
pattern: '\\n\\n',
chunkSize: 100,
chunkOverlap: 50,
strictBoundaries: true,
})
const text = 'First section content.\n\nSecond section content.\n\nThird section content.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(3)
expect(chunks[0].text).toBe('First section content.')
expect(chunks[1].text).toBe('Second section content.')
expect(chunks[2].text).toBe('Third section content.')
})
it.concurrent(
'should still split when content fits in single chunk if matches exist',
async () => {
const chunker = new RegexChunker({
pattern: '\\n\\n',
chunkSize: 1024,
strictBoundaries: true,
})
const text = 'A.\n\nB.\n\nC.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(3)
}
)
it.concurrent('should sub-chunk a single oversized segment at word boundaries', async () => {
const chunker = new RegexChunker({
pattern: '---',
chunkSize: 10,
strictBoundaries: true,
})
const longSegment =
'This is a very long segment with many words that exceeds the chunk size limit significantly.'
const text = `${longSegment}---short`
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(2)
expect(chunks[chunks.length - 1].text).toBe('short')
})
it.concurrent('should return single chunk when regex finds no matches', async () => {
const chunker = new RegexChunker({
pattern: '###NOMATCH###',
chunkSize: 1024,
strictBoundaries: true,
})
const text = 'Plain text with no delimiter at all.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
})
it.concurrent('should return empty array for empty input', async () => {
const chunker = new RegexChunker({
pattern: '\\n\\n',
strictBoundaries: true,
})
const chunks = await chunker.chunk('')
expect(chunks).toEqual([])
})
it.concurrent(
'should default to merging behavior when strictBoundaries is omitted',
async () => {
const chunker = new RegexChunker({ pattern: '\\n\\n', chunkSize: 100 })
const text = 'Short.\n\nAlso short.\n\nTiny.\n\nSmall too.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
}
)
it.concurrent('should produce non-overlapping startIndex/endIndex metadata', async () => {
const chunker = new RegexChunker({
pattern: '\\n\\n',
chunkSize: 1024,
chunkOverlap: 50,
strictBoundaries: true,
})
const text = 'First.\n\nSecond.\n\nThird.'
const chunks = await chunker.chunk(text)
for (let i = 1; i < chunks.length; i++) {
expect(chunks[i].metadata.startIndex).toBeGreaterThanOrEqual(
chunks[i - 1].metadata.endIndex
)
}
})
})
})
+219
View File
@@ -0,0 +1,219 @@
import { createLogger } from '@sim/logger'
import { toError } from '@sim/utils/errors'
import type { Chunk, RegexChunkerOptions } from '@/lib/chunkers/types'
import {
addOverlap,
buildChunks,
cleanText,
estimateTokens,
resolveChunkerOptions,
splitAtWordBoundaries,
tokensToChars,
} from '@/lib/chunkers/utils'
const logger = createLogger('RegexChunker')
const MAX_PATTERN_LENGTH = 500
const NAMED_GROUP_PREFIX = /^\(\?<(?![=!])[^>]+>/
/**
* Converts unescaped capturing groups `(...)` and named capturing groups
* `(?<name>...)` into non-capturing groups `(?:...)`. `String.prototype.split()`
* interleaves captured text (named or otherwise) into the result array, which
* would surface delimiter text as spurious chunks. Lookarounds (`(?=`, `(?!`,
* `(?<=`, `(?<!`) and other `(?...)` constructs are left untouched.
*/
function toNonCapturing(pattern: string): string {
let result = ''
let inClass = false
for (let i = 0; i < pattern.length; i++) {
const c = pattern[i]
if (c === '\\' && i + 1 < pattern.length) {
result += c + pattern[i + 1]
i++
continue
}
if (c === '[') inClass = true
else if (c === ']') inClass = false
if (!inClass && c === '(') {
if (pattern[i + 1] !== '?') {
result += '(?:'
continue
}
const namedMatch = pattern.slice(i).match(NAMED_GROUP_PREFIX)
if (namedMatch) {
result += '(?:'
i += namedMatch[0].length - 1
continue
}
}
result += c
}
return result
}
export class RegexChunker {
private readonly chunkSize: number
private readonly chunkOverlap: number
private readonly regex: RegExp
private readonly strictBoundaries: boolean
constructor(options: RegexChunkerOptions) {
const resolved = resolveChunkerOptions(options)
this.chunkSize = resolved.chunkSize
this.chunkOverlap = resolved.chunkOverlap
this.regex = this.compilePattern(options.pattern)
this.strictBoundaries = options.strictBoundaries ?? false
}
private compilePattern(pattern: string): RegExp {
if (!pattern) {
throw new Error('Regex pattern is required')
}
if (pattern.length > MAX_PATTERN_LENGTH) {
throw new Error(`Regex pattern exceeds maximum length of ${MAX_PATTERN_LENGTH} characters`)
}
try {
const regex = new RegExp(toNonCapturing(pattern), 'g')
const testStrings = [
'a'.repeat(10000),
' '.repeat(10000),
'a '.repeat(5000),
'aB1 xY2\n'.repeat(1250),
`${'a'.repeat(30)}!`,
`${'a b '.repeat(25)}!`,
]
for (const testStr of testStrings) {
regex.lastIndex = 0
const start = Date.now()
regex.test(testStr)
const elapsed = Date.now() - start
if (elapsed > 50) {
throw new Error('Regex pattern appears to have catastrophic backtracking')
}
}
regex.lastIndex = 0
return regex
} catch (error) {
if (error instanceof Error && error.message.includes('catastrophic')) {
throw error
}
throw new Error(`Invalid regex pattern "${pattern}": ${toError(error).message}`)
}
}
async chunk(content: string): Promise<Chunk[]> {
if (!content?.trim()) {
return []
}
const cleaned = cleanText(content)
if (!this.strictBoundaries && estimateTokens(cleaned) <= this.chunkSize) {
logger.info('Content fits in single chunk')
return buildChunks([cleaned], 0)
}
this.regex.lastIndex = 0
const segments = cleaned.split(this.regex).filter((s) => s.trim().length > 0)
if (segments.length <= 1) {
if (this.strictBoundaries) {
logger.info('Regex pattern produced no splits in strict mode, returning single chunk')
return buildChunks([cleaned.trim()], 0)
}
logger.warn(
'Regex pattern did not produce any splits, falling back to word-boundary splitting'
)
const chunkSizeChars = tokensToChars(this.chunkSize)
let chunks = splitAtWordBoundaries(cleaned, chunkSizeChars)
if (this.chunkOverlap > 0) {
const overlapChars = tokensToChars(this.chunkOverlap)
chunks = addOverlap(chunks, overlapChars)
}
return buildChunks(chunks, this.chunkOverlap)
}
if (this.strictBoundaries) {
const chunks = this.expandOversizedSegments(segments)
logger.info(`Chunked into ${chunks.length} strict-boundary regex chunks`)
return buildChunks(chunks, 0)
}
const merged = this.mergeSegments(segments)
let chunks = merged
if (this.chunkOverlap > 0) {
const overlapChars = tokensToChars(this.chunkOverlap)
chunks = addOverlap(chunks, overlapChars)
}
logger.info(`Chunked into ${chunks.length} regex-based chunks`)
return buildChunks(chunks, this.chunkOverlap)
}
/**
* In strict-boundary mode each segment becomes its own chunk. Segments that
* exceed chunkSize are still split at word boundaries to preserve the token
* limit invariant; this is a safety floor, not a merge.
*/
private expandOversizedSegments(segments: string[]): string[] {
const result: string[] = []
const chunkSizeChars = tokensToChars(this.chunkSize)
for (const segment of segments) {
const trimmed = segment.trim()
if (!trimmed) continue
if (estimateTokens(trimmed) <= this.chunkSize) {
result.push(trimmed)
} else {
const subChunks = splitAtWordBoundaries(trimmed, chunkSizeChars)
for (const sub of subChunks) {
if (sub.trim()) result.push(sub)
}
}
}
return result
}
private mergeSegments(segments: string[]): string[] {
const chunks: string[] = []
let current = ''
for (const segment of segments) {
const test = current ? `${current}\n${segment}` : segment
if (estimateTokens(test) <= this.chunkSize) {
current = test
} else {
if (current.trim()) {
chunks.push(current.trim())
}
if (estimateTokens(segment) > this.chunkSize) {
const chunkSizeChars = tokensToChars(this.chunkSize)
const subChunks = splitAtWordBoundaries(segment, chunkSizeChars)
for (const sub of subChunks) {
chunks.push(sub)
}
current = ''
} else {
current = segment
}
}
}
if (current.trim()) {
chunks.push(current.trim())
}
return chunks
}
}
@@ -0,0 +1,283 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import { SentenceChunker } from './sentence-chunker'
describe('SentenceChunker', () => {
describe('empty and whitespace input', () => {
it.concurrent('should return empty array for empty string', async () => {
const chunker = new SentenceChunker({ chunkSize: 100 })
const chunks = await chunker.chunk('')
expect(chunks).toEqual([])
})
it.concurrent('should return empty array for whitespace-only input', async () => {
const chunker = new SentenceChunker({ chunkSize: 100 })
const chunks = await chunker.chunk(' \n\n\t ')
expect(chunks).toEqual([])
})
it.concurrent('should return empty array for null-ish content', async () => {
const chunker = new SentenceChunker({ chunkSize: 100 })
const chunks = await chunker.chunk(undefined as unknown as string)
expect(chunks).toEqual([])
})
})
describe('small content (single chunk)', () => {
it.concurrent('should return single chunk when content fits within chunk size', async () => {
const chunker = new SentenceChunker({ chunkSize: 100 })
const text = 'This is a short sentence. Another short one.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
expect(chunks[0].tokenCount).toBe(Math.ceil(text.length / 4))
})
})
describe('sentence boundary splitting', () => {
it.concurrent('should split text at sentence boundaries', async () => {
const chunker = new SentenceChunker({ chunkSize: 20 })
const text =
'First sentence here. Second sentence here. Third sentence here. Fourth sentence here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
for (let i = 0; i < chunks.length - 1; i++) {
const trimmed = chunks[i].text.trim()
const lastChar = trimmed[trimmed.length - 1]
expect(['.', '!', '?']).toContain(lastChar)
}
})
})
describe('abbreviation handling', () => {
it.concurrent('should not split at common abbreviations', async () => {
const chunker = new SentenceChunker({ chunkSize: 200 })
const text = 'Mr. Smith went to Washington. He arrived on Jan. 5th.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toContain('Mr. Smith')
expect(chunks[0].text).toContain('Jan. 5th')
})
it.concurrent('should not split at Dr., Mrs., Ms., Prof., Jr., Sr., St.', async () => {
const chunker = new SentenceChunker({ chunkSize: 500 })
const text =
'Dr. Jones and Mrs. Brown met Prof. Davis at St. Mary hospital. Jr. members joined Sr. staff in Feb. for a review.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
})
})
describe('single capital initial handling', () => {
it.concurrent('should not split at single capital letter initials', async () => {
const chunker = new SentenceChunker({ chunkSize: 200 })
const text = 'J. K. Rowling wrote books. They are popular.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toContain('J. K. Rowling')
})
})
describe('decimal handling', () => {
it.concurrent('should not split at decimal numbers', async () => {
const chunker = new SentenceChunker({ chunkSize: 20 })
const text = 'The value is 3.14. That is pi.'
const chunks = await chunker.chunk(text)
const allText = chunks.map((c) => c.text).join(' ')
expect(allText).toContain('3.14')
const largeChunker = new SentenceChunker({ chunkSize: 200 })
const largeChunks = await largeChunker.chunk(text)
expect(largeChunks).toHaveLength(1)
})
})
describe('ellipsis handling', () => {
it.concurrent('should not split at ellipsis', async () => {
const chunker = new SentenceChunker({ chunkSize: 200 })
const text = 'Wait for it... The answer is here. Done.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toContain('Wait for it...')
})
})
describe('exclamation and question marks', () => {
it.concurrent('should split at exclamation and question marks', async () => {
const chunker = new SentenceChunker({ chunkSize: 10 })
const text = 'What is this? It is great! I agree.'
const chunks = await chunker.chunk(text)
const allText = chunks.map((c) => c.text).join(' ')
expect(allText).toContain('What is this?')
expect(allText).toContain('It is great!')
expect(allText).toContain('I agree.')
})
it.concurrent('should treat ? and ! as sentence boundaries', async () => {
const chunker = new SentenceChunker({ chunkSize: 15 })
const text = 'What is this thing? It is really great! I strongly agree.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThanOrEqual(1)
const allText = chunks.map((c) => c.text).join(' ')
expect(allText).toContain('?')
expect(allText).toContain('!')
})
})
describe('minSentencesPerChunk', () => {
it.concurrent('should group at least minSentencesPerChunk sentences per chunk', async () => {
const chunker = new SentenceChunker({ chunkSize: 100, minSentencesPerChunk: 2 })
const text =
'First sentence. Second sentence. Third sentence. Fourth sentence. Fifth sentence.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks).toHaveLength(1)
})
it.concurrent('should enforce min sentences even when token limit is reached', async () => {
const chunker = new SentenceChunker({ chunkSize: 6, minSentencesPerChunk: 2 })
const text = 'Short one. Another one. Third one here. Fourth one here.'
const chunks = await chunker.chunk(text)
const firstChunkSentences = chunks[0].text
.split(/(?<=[.!?])\s+/)
.filter((s) => s.trim().length > 0)
expect(firstChunkSentences.length).toBeGreaterThanOrEqual(2)
})
})
describe('oversized sentence fallback', () => {
it.concurrent(
'should chunk a single very long sentence via word-boundary splitting',
async () => {
const chunker = new SentenceChunker({ chunkSize: 10 })
const longSentence = `${'word '.repeat(50).trim()}.`
const chunks = await chunker.chunk(longSentence)
expect(chunks.length).toBeGreaterThan(1)
const allText = chunks.map((c) => c.text).join(' ')
expect(allText).toContain('word')
}
)
it.concurrent('should handle oversized sentence mixed with normal sentences', async () => {
const chunker = new SentenceChunker({ chunkSize: 10 })
const longSentence = `${'word '.repeat(50).trim()}.`
const text = `Short sentence. ${longSentence} Another short one.`
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(2)
const allText = chunks.map((c) => c.text).join(' ')
expect(allText).toContain('Short sentence.')
expect(allText).toContain('Another short one.')
})
})
describe('sentence-level overlap', () => {
it.concurrent('should include overlap from previous chunk when chunkOverlap > 0', async () => {
const chunker = new SentenceChunker({ chunkSize: 15, chunkOverlap: 10 })
const text =
'First sentence here. Second sentence here. Third sentence here. Fourth sentence here.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
expect(chunks[1].text.length).toBeGreaterThan(0)
}
})
it.concurrent('should not add overlap when chunkOverlap is 0', async () => {
const chunker = new SentenceChunker({ chunkSize: 15, chunkOverlap: 0 })
const text = 'First sentence here. Second sentence here. Third sentence here.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
const chunk1End = chunks[0].text.slice(-20)
expect(chunks[1].text.startsWith(chunk1End)).toBe(false)
}
})
})
describe('chunk metadata', () => {
it.concurrent('should include text, tokenCount, and metadata in each chunk', async () => {
const chunker = new SentenceChunker({ chunkSize: 100 })
const text = 'This is a test sentence. Another sentence follows.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0]).toHaveProperty('text')
expect(chunks[0]).toHaveProperty('tokenCount')
expect(chunks[0]).toHaveProperty('metadata')
expect(chunks[0].metadata).toHaveProperty('startIndex')
expect(chunks[0].metadata).toHaveProperty('endIndex')
})
it.concurrent('should have startIndex of 0 for the first chunk', async () => {
const chunker = new SentenceChunker({ chunkSize: 10 })
const text = 'First sentence. Second sentence. Third sentence.'
const chunks = await chunker.chunk(text)
expect(chunks[0].metadata.startIndex).toBe(0)
})
it.concurrent('should have non-negative indices for all chunks', async () => {
const chunker = new SentenceChunker({ chunkSize: 10, chunkOverlap: 5 })
const text =
'First sentence here. Second sentence here. Third sentence here. Fourth sentence.'
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.metadata.startIndex).toBeGreaterThanOrEqual(0)
expect(chunk.metadata.endIndex).toBeGreaterThanOrEqual(chunk.metadata.startIndex)
}
})
it.concurrent('should have correct tokenCount based on text length', async () => {
const chunker = new SentenceChunker({ chunkSize: 100 })
const text = 'Hello world test.'
const chunks = await chunker.chunk(text)
expect(chunks[0].tokenCount).toBe(Math.ceil(text.length / 4))
})
})
describe('respects chunk size', () => {
it.concurrent('should produce chunks within approximate token limit', async () => {
const chunkSize = 20
const chunker = new SentenceChunker({ chunkSize })
const text =
'This is the first sentence. Here is the second one. And the third sentence follows. Then comes the fourth. Finally the fifth sentence.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
for (const chunk of chunks) {
expect(chunk.tokenCount).toBeLessThanOrEqual(chunkSize * 2)
}
})
it.concurrent('should create more chunks with smaller chunk size', async () => {
const text =
'Sentence number one. Sentence number two. Sentence number three. Sentence number four. Sentence number five. Sentence number six.'
const largeChunker = new SentenceChunker({ chunkSize: 200 })
const smallChunker = new SentenceChunker({ chunkSize: 10 })
const largeChunks = await largeChunker.chunk(text)
const smallChunks = await smallChunker.chunk(text)
expect(smallChunks.length).toBeGreaterThan(largeChunks.length)
})
})
})
+141
View File
@@ -0,0 +1,141 @@
import { createLogger } from '@sim/logger'
import type { Chunk, SentenceChunkerOptions } from '@/lib/chunkers/types'
import {
buildChunks,
cleanText,
estimateTokens,
resolveChunkerOptions,
splitAtWordBoundaries,
tokensToChars,
} from '@/lib/chunkers/utils'
const logger = createLogger('SentenceChunker')
/** Never splits mid-sentence unless a single sentence exceeds the limit. */
export class SentenceChunker {
private readonly chunkSize: number
private readonly chunkOverlap: number
private readonly minSentencesPerChunk: number
constructor(options: SentenceChunkerOptions = {}) {
const resolved = resolveChunkerOptions(options)
this.chunkSize = resolved.chunkSize
this.chunkOverlap = resolved.chunkOverlap
this.minSentencesPerChunk = options.minSentencesPerChunk ?? 1
}
/** Splits on sentence boundaries while avoiding abbreviations, decimals, and ellipses. */
private splitSentences(text: string): string[] {
return text
.split(
/(?<!\b(?:Mr|Mrs|Ms|Dr|Prof|Sr|Jr|St|Rev|Gen|Sgt|Capt|Lt|Col|Maj|No|Fig|Vol|Ch|vs|etc|Inc|Ltd|Corp|Co|approx|dept|est|govt|Ave|Blvd|Rd|Jan|Feb|Mar|Apr|Jun|Jul|Aug|Sep|Oct|Nov|Dec|i\.e|e\.g)\.)(?<![A-Z]\.)(?<!\.\.)(?<!\d\.)(?<=[.!?])\s+/
)
.filter((s) => s.trim().length > 0)
}
async chunk(content: string): Promise<Chunk[]> {
if (!content?.trim()) {
return []
}
const cleaned = cleanText(content)
const sentences = this.splitSentences(cleaned)
if (sentences.length === 0) {
return []
}
if (estimateTokens(cleaned) <= this.chunkSize) {
logger.info('Content fits in single chunk')
return buildChunks([cleaned], 0)
}
const chunkSentenceGroups: string[][] = []
let currentGroup: string[] = []
let currentTokens = 0
const chunkSizeChars = tokensToChars(this.chunkSize)
for (const sentence of sentences) {
const sentenceTokens = estimateTokens(sentence)
if (sentenceTokens > this.chunkSize) {
if (currentGroup.length > 0) {
chunkSentenceGroups.push(currentGroup)
currentGroup = []
currentTokens = 0
}
const parts = splitAtWordBoundaries(sentence, chunkSizeChars)
for (const part of parts) {
chunkSentenceGroups.push([part])
}
continue
}
const wouldExceed = currentTokens + sentenceTokens > this.chunkSize
const hasMinSentences = currentGroup.length >= this.minSentencesPerChunk
if (wouldExceed && hasMinSentences) {
chunkSentenceGroups.push(currentGroup)
currentGroup = [sentence]
currentTokens = sentenceTokens
} else {
currentGroup.push(sentence)
currentTokens += sentenceTokens
}
}
if (currentGroup.length > 0) {
chunkSentenceGroups.push(currentGroup)
}
const rawChunks = this.applyOverlapFromGroups(chunkSentenceGroups)
logger.info(`Chunked into ${rawChunks.length} sentence-based chunks`)
return buildChunks(rawChunks, this.chunkOverlap)
}
/** Applies overlap at the sentence level using original groups to avoid re-splitting. */
private applyOverlapFromGroups(groups: string[][]): string[] {
if (this.chunkOverlap <= 0 || groups.length <= 1) {
return groups.map((g) => g.join(' '))
}
const overlapChars = tokensToChars(this.chunkOverlap)
const result: string[] = []
for (let i = 0; i < groups.length; i++) {
if (i === 0) {
result.push(groups[i].join(' '))
continue
}
const prevGroup = groups[i - 1]
const overlapSentences: string[] = []
let overlapLen = 0
for (let j = prevGroup.length - 1; j >= 0; j--) {
if (overlapLen + prevGroup[j].length > overlapChars) break
overlapSentences.unshift(prevGroup[j])
overlapLen += prevGroup[j].length
}
const currentText = groups[i].join(' ')
if (overlapSentences.length > 0) {
result.push(`${overlapSentences.join(' ')} ${currentText}`)
} else {
// No complete sentence fits — fall back to character-level overlap
const prevText = prevGroup.join(' ')
const tail = prevText.slice(-overlapChars)
const wordMatch = tail.match(/^\s*\S/)
const cleanTail = wordMatch ? tail.slice(tail.indexOf(wordMatch[0].trim())) : tail
if (cleanTail.trim()) {
result.push(`${cleanTail.trim()} ${currentText}`)
} else {
result.push(currentText)
}
}
}
return result
}
}
@@ -0,0 +1,334 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import { StructuredDataChunker } from './structured-data-chunker'
describe('StructuredDataChunker', () => {
describe('isStructuredData', () => {
it('should detect CSV content with many columns', () => {
const csv = 'name,age,city,country\nAlice,30,NYC,USA\nBob,25,LA,USA'
expect(StructuredDataChunker.isStructuredData(csv)).toBe(true)
})
it('should detect TSV content with many columns', () => {
const tsv = 'name\tage\tcity\tcountry\nAlice\t30\tNYC\tUSA\nBob\t25\tLA\tUSA'
expect(StructuredDataChunker.isStructuredData(tsv)).toBe(true)
})
it('should detect pipe-delimited content with many columns', () => {
const piped = 'name|age|city|country\nAlice|30|NYC|USA\nBob|25|LA|USA'
expect(StructuredDataChunker.isStructuredData(piped)).toBe(true)
})
it('should detect CSV by mime type', () => {
expect(StructuredDataChunker.isStructuredData('any content', 'text/csv')).toBe(true)
})
it('should detect XLSX by mime type', () => {
expect(
StructuredDataChunker.isStructuredData(
'any content',
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
)
).toBe(true)
})
it('should detect XLS by mime type', () => {
expect(
StructuredDataChunker.isStructuredData('any content', 'application/vnd.ms-excel')
).toBe(true)
})
it('should detect TSV by mime type', () => {
expect(
StructuredDataChunker.isStructuredData('any content', 'text/tab-separated-values')
).toBe(true)
})
it('should return false for plain text', () => {
const plainText = 'This is just regular text.\nWith some lines.\nNo structure here.'
expect(StructuredDataChunker.isStructuredData(plainText)).toBe(false)
})
it('should return false for single line', () => {
expect(StructuredDataChunker.isStructuredData('just one line')).toBe(false)
})
it('should handle inconsistent delimiter counts', () => {
const inconsistent = 'name,age\nAlice,30,extra\nBob'
const result = StructuredDataChunker.isStructuredData(inconsistent)
expect(typeof result).toBe('boolean')
})
})
describe('chunkStructuredData', () => {
it.concurrent('should return empty array for empty content', async () => {
const chunks = await StructuredDataChunker.chunkStructuredData('')
expect(chunks).toEqual([])
})
it.concurrent('should return empty array for whitespace only', async () => {
const chunks = await StructuredDataChunker.chunkStructuredData(' \n\n ')
expect(chunks).toEqual([])
})
it.concurrent('should chunk basic CSV data', async () => {
const csv = `name,age,city
Alice,30,New York
Bob,25,Los Angeles
Charlie,35,Chicago`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('Headers:')
expect(chunks[0].text).toContain('name,age,city')
})
it.concurrent('should include row count in chunks', async () => {
const csv = `name,age
Alice,30
Bob,25`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('rows of data')
})
it.concurrent('should include sheet name when provided', async () => {
const csv = `name,age
Alice,30`
const chunks = await StructuredDataChunker.chunkStructuredData(csv, { sheetName: 'Users' })
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('Users')
})
it.concurrent('should use provided headers when available', async () => {
const data = `Alice,30
Bob,25`
const chunks = await StructuredDataChunker.chunkStructuredData(data, {
headers: ['Name', 'Age'],
})
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('Name\tAge')
})
it.concurrent('should chunk large datasets into multiple chunks', async () => {
const rows = ['name,value']
for (let i = 0; i < 500; i++) {
rows.push(`Item${i},Value${i}`)
}
const csv = rows.join('\n')
const chunks = await StructuredDataChunker.chunkStructuredData(csv, { chunkSize: 200 })
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should include token count in chunk metadata', async () => {
const csv = `name,age
Alice,30
Bob,25`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].tokenCount).toBeGreaterThan(0)
})
})
describe('chunk metadata', () => {
it.concurrent('should include startIndex as row index', async () => {
const csv = `header1,header2
row1,data1
row2,data2
row3,data3`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].metadata.startIndex).toBeDefined()
expect(chunks[0].metadata.startIndex).toBeGreaterThanOrEqual(0)
})
it.concurrent('should include endIndex as row index', async () => {
const csv = `header1,header2
row1,data1
row2,data2`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].metadata.endIndex).toBeDefined()
expect(chunks[0].metadata.endIndex).toBeGreaterThanOrEqual(chunks[0].metadata.startIndex)
})
})
describe('edge cases', () => {
it.concurrent('should handle single data row', async () => {
const csv = `name,age
Alice,30`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBe(1)
})
it.concurrent('should handle header only', async () => {
const csv = 'name,age,city'
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThanOrEqual(0)
})
it.concurrent('should handle unicode content', async () => {
const csv = `名前,年齢,市
田中,30,東京
鈴木,25,大阪`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('田中')
})
it.concurrent('should handle quoted CSV fields', async () => {
const csv = `name,description
Alice,"Has a comma, in description"
Bob,"Multiple
lines"`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle empty cells', async () => {
const csv = `name,age,city
Alice,,NYC
,25,LA
Charlie,35,`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle long cell values', async () => {
const csv = `name,description
Alice,${'A'.repeat(1000)}
Bob,${'B'.repeat(1000)}`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle many columns', async () => {
const headers = Array.from({ length: 50 }, (_, i) => `col${i}`).join(',')
const row = Array.from({ length: 50 }, (_, i) => `val${i}`).join(',')
const csv = `${headers}\n${row}`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('options', () => {
it.concurrent('should respect custom chunkSize', async () => {
const rows = ['name,value']
for (let i = 0; i < 200; i++) {
rows.push(`Item${i},Value${i}`)
}
const csv = rows.join('\n')
const smallChunks = await StructuredDataChunker.chunkStructuredData(csv, { chunkSize: 100 })
const largeChunks = await StructuredDataChunker.chunkStructuredData(csv, { chunkSize: 2000 })
expect(smallChunks.length).toBeGreaterThan(largeChunks.length)
})
it.concurrent('should handle default options', async () => {
const csv = `name,age
Alice,30`
const chunks = await StructuredDataChunker.chunkStructuredData(csv)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('large inputs', () => {
it.concurrent('should handle 10,000 rows', async () => {
const rows = ['id,name,value']
for (let i = 0; i < 10000; i++) {
rows.push(`${i},Item${i},Value${i}`)
}
const csv = rows.join('\n')
const chunks = await StructuredDataChunker.chunkStructuredData(csv, { chunkSize: 500 })
expect(chunks.length).toBeGreaterThan(1)
const totalRowCount = chunks.reduce((sum, chunk) => {
const match = chunk.text.match(/\[(\d+) rows of data\]/)
return sum + (match ? Number.parseInt(match[1]) : 0)
}, 0)
expect(totalRowCount).toBeGreaterThan(0)
})
it.concurrent('should handle very wide rows', async () => {
const columns = 100
const headers = Array.from({ length: columns }, (_, i) => `column${i}`).join(',')
const rows = [headers]
for (let i = 0; i < 50; i++) {
rows.push(Array.from({ length: columns }, (_, j) => `r${i}c${j}`).join(','))
}
const csv = rows.join('\n')
const chunks = await StructuredDataChunker.chunkStructuredData(csv, { chunkSize: 300 })
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('delimiter detection', () => {
it.concurrent('should handle comma delimiter', async () => {
const csv = `a,b,c,d
1,2,3,4
5,6,7,8`
expect(StructuredDataChunker.isStructuredData(csv)).toBe(true)
})
it.concurrent('should handle tab delimiter', async () => {
const tsv = `a\tb\tc\td
1\t2\t3\t4
5\t6\t7\t8`
expect(StructuredDataChunker.isStructuredData(tsv)).toBe(true)
})
it.concurrent('should handle pipe delimiter', async () => {
const piped = `a|b|c|d
1|2|3|4
5|6|7|8`
expect(StructuredDataChunker.isStructuredData(piped)).toBe(true)
})
it.concurrent('should not detect with fewer than 3 delimiters per line', async () => {
const sparse = `a,b
1,2`
const result = StructuredDataChunker.isStructuredData(sparse)
expect(typeof result).toBe('boolean')
})
})
describe('header handling', () => {
it.concurrent('should include headers in each chunk by default', async () => {
const rows = ['name,value']
for (let i = 0; i < 100; i++) {
rows.push(`Item${i},Value${i}`)
}
const csv = rows.join('\n')
const chunks = await StructuredDataChunker.chunkStructuredData(csv, { chunkSize: 200 })
expect(chunks.length).toBeGreaterThan(1)
for (const chunk of chunks) {
expect(chunk.text).toContain('Headers:')
}
})
})
})
@@ -0,0 +1,176 @@
import { createLogger } from '@sim/logger'
import type { Chunk, StructuredDataOptions } from '@/lib/chunkers/types'
/** Structured data is denser in tokens (~3 chars/token vs ~4 for prose) */
function estimateStructuredTokens(text: string): number {
if (!text?.trim()) return 0
return Math.ceil(text.length / 3)
}
const logger = createLogger('StructuredDataChunker')
const DEFAULT_CONFIG = {
TARGET_CHUNK_SIZE: 1024,
MIN_ROWS_PER_CHUNK: 5,
MAX_ROWS_PER_CHUNK: 500,
INCLUDE_HEADERS_IN_EACH_CHUNK: true,
} as const
export class StructuredDataChunker {
static async chunkStructuredData(
content: string,
options: StructuredDataOptions = {}
): Promise<Chunk[]> {
const chunks: Chunk[] = []
const lines = content.split('\n').filter((line) => line.trim())
if (lines.length === 0) {
return chunks
}
const targetChunkSize = options.chunkSize ?? DEFAULT_CONFIG.TARGET_CHUNK_SIZE
const headerLine = options.headers?.join('\t') || lines[0]
const dataStartIndex = options.headers ? 0 : 1
const estimatedTokensPerRow = StructuredDataChunker.estimateStructuredTokensPerRow(
lines.slice(dataStartIndex, Math.min(10, lines.length))
)
const optimalRowsPerChunk = StructuredDataChunker.calculateOptimalRowsPerChunk(
estimatedTokensPerRow,
targetChunkSize
)
logger.info(
`Structured data chunking: ${lines.length} rows, ~${estimatedTokensPerRow} tokens/row, ${optimalRowsPerChunk} rows/chunk, target: ${targetChunkSize} tokens`
)
let currentChunkRows: string[] = []
let currentTokenEstimate = 0
const headerTokens = estimateStructuredTokens(headerLine)
let chunkStartRow = dataStartIndex
for (let i = dataStartIndex; i < lines.length; i++) {
const row = lines[i]
const rowTokens = estimateStructuredTokens(row)
const projectedTokens =
currentTokenEstimate +
rowTokens +
(DEFAULT_CONFIG.INCLUDE_HEADERS_IN_EACH_CHUNK ? headerTokens : 0)
const shouldCreateChunk =
(projectedTokens > targetChunkSize &&
currentChunkRows.length >= DEFAULT_CONFIG.MIN_ROWS_PER_CHUNK) ||
currentChunkRows.length >= optimalRowsPerChunk
if (shouldCreateChunk && currentChunkRows.length > 0) {
const chunkContent = StructuredDataChunker.formatChunk(
headerLine,
currentChunkRows,
options.sheetName
)
chunks.push(StructuredDataChunker.createChunk(chunkContent, chunkStartRow, i - 1))
currentChunkRows = []
currentTokenEstimate = 0
chunkStartRow = i
}
currentChunkRows.push(row)
currentTokenEstimate += rowTokens
}
if (currentChunkRows.length > 0) {
const chunkContent = StructuredDataChunker.formatChunk(
headerLine,
currentChunkRows,
options.sheetName
)
chunks.push(StructuredDataChunker.createChunk(chunkContent, chunkStartRow, lines.length - 1))
}
logger.info(`Created ${chunks.length} chunks from ${lines.length} rows of structured data`)
return chunks
}
private static formatChunk(headerLine: string, rows: string[], sheetName?: string): string {
let content = ''
if (sheetName) {
content += `=== ${sheetName} ===\n\n`
}
if (DEFAULT_CONFIG.INCLUDE_HEADERS_IN_EACH_CHUNK) {
content += `Headers: ${headerLine}\n`
content += `${'-'.repeat(Math.min(80, headerLine.length))}\n`
}
content += rows.join('\n')
content += `\n\n[${rows.length} rows of data]`
return content
}
private static createChunk(content: string, startRow: number, endRow: number): Chunk {
return {
text: content,
tokenCount: estimateStructuredTokens(content),
metadata: {
startIndex: startRow,
endIndex: endRow,
},
}
}
private static estimateStructuredTokensPerRow(sampleRows: string[]): number {
if (sampleRows.length === 0) return 50
const totalTokens = sampleRows.reduce((sum, row) => sum + estimateStructuredTokens(row), 0)
return Math.ceil(totalTokens / sampleRows.length)
}
private static calculateOptimalRowsPerChunk(
tokensPerRow: number,
targetChunkSize: number
): number {
const optimal = Math.floor(targetChunkSize / tokensPerRow)
return Math.min(
Math.max(optimal, DEFAULT_CONFIG.MIN_ROWS_PER_CHUNK),
DEFAULT_CONFIG.MAX_ROWS_PER_CHUNK
)
}
static isStructuredData(content: string, mimeType?: string): boolean {
if (mimeType) {
const structuredMimeTypes = [
'text/csv',
'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet',
'application/vnd.ms-excel',
'text/tab-separated-values',
]
if (structuredMimeTypes.includes(mimeType)) {
return true
}
}
const lines = content.split('\n').slice(0, 10)
if (lines.length < 2) return false
const delimiters = [',', '\t', '|']
for (const delimiter of delimiters) {
const escaped = delimiter.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')
const counts = lines.map((line) => (line.match(new RegExp(escaped, 'g')) || []).length)
const avgCount = counts.reduce((a, b) => a + b, 0) / counts.length
const tolerance = Math.max(1, Math.ceil(avgCount * 0.2))
if (avgCount > 2 && counts.every((c) => Math.abs(c - avgCount) <= tolerance)) {
return true
}
}
return false
}
}
+532
View File
@@ -0,0 +1,532 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import { TextChunker } from './text-chunker'
describe('TextChunker', () => {
describe('basic chunking', () => {
it.concurrent('should return empty array for empty text', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const chunks = await chunker.chunk('')
expect(chunks).toEqual([])
})
it.concurrent('should return empty array for whitespace-only text', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const chunks = await chunker.chunk(' \n\n ')
expect(chunks).toEqual([])
})
it.concurrent('should return single chunk for small text', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'This is a short text.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
})
it.concurrent('should include token count in chunk metadata', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Hello world'
const chunks = await chunker.chunk(text)
expect(chunks[0].tokenCount).toBe(3)
})
})
describe('chunk size limits (tokens)', () => {
it.concurrent('should split text that exceeds chunk size', async () => {
const chunker = new TextChunker({ chunkSize: 50 })
const text = 'This is a test sentence. '.repeat(20)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should respect chunk size limit', async () => {
const chunkSize = 50
const chunker = new TextChunker({ chunkSize })
const text = 'This is a test sentence. '.repeat(20)
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.tokenCount).toBeLessThanOrEqual(chunkSize + 5)
}
})
it.concurrent('should create more chunks with smaller chunk size', async () => {
const text = 'This is a test sentence. '.repeat(20)
const largeChunker = new TextChunker({ chunkSize: 200 })
const smallChunker = new TextChunker({ chunkSize: 50 })
const largeChunks = await largeChunker.chunk(text)
const smallChunks = await smallChunker.chunk(text)
expect(smallChunks.length).toBeGreaterThan(largeChunks.length)
})
})
describe('overlap (tokens)', () => {
it.concurrent('should not add overlap when overlap is 0', async () => {
const chunker = new TextChunker({ chunkSize: 20, chunkOverlap: 0 })
const text = 'First sentence here. Second sentence here. Third sentence here.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
const firstChunkEnd = chunks[0].text.slice(-10)
const secondChunkStart = chunks[1].text.slice(0, 10)
expect(chunks[1].text.startsWith(firstChunkEnd)).toBe(false)
}
})
it.concurrent('should add overlap from previous chunk', async () => {
const chunker = new TextChunker({ chunkSize: 20, chunkOverlap: 50 })
const text =
'First paragraph with some content here.\n\nSecond paragraph with different content here.\n\nThird paragraph with more content.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
expect(chunks[1].text.length).toBeGreaterThan(0)
}
})
it.concurrent('should handle overlap larger than previous chunk', async () => {
const chunker = new TextChunker({ chunkSize: 10, chunkOverlap: 500 })
const text = 'Short. Another short.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('hierarchical splitting', () => {
it.concurrent('should prefer splitting on paragraph boundaries', async () => {
const chunker = new TextChunker({ chunkSize: 30 })
const text =
'First paragraph content here.\n\nSecond paragraph content here.\n\nThird paragraph content here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThanOrEqual(1)
})
it.concurrent('should split on sentences when paragraphs are too large', async () => {
const chunker = new TextChunker({ chunkSize: 20 })
const text =
'First sentence. Second sentence. Third sentence. Fourth sentence. Fifth sentence.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should split on words as last resort', async () => {
const chunker = new TextChunker({ chunkSize: 5 })
const text = 'Supercalifragilisticexpialidocious is a very long word that needs splitting.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle markdown headings', async () => {
const chunker = new TextChunker({ chunkSize: 30 })
const text =
'# Heading 1\n\nContent under heading 1.\n\n## Heading 2\n\nContent under heading 2.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThanOrEqual(1)
})
})
describe('text cleaning', () => {
it.concurrent('should normalize Windows line endings', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Line 1\r\nLine 2\r\nLine 3'
const chunks = await chunker.chunk(text)
expect(chunks[0].text).not.toContain('\r')
})
it.concurrent('should collapse multiple spaces', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Word1 Word2 Word3'
const chunks = await chunker.chunk(text)
expect(chunks[0].text).not.toContain(' ')
})
it.concurrent('should limit consecutive newlines', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Para 1\n\n\n\n\nPara 2'
const chunks = await chunker.chunk(text)
expect(chunks[0].text).not.toContain('\n\n\n')
})
})
describe('chunk metadata', () => {
it.concurrent('should include startIndex and endIndex', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'This is test content.'
const chunks = await chunker.chunk(text)
expect(chunks[0].metadata.startIndex).toBeDefined()
expect(chunks[0].metadata.endIndex).toBeDefined()
expect(chunks[0].metadata.startIndex).toBe(0)
})
it.concurrent('should have sequential indices across chunks', async () => {
const chunker = new TextChunker({ chunkSize: 20, chunkOverlap: 0 })
const text = 'First part of text. Second part of text. Third part of text.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
for (let i = 1; i < chunks.length; i++) {
expect(chunks[i].metadata.startIndex).toBeGreaterThanOrEqual(0)
}
}
})
})
describe('default values', () => {
it.concurrent('should use default chunkSize of 1024 tokens', async () => {
const chunker = new TextChunker({})
const text = 'Word '.repeat(400)
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
})
it.concurrent('should use default minCharactersPerChunk of 100', async () => {
const chunker = new TextChunker({ chunkSize: 10 })
const text = 'This is a longer sentence with more content. '.repeat(5)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should use default overlap of 0', async () => {
const chunker = new TextChunker({ chunkSize: 20 })
const text = 'First sentence here. Second sentence here.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
const firstEnd = chunks[0].text.slice(-5)
expect(chunks[1].text.startsWith(firstEnd)).toBe(false)
}
})
})
describe('edge cases', () => {
it.concurrent('should handle very long text', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'This is a sentence. '.repeat(1000)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
const totalLength = chunks.reduce((sum, c) => sum + c.text.length, 0)
expect(totalLength).toBeGreaterThan(0)
})
it.concurrent('should handle text with only special characters', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = '!@#$%^&*()_+-=[]{}|;:,.<>?'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
})
it.concurrent('should handle unicode text', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = '这是中文测试。日本語テスト。한국어 테스트.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('中文')
})
it.concurrent('should not lose any content during chunking', async () => {
const chunker = new TextChunker({ chunkSize: 30, chunkOverlap: 0 })
const originalText =
'The quick brown fox jumps over the lazy dog. Pack my box with five dozen liquor jugs.'
const chunks = await chunker.chunk(originalText)
const allText = chunks.map((c) => c.text).join(' ')
expect(allText).toContain('quick')
expect(allText).toContain('fox')
expect(allText).toContain('lazy')
expect(allText).toContain('dog')
})
})
describe('boundary conditions', () => {
it.concurrent('should handle text exactly at chunk size boundary', async () => {
const chunker = new TextChunker({ chunkSize: 10 })
const text = 'A'.repeat(40)
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].tokenCount).toBe(10)
})
it.concurrent('should handle text one token over chunk size', async () => {
const chunker = new TextChunker({ chunkSize: 10 })
const text = 'A'.repeat(44)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThanOrEqual(1)
})
it.concurrent('should handle chunkSize of 1 token', async () => {
const chunker = new TextChunker({ chunkSize: 1 })
const text = 'Hello world test'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle overlap equal to half of chunk size', async () => {
const chunker = new TextChunker({ chunkSize: 20, chunkOverlap: 10 })
const text = 'First sentence here. Second sentence here. Third sentence here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should clamp overlap to max 50% of chunk size', async () => {
const chunker = new TextChunker({ chunkSize: 20, chunkOverlap: 60 })
const text = 'First paragraph here.\n\nSecond paragraph here.\n\nThird paragraph here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle zero minCharactersPerChunk', async () => {
const chunker = new TextChunker({ chunkSize: 10, minCharactersPerChunk: 0 })
const text = 'A B C'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('encoding and special characters', () => {
it.concurrent('should handle emoji characters', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Hello 👋 World 🌍! This has emojis 🎉🎊🎈'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toContain('👋')
expect(chunks[0].text).toContain('🌍')
})
it.concurrent('should handle mixed language text', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'English text. 中文文本。日本語テキスト。한국어 텍스트. العربية'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('English')
expect(chunks[0].text).toContain('中文')
expect(chunks[0].text).toContain('日本語')
})
it.concurrent('should handle RTL text (Arabic/Hebrew)', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'مرحبا بالعالم - שלום עולם - Hello World'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
expect(chunks[0].text).toContain('مرحبا')
expect(chunks[0].text).toContain('שלום')
})
it.concurrent('should handle null characters in text', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Hello\0World\0Test'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle combining diacritics', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'cafe\u0301 resume\u0301 naive\u0308'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle zero-width characters', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Hello\u200B\u200C\u200DWorld'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle old Mac line endings (\\r)', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Line 1\rLine 2\rLine 3'
const chunks = await chunker.chunk(text)
expect(chunks[0].text).not.toContain('\r')
})
})
describe('large inputs', () => {
it.concurrent('should handle 10,000 word document', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'This is a test sentence with several words. '.repeat(2000)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
const totalChars = chunks.reduce((sum, c) => sum + c.text.length, 0)
expect(totalChars).toBeGreaterThan(0)
})
it.concurrent('should handle 1MB of text', async () => {
const chunker = new TextChunker({ chunkSize: 500 })
const text = 'Lorem ipsum dolor sit amet. '.repeat(40000)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle very long single line', async () => {
const chunker = new TextChunker({ chunkSize: 50 })
const text = 'Word'.repeat(10000)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should handle many short paragraphs', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = Array(500)
.fill(0)
.map((_, i) => `Paragraph ${i}.`)
.join('\n\n')
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
})
describe('markdown and code handling', () => {
it.concurrent('should handle code blocks', async () => {
const chunker = new TextChunker({ chunkSize: 50 })
const text = `
# Code Example
\`\`\`javascript
function hello() {
console.log("Hello World");
}
\`\`\`
Some explanation text after the code.
`
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle nested lists', async () => {
const chunker = new TextChunker({ chunkSize: 50 })
const text = `
- Item 1
- Nested 1.1
- Nested 1.2
- Deep nested 1.2.1
- Item 2
- Nested 2.1
`
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle markdown tables', async () => {
const chunker = new TextChunker({ chunkSize: 50 })
const text = `
| Header 1 | Header 2 | Header 3 |
|----------|----------|----------|
| Cell 1 | Cell 2 | Cell 3 |
| Cell 4 | Cell 5 | Cell 6 |
`
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should handle inline code', async () => {
const chunker = new TextChunker({ chunkSize: 100 })
const text = 'Use `const` for constants and `let` for variables. Call `myFunction()` here.'
const chunks = await chunker.chunk(text)
expect(chunks[0].text).toContain('`const`')
})
})
describe('separator hierarchy', () => {
it.concurrent('should split on horizontal rules', async () => {
const chunker = new TextChunker({ chunkSize: 30 })
const text = 'Section 1 content here.\n---\nSection 2 content here.\n---\nSection 3 content.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should split on question marks', async () => {
const chunker = new TextChunker({ chunkSize: 20 })
const text = 'What is this? How does it work? Why is it important? When to use it?'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should split on exclamation marks', async () => {
const chunker = new TextChunker({ chunkSize: 20 })
const text = 'Amazing! Incredible! Fantastic! Wonderful! Great!'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should split on semicolons', async () => {
const chunker = new TextChunker({ chunkSize: 20 })
const text = 'First clause; second clause; third clause; fourth clause'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
})
describe('chunk index accuracy', () => {
it.concurrent('should have non-negative indices', async () => {
const chunker = new TextChunker({ chunkSize: 30, chunkOverlap: 10 })
const text = 'First part. Second part. Third part. Fourth part. Fifth part.'
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.metadata.startIndex).toBeGreaterThanOrEqual(0)
expect(chunk.metadata.endIndex).toBeGreaterThanOrEqual(chunk.metadata.startIndex)
}
})
it.concurrent('should have endIndex greater than or equal to startIndex', async () => {
const chunker = new TextChunker({ chunkSize: 20 })
const text = 'Multiple sentences here. Another one here. And another. And more.'
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.metadata.endIndex).toBeGreaterThanOrEqual(chunk.metadata.startIndex)
}
})
})
})
+108
View File
@@ -0,0 +1,108 @@
import type { Chunk, ChunkerOptions } from '@/lib/chunkers/types'
import {
addOverlap,
buildChunks,
cleanText,
estimateTokens,
resolveChunkerOptions,
splitAtWordBoundaries,
tokensToChars,
} from '@/lib/chunkers/utils'
export class TextChunker {
private readonly chunkSize: number
private readonly chunkOverlap: number
private readonly separators = [
'\n---\n',
'\n***\n',
'\n___\n',
'\n# ',
'\n## ',
'\n### ',
'\n#### ',
'\n##### ',
'\n###### ',
'\n\n',
'\n',
'. ',
'! ',
'? ',
'; ',
', ',
' ',
]
constructor(options: ChunkerOptions = {}) {
const resolved = resolveChunkerOptions(options)
this.chunkSize = resolved.chunkSize
this.chunkOverlap = resolved.chunkOverlap
}
private splitRecursively(text: string, separatorIndex = 0): string[] {
const tokenCount = estimateTokens(text)
if (tokenCount <= this.chunkSize) {
return text.trim() ? [text] : []
}
if (separatorIndex >= this.separators.length) {
const chunkSizeChars = tokensToChars(this.chunkSize)
return splitAtWordBoundaries(text, chunkSizeChars)
}
const separator = this.separators[separatorIndex]
const parts = text.split(separator).filter((part) => part.trim())
if (parts.length <= 1) {
return this.splitRecursively(text, separatorIndex + 1)
}
const chunks: string[] = []
let currentChunk = ''
for (const part of parts) {
const testChunk = currentChunk + (currentChunk ? separator : '') + part
if (estimateTokens(testChunk) <= this.chunkSize) {
currentChunk = testChunk
} else {
if (currentChunk.trim()) {
chunks.push(currentChunk.trim())
}
if (estimateTokens(part) > this.chunkSize) {
const subChunks = this.splitRecursively(part, separatorIndex + 1)
for (const subChunk of subChunks) {
chunks.push(subChunk)
}
currentChunk = ''
} else {
currentChunk = part
}
}
}
if (currentChunk.trim()) {
chunks.push(currentChunk.trim())
}
return chunks
}
async chunk(text: string): Promise<Chunk[]> {
if (!text?.trim()) {
return []
}
const cleaned = cleanText(text)
let chunks = this.splitRecursively(cleaned)
if (this.chunkOverlap > 0) {
const overlapChars = tokensToChars(this.chunkOverlap)
chunks = addOverlap(chunks, overlapChars)
}
return buildChunks(chunks, this.chunkOverlap)
}
}
+236
View File
@@ -0,0 +1,236 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import { TokenChunker } from './token-chunker'
describe('TokenChunker', () => {
describe('empty and whitespace input', () => {
it.concurrent('should return empty array for empty string', async () => {
const chunker = new TokenChunker({ chunkSize: 100 })
const chunks = await chunker.chunk('')
expect(chunks).toEqual([])
})
it.concurrent('should return empty array for whitespace-only input', async () => {
const chunker = new TokenChunker({ chunkSize: 100 })
const chunks = await chunker.chunk(' \n\n\t ')
expect(chunks).toEqual([])
})
})
describe('small content', () => {
it.concurrent('should return single chunk when content fits within chunkSize', async () => {
const chunker = new TokenChunker({ chunkSize: 100 })
const text = 'This is a short text.'
const chunks = await chunker.chunk(text)
expect(chunks).toHaveLength(1)
expect(chunks[0].text).toBe(text)
})
})
describe('token count accuracy', () => {
it.concurrent('should compute tokenCount as Math.ceil(text.length / 4)', async () => {
const chunker = new TokenChunker({ chunkSize: 100 })
const text = 'Hello world'
const chunks = await chunker.chunk(text)
expect(chunks[0].tokenCount).toBe(Math.ceil(text.length / 4))
})
it.concurrent('should compute tokenCount correctly for longer text', async () => {
const chunker = new TokenChunker({ chunkSize: 100 })
const text = 'The quick brown fox jumps over the lazy dog.'
const chunks = await chunker.chunk(text)
expect(chunks[0].tokenCount).toBe(11)
})
})
describe('chunk metadata', () => {
it.concurrent(
'should include text, tokenCount, and metadata with startIndex and endIndex',
async () => {
const chunker = new TokenChunker({ chunkSize: 100 })
const text = 'Some test content here.'
const chunks = await chunker.chunk(text)
expect(chunks[0]).toHaveProperty('text')
expect(chunks[0]).toHaveProperty('tokenCount')
expect(chunks[0].metadata).toHaveProperty('startIndex')
expect(chunks[0].metadata).toHaveProperty('endIndex')
expect(chunks[0].metadata.startIndex).toBe(0)
expect(chunks[0].metadata.endIndex).toBeGreaterThan(0)
}
)
it.concurrent('should have non-negative indices across all chunks', async () => {
const chunker = new TokenChunker({ chunkSize: 20, chunkOverlap: 0 })
const text = 'First part of the text. Second part of the text. Third part of the text.'
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.metadata.startIndex).toBeGreaterThanOrEqual(0)
expect(chunk.metadata.endIndex).toBeGreaterThanOrEqual(chunk.metadata.startIndex)
}
})
})
describe('respects chunk size', () => {
it.concurrent('should not produce chunks exceeding chunkSize tokens', async () => {
const chunkSize = 50
const chunker = new TokenChunker({ chunkSize })
const text = 'This is a test sentence with several words. '.repeat(30)
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
expect(chunk.tokenCount).toBeLessThanOrEqual(chunkSize)
}
})
})
describe('splitting behavior', () => {
it.concurrent('should produce multiple chunks for long text', async () => {
const chunker = new TokenChunker({ chunkSize: 50 })
const text = 'This is a test sentence. '.repeat(30)
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
})
it.concurrent('should create more chunks with smaller chunkSize', async () => {
const text = 'This is a test sentence with content. '.repeat(20)
const largeChunker = new TokenChunker({ chunkSize: 200 })
const smallChunker = new TokenChunker({ chunkSize: 50 })
const largeChunks = await largeChunker.chunk(text)
const smallChunks = await smallChunker.chunk(text)
expect(smallChunks.length).toBeGreaterThan(largeChunks.length)
})
})
describe('sliding window overlap', () => {
it.concurrent('should produce more chunks with overlap than without', async () => {
const text =
'Alpha bravo charlie delta echo foxtrot golf hotel india juliet kilo lima mike november oscar papa quebec romeo sierra tango uniform victor whiskey xray yankee zulu. '.repeat(
5
)
const withOverlap = new TokenChunker({ chunkSize: 30, chunkOverlap: 10 })
const withoutOverlap = new TokenChunker({ chunkSize: 30, chunkOverlap: 0 })
const overlapChunks = await withOverlap.chunk(text)
const noOverlapChunks = await withoutOverlap.chunk(text)
expect(overlapChunks.length).toBeGreaterThan(noOverlapChunks.length)
})
it.concurrent('should not share text between chunks when chunkOverlap is 0', async () => {
const chunker = new TokenChunker({ chunkSize: 20, chunkOverlap: 0 })
const text =
'First sentence here. Second sentence here. Third sentence here. Fourth sentence here.'
const chunks = await chunker.chunk(text)
if (chunks.length > 1) {
const firstChunkEnd = chunks[0].text.slice(-10)
expect(chunks[1].text.startsWith(firstChunkEnd)).toBe(false)
}
})
})
describe('overlap clamped to 50%', () => {
it.concurrent('should still work when overlap is set >= chunkSize', async () => {
const chunker = new TokenChunker({ chunkSize: 20, chunkOverlap: 100 })
const text =
'First paragraph content here. Second paragraph content here. Third paragraph here.'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(0)
})
it.concurrent('should clamp overlap to 50% of chunkSize', async () => {
const chunkerClamped = new TokenChunker({ chunkSize: 20, chunkOverlap: 100 })
const chunkerHalf = new TokenChunker({ chunkSize: 20, chunkOverlap: 10 })
const text =
'Word one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen sixteen seventeen eighteen nineteen twenty. '.repeat(
5
)
const clampedChunks = await chunkerClamped.chunk(text)
const halfChunks = await chunkerHalf.chunk(text)
expect(clampedChunks.length).toBe(halfChunks.length)
})
})
describe('word boundary snapping', () => {
it.concurrent('should produce trimmed chunks without leading or trailing spaces', async () => {
const chunker = new TokenChunker({ chunkSize: 20 })
const text =
'the cat sat on the mat and the dog ran fast over the big red fox and then the bird flew high up in the clear blue sky above the green hill'
const chunks = await chunker.chunk(text)
expect(chunks.length).toBeGreaterThan(1)
for (const chunk of chunks) {
const trimmed = chunk.text.trim()
expect(trimmed).toBe(chunk.text)
expect(trimmed.length).toBeGreaterThan(0)
}
})
it.concurrent('should produce chunks that start and end on word boundaries', async () => {
const chunker = new TokenChunker({ chunkSize: 15 })
const text =
'The quick brown fox jumps over the lazy dog and then runs away quickly into the forest'
const chunks = await chunker.chunk(text)
for (const chunk of chunks) {
const trimmed = chunk.text.trim()
expect(trimmed).toBe(chunk.text)
}
})
})
describe('consistent coverage', () => {
it.concurrent('should represent all content from original text across chunks', async () => {
const chunker = new TokenChunker({ chunkSize: 30, chunkOverlap: 0 })
const originalText =
'The quick brown fox jumps over the lazy dog. Pack my box with five dozen liquor jugs.'
const chunks = await chunker.chunk(originalText)
const allText = chunks.map((c) => c.text).join(' ')
expect(allText).toContain('quick')
expect(allText).toContain('fox')
expect(allText).toContain('lazy')
expect(allText).toContain('dog')
expect(allText).toContain('liquor')
expect(allText).toContain('jugs')
})
it.concurrent('should preserve all words across chunks for longer text', async () => {
const chunker = new TokenChunker({ chunkSize: 20, chunkOverlap: 0 })
const words = [
'alpha',
'bravo',
'charlie',
'delta',
'echo',
'foxtrot',
'golf',
'hotel',
'india',
'juliet',
]
const originalText = `${words.join(' is a word and ')} is also a word.`
const chunks = await chunker.chunk(originalText)
const combined = chunks.map((c) => c.text).join(' ')
for (const word of words) {
expect(combined).toContain(word)
}
})
})
})
+54
View File
@@ -0,0 +1,54 @@
import { createLogger } from '@sim/logger'
import type { Chunk, ChunkerOptions } from '@/lib/chunkers/types'
import {
buildChunks,
cleanText,
estimateTokens,
resolveChunkerOptions,
splitAtWordBoundaries,
tokensToChars,
} from '@/lib/chunkers/utils'
const logger = createLogger('TokenChunker')
export class TokenChunker {
private readonly chunkSize: number
private readonly chunkOverlap: number
private readonly minCharactersPerChunk: number
constructor(options: ChunkerOptions = {}) {
const resolved = resolveChunkerOptions(options)
this.chunkSize = resolved.chunkSize
this.chunkOverlap = resolved.chunkOverlap
this.minCharactersPerChunk = resolved.minCharactersPerChunk
}
async chunk(content: string): Promise<Chunk[]> {
if (!content?.trim()) {
return []
}
const cleaned = cleanText(content)
if (estimateTokens(cleaned) <= this.chunkSize) {
logger.info('Content fits in single chunk')
return buildChunks([cleaned], 0)
}
const chunkSizeChars = tokensToChars(this.chunkSize)
const overlapChars = tokensToChars(this.chunkOverlap)
const stepChars = this.chunkOverlap > 0 ? chunkSizeChars - overlapChars : undefined
const rawChunks = splitAtWordBoundaries(cleaned, chunkSizeChars, stepChars)
const filtered =
rawChunks.length > 1
? rawChunks.filter((c) => c.length >= this.minCharactersPerChunk)
: rawChunks
const chunks = filtered.length > 0 ? filtered : rawChunks
logger.info(`Chunked into ${chunks.length} token-based chunks`)
return buildChunks(chunks, this.chunkOverlap)
}
}
+78
View File
@@ -0,0 +1,78 @@
/**
* Units:
* - chunkSize/chunkOverlap: TOKENS (1 token ≈ 4 characters)
* - minCharactersPerChunk: CHARACTERS
*/
export interface ChunkerOptions {
chunkSize?: number
chunkOverlap?: number
minCharactersPerChunk?: number
}
export interface Chunk {
text: string
tokenCount: number
metadata: {
startIndex: number
endIndex: number
}
}
export interface StructuredDataOptions extends ChunkerOptions {
headers?: string[]
totalRows?: number
sheetName?: string
}
export interface DocChunk {
text: string
tokenCount: number
sourceDocument: string
headerLink: string
headerText: string
headerLevel: number
embedding: number[]
embeddingModel: string
metadata: {
sourceUrl?: string
headers?: string[]
title?: string
startIndex: number
endIndex: number
}
}
export interface DocsChunkerOptions extends ChunkerOptions {
baseUrl?: string
}
export type ChunkingStrategy = 'auto' | 'text' | 'regex' | 'recursive' | 'sentence' | 'token'
export type RecursiveRecipe = 'plain' | 'markdown' | 'code'
export interface StrategyOptions {
pattern?: string
separators?: string[]
recipe?: RecursiveRecipe
strictBoundaries?: boolean
}
export interface SentenceChunkerOptions extends ChunkerOptions {
minSentencesPerChunk?: number
}
export interface RecursiveChunkerOptions extends ChunkerOptions {
separators?: string[]
recipe?: RecursiveRecipe
}
export interface RegexChunkerOptions extends ChunkerOptions {
pattern: string
/**
* When true, each regex match becomes its own chunk and small adjacent
* segments are not merged together. Overlap is also disabled. Useful for
* structural inputs where boundaries (e.g. one record per match) must be
* preserved exactly.
*/
strictBoundaries?: boolean
}
+217
View File
@@ -0,0 +1,217 @@
/**
* @vitest-environment node
*/
import { describe, expect, it } from 'vitest'
import {
addOverlap,
buildChunks,
cleanText,
estimateTokens,
resolveChunkerOptions,
splitAtWordBoundaries,
tokensToChars,
} from './utils'
describe('estimateTokens', () => {
it('returns 0 for empty string', () => {
expect(estimateTokens('')).toBe(0)
})
it('returns 0 for whitespace-only string', () => {
expect(estimateTokens(' ')).toBe(0)
})
it('returns 0 for null or undefined via optional chaining', () => {
expect(estimateTokens(null as unknown as string)).toBe(0)
expect(estimateTokens(undefined as unknown as string)).toBe(0)
})
it('returns Math.ceil(text.length / 4) for normal text', () => {
const text = 'Hello world'
expect(estimateTokens(text)).toBe(Math.ceil(text.length / 4))
})
it('estimates "Hello world" (11 chars) as 3 tokens', () => {
expect(estimateTokens('Hello world')).toBe(3)
})
})
describe('tokensToChars', () => {
it('returns tokens * 4', () => {
expect(tokensToChars(1)).toBe(4)
expect(tokensToChars(5)).toBe(20)
})
it('converts 10 tokens to 40 chars', () => {
expect(tokensToChars(10)).toBe(40)
})
})
describe('cleanText', () => {
it('normalizes \\r\\n to \\n', () => {
expect(cleanText('hello\r\nworld')).toBe('hello\nworld')
})
it('normalizes \\r to \\n', () => {
expect(cleanText('hello\rworld')).toBe('hello\nworld')
})
it('collapses 3+ newlines to \\n\\n', () => {
expect(cleanText('hello\n\n\n\nworld')).toBe('hello\n\nworld')
})
it('replaces tabs with spaces', () => {
expect(cleanText('hello\tworld')).toBe('hello world')
})
it('collapses multiple spaces to single space', () => {
expect(cleanText('hello world')).toBe('hello world')
})
it('trims leading and trailing whitespace', () => {
expect(cleanText(' hello world ')).toBe('hello world')
})
})
describe('addOverlap', () => {
it('returns unchanged chunks when overlapChars <= 0', () => {
const chunks = ['chunk one', 'chunk two']
expect(addOverlap(chunks, 0)).toEqual(chunks)
expect(addOverlap(chunks, -5)).toEqual(chunks)
})
it('returns unchanged chunks when only 1 chunk', () => {
const chunks = ['only chunk']
expect(addOverlap(chunks, 10)).toEqual(chunks)
})
it('prepends tail of previous chunk to next chunk with overlap > 0', () => {
const chunks = ['first chunk here', 'second chunk here']
const result = addOverlap(chunks, 10)
expect(result[0]).toBe('first chunk here')
expect(result[1]).toContain('second chunk here')
expect(result[1].length).toBeGreaterThan('second chunk here'.length)
})
it('joins overlap text with space', () => {
const chunks = ['first chunk here', 'second chunk here']
const result = addOverlap(chunks, 10)
expect(result[1]).toContain('here second')
})
it('snaps overlap to word boundary', () => {
const chunks = ['hello beautiful world', 'next chunk']
const result = addOverlap(chunks, 15)
expect(result[1]).toBe('beautiful world next chunk')
})
})
describe('splitAtWordBoundaries', () => {
it('returns single element for short text', () => {
const result = splitAtWordBoundaries('short text', 100)
expect(result).toHaveLength(1)
expect(result[0]).toBe('short text')
})
it('produces multiple chunks for long text', () => {
const text = 'word '.repeat(100).trim()
const result = splitAtWordBoundaries(text, 20)
expect(result.length).toBeGreaterThan(1)
})
it('respects chunk size limit', () => {
const text = 'word '.repeat(100).trim()
const chunkSize = 25
const result = splitAtWordBoundaries(text, chunkSize)
for (const chunk of result) {
expect(chunk.length).toBeLessThanOrEqual(chunkSize)
}
})
it('does not break mid-word', () => {
const text = 'internationalization globalization modernization'
const result = splitAtWordBoundaries(text, 25)
for (const chunk of result) {
expect(chunk).not.toMatch(/^\S+\s\S+$.*\S$/)
const words = chunk.split(' ')
for (const word of words) {
expect(text).toContain(word)
}
}
})
it('produces overlapping chunks with stepChars < chunkSizeChars', () => {
const text = 'one two three four five six seven eight nine ten'
const result = splitAtWordBoundaries(text, 20, 10)
expect(result.length).toBeGreaterThan(1)
const combined = result.join(' ')
for (const word of text.split(' ')) {
expect(combined).toContain(word)
}
})
it('ensures step is at least 1 to prevent infinite loops', () => {
const text = 'hello world test'
const result = splitAtWordBoundaries(text, 10, 0)
expect(result.length).toBeGreaterThan(0)
})
})
describe('buildChunks', () => {
it('creates Chunk objects with text, tokenCount, and metadata', () => {
const texts = ['hello world', 'foo bar']
const chunks = buildChunks(texts, 0)
for (const chunk of chunks) {
expect(chunk).toHaveProperty('text')
expect(chunk).toHaveProperty('tokenCount')
expect(chunk).toHaveProperty('metadata')
expect(chunk.metadata).toHaveProperty('startIndex')
expect(chunk.metadata).toHaveProperty('endIndex')
}
})
it('sets metadata with startIndex and endIndex', () => {
const texts = ['chunk one', 'chunk two']
const chunks = buildChunks(texts, 0)
expect(typeof chunks[0].metadata.startIndex).toBe('number')
expect(typeof chunks[0].metadata.endIndex).toBe('number')
})
it('sets startIndex of first chunk to 0', () => {
const texts = ['first chunk', 'second chunk']
const chunks = buildChunks(texts, 0)
expect(chunks[0].metadata.startIndex).toBe(0)
})
it('produces contiguous chunks with overlapTokens=0', () => {
const texts = ['hello world', 'foo bar baz']
const chunks = buildChunks(texts, 0)
expect(chunks[0].metadata.endIndex).toBe(chunks[1].metadata.startIndex)
})
})
describe('resolveChunkerOptions', () => {
it('applies defaults: chunkSize=1024, chunkOverlap=0, minCharactersPerChunk=100', () => {
const result = resolveChunkerOptions({})
expect(result.chunkSize).toBe(1024)
expect(result.chunkOverlap).toBe(0)
expect(result.minCharactersPerChunk).toBe(100)
})
it('clamps overlap to max 50% of chunkSize', () => {
const result = resolveChunkerOptions({ chunkSize: 100, chunkOverlap: 80 })
expect(result.chunkOverlap).toBe(50)
})
it('respects provided values when within limits', () => {
const result = resolveChunkerOptions({
chunkSize: 500,
chunkOverlap: 100,
minCharactersPerChunk: 50,
})
expect(result.chunkSize).toBe(500)
expect(result.chunkOverlap).toBe(100)
expect(result.minCharactersPerChunk).toBe(50)
})
})
+143
View File
@@ -0,0 +1,143 @@
import type { Chunk } from '@/lib/chunkers/types'
/** 1 token ≈ 4 characters for English text */
export function estimateTokens(text: string): number {
if (!text?.trim()) return 0
return Math.ceil(text.length / 4)
}
export function tokensToChars(tokens: number): number {
return tokens * 4
}
export function cleanText(text: string): string {
return text
.replace(/\r\n/g, '\n')
.replace(/\r/g, '\n')
.replace(/\n{3,}/g, '\n\n')
.replace(/\t/g, ' ')
.replace(/ {2,}/g, ' ')
.trim()
}
export function addOverlap(chunks: string[], overlapChars: number): string[] {
if (overlapChars <= 0 || chunks.length <= 1) {
return chunks
}
const result: string[] = []
for (let i = 0; i < chunks.length; i++) {
let chunk = chunks[i]
if (i > 0) {
const prevChunk = chunks[i - 1]
const overlapLength = Math.min(overlapChars, prevChunk.length)
const overlapText = prevChunk.slice(-overlapLength)
const wordBoundaryMatch = overlapText.match(/^\s*\S/)
const cleanOverlap = wordBoundaryMatch
? overlapText.slice(overlapText.indexOf(wordBoundaryMatch[0].trim()))
: overlapText
if (cleanOverlap.trim()) {
chunk = `${cleanOverlap.trim()} ${chunk}`
}
}
result.push(chunk)
}
return result
}
/**
* When stepChars is provided (< chunkSizeChars), produces overlapping chunks
* using a sliding window where chunks stay within the size limit.
*/
export function splitAtWordBoundaries(
text: string,
chunkSizeChars: number,
stepChars?: number
): string[] {
const parts: string[] = []
let pos = 0
while (pos < text.length) {
let end = Math.min(pos + chunkSizeChars, text.length)
if (end < text.length) {
const lastSpace = text.lastIndexOf(' ', end)
if (lastSpace > pos) {
end = lastSpace
}
}
const part = text.slice(pos, end).trim()
if (part) {
parts.push(part)
}
if (stepChars !== undefined) {
// Sliding window: advance by step for predictable overlap
const nextPos = pos + Math.max(1, stepChars)
if (nextPos >= text.length) break
pos = nextPos
} else {
// Non-overlapping: advance from end of extracted content
if (end >= text.length) break
pos = end
}
while (pos < text.length && text[pos] === ' ') pos++
}
return parts
}
export function buildChunks(texts: string[], overlapTokens: number): Chunk[] {
let previousEndIndex = 0
const overlapChars = tokensToChars(overlapTokens)
return texts.map((text, index) => {
let startIndex: number
let actualContentLength: number
if (index === 0 || overlapTokens <= 0) {
startIndex = previousEndIndex
actualContentLength = text.length
} else {
const prevChunk = texts[index - 1]
const overlapLength = Math.min(overlapChars, prevChunk.length, text.length)
startIndex = previousEndIndex - overlapLength
actualContentLength = text.length - overlapLength
}
const safeStart = Math.max(0, startIndex)
const endIndex = safeStart + Math.max(0, actualContentLength)
previousEndIndex = endIndex
return {
text,
tokenCount: estimateTokens(text),
metadata: {
startIndex: safeStart,
endIndex,
},
}
})
}
export function resolveChunkerOptions(options: {
chunkSize?: number
chunkOverlap?: number
minCharactersPerChunk?: number
}): { chunkSize: number; chunkOverlap: number; minCharactersPerChunk: number } {
const chunkSize = options.chunkSize ?? 1024
const maxOverlap = Math.floor(chunkSize * 0.5)
return {
chunkSize,
chunkOverlap: Math.min(options.chunkOverlap ?? 0, maxOverlap),
minCharactersPerChunk: options.minCharactersPerChunk ?? 100,
}
}