Skip to content

Commit aba6fe2

Browse files
waleedlatif1waleedlatif1
committed
fix(knowledge): bound JSON/YAML chunker expansion
JsonYamlChunker re-parsed and re-serialized document content with no expansion limit. ChunkBudget only counts emitted chunks, so every parse and full-object stringify ran before it could fire — a small aliased YAML source expands to tens of MB, and the same content was parsed twice because isStructuredData and chunkJsonYaml each parsed it. - Measure the parsed value with the shared measureYamlExpansion guard before anything materializes it, and skip parsing entirely when the source is already larger than the ceiling - Size the ceiling to the most text the chunker could ever emit (maxChunks x chunkSize), floored at 4MB and capped at what the YAML file parser itself permits, so documents that fit the budget chunk exactly as before - Replace isStructuredData + chunkJsonYaml with one chunkStructured entry point that parses once and returns null when the content is not structured, leaving chunker selection with the document processor
1 parent 178edba commit aba6fe2

4 files changed

Lines changed: 202 additions & 63 deletions

File tree

‎apps/sim/lib/chunkers/json-yaml-chunker.test.ts‎

Lines changed: 80 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
*/
44

55
import { describe, expect, it, vi } from 'vitest'
6-
import { JsonYamlChunker } from './json-yaml-chunker'
6+
import { JsonYamlChunker } from '@/lib/chunkers/json-yaml-chunker'
77

88
vi.mock('@/lib/tokenization', () => ({
99
getAccurateTokenCount: (text: string) => Math.ceil(text.length / 4),
@@ -37,30 +37,96 @@ describe('JsonYamlChunker', () => {
3737
expect(chunks.every((chunk) => chunk.tokenCount <= 100)).toBe(true)
3838
})
3939

40-
describe('isStructuredData', () => {
41-
it('should detect valid JSON', () => {
42-
expect(JsonYamlChunker.isStructuredData('{"key": "value"}')).toBe(true)
40+
describe('chunkStructured', () => {
41+
it('chunks valid JSON', async () => {
42+
await expect(JsonYamlChunker.chunkStructured('{"key": "value"}')).resolves.not.toBeNull()
4343
})
4444

45-
it('should detect valid JSON array', () => {
46-
expect(JsonYamlChunker.isStructuredData('[1, 2, 3]')).toBe(true)
45+
it('chunks a valid JSON array', async () => {
46+
await expect(JsonYamlChunker.chunkStructured('[1, 2, 3]')).resolves.not.toBeNull()
4747
})
4848

49-
it('should detect valid YAML', () => {
50-
expect(JsonYamlChunker.isStructuredData('key: value\nother: data')).toBe(true)
49+
it('chunks valid YAML', async () => {
50+
await expect(
51+
JsonYamlChunker.chunkStructured('key: value\nother: data')
52+
).resolves.not.toBeNull()
5153
})
5254

53-
it('should return false for plain text parsed as YAML scalar', () => {
54-
expect(JsonYamlChunker.isStructuredData('Hello, this is plain text.')).toBe(false)
55+
it('declines plain text that parses as a YAML scalar', async () => {
56+
await expect(
57+
JsonYamlChunker.chunkStructured('Hello, this is plain text.')
58+
).resolves.toBeNull()
5559
})
5660

57-
it('should return false for invalid JSON/YAML with unbalanced braces', () => {
58-
expect(JsonYamlChunker.isStructuredData('{invalid: json: content: {{')).toBe(false)
61+
it('declines invalid JSON/YAML with unbalanced braces', async () => {
62+
await expect(
63+
JsonYamlChunker.chunkStructured('{invalid: json: content: {{')
64+
).resolves.toBeNull()
5965
})
6066

61-
it('should detect nested JSON objects', () => {
67+
it('chunks nested JSON objects', async () => {
6268
const nested = JSON.stringify({ level1: { level2: { level3: 'value' } } })
63-
expect(JsonYamlChunker.isStructuredData(nested)).toBe(true)
69+
await expect(JsonYamlChunker.chunkStructured(nested)).resolves.not.toBeNull()
70+
})
71+
72+
it('declines an alias-expansion bomb instead of expanding it', async () => {
73+
const lines = ['a0: &a0 "lol"']
74+
for (let level = 1; level <= 7; level++) {
75+
lines.push(
76+
`a${level}: &a${level} [${Array(7)
77+
.fill(`*a${level - 1}`)
78+
.join(',')}]`
79+
)
80+
}
81+
lines.push('top: *a7')
82+
const bomb = lines.join('\n')
83+
84+
const chunks = await JsonYamlChunker.chunkStructured(bomb, {
85+
chunkSize: 1024,
86+
minCharactersPerChunk: 1,
87+
maxChunks: 5000,
88+
})
89+
90+
expect(chunks).toBeNull()
91+
})
92+
93+
it('never parses source larger than one output budget', async () => {
94+
const oversized = JSON.stringify({ value: 'x'.repeat(5 * 1024 * 1024) })
95+
const parse = vi.spyOn(JSON, 'parse')
96+
97+
try {
98+
await expect(
99+
JsonYamlChunker.chunkStructured(oversized, {
100+
chunkSize: 1024,
101+
minCharactersPerChunk: 1,
102+
maxChunks: 1024,
103+
})
104+
).resolves.toBeNull()
105+
expect(parse).not.toHaveBeenCalled()
106+
} finally {
107+
parse.mockRestore()
108+
}
109+
})
110+
111+
it('chunks with default options', async () => {
112+
const chunks = await JsonYamlChunker.chunkStructured(JSON.stringify({ test: 'value' }))
113+
114+
expect(chunks?.length).toBeGreaterThan(0)
115+
})
116+
117+
it('honors a custom chunk size', async () => {
118+
const largeObject: Record<string, string> = {}
119+
for (let i = 0; i < 50; i++) {
120+
largeObject[`key${i}`] = `value${i}`.repeat(20)
121+
}
122+
const json = JSON.stringify(largeObject)
123+
124+
const chunksSmall = await JsonYamlChunker.chunkStructured(json, { chunkSize: 50 })
125+
const chunksLarge = await JsonYamlChunker.chunkStructured(json, { chunkSize: 500 })
126+
127+
expect(chunksSmall).not.toBeNull()
128+
expect(chunksLarge).not.toBeNull()
129+
expect(chunksSmall?.length).toBeGreaterThan(chunksLarge?.length as number)
64130
})
65131
})
66132

@@ -368,28 +434,6 @@ server:
368434
})
369435
})
370436

371-
describe('static chunkJsonYaml method', () => {
372-
it.concurrent('should work with default options', async () => {
373-
const json = JSON.stringify({ test: 'value' })
374-
const chunks = await JsonYamlChunker.chunkJsonYaml(json)
375-
376-
expect(chunks.length).toBeGreaterThan(0)
377-
})
378-
379-
it.concurrent('should accept custom options', async () => {
380-
const largeObject: Record<string, string> = {}
381-
for (let i = 0; i < 50; i++) {
382-
largeObject[`key${i}`] = `value${i}`.repeat(20)
383-
}
384-
const json = JSON.stringify(largeObject)
385-
386-
const chunksSmall = await JsonYamlChunker.chunkJsonYaml(json, { chunkSize: 50 })
387-
const chunksLarge = await JsonYamlChunker.chunkJsonYaml(json, { chunkSize: 500 })
388-
389-
expect(chunksSmall.length).toBeGreaterThan(chunksLarge.length)
390-
})
391-
})
392-
393437
describe('chunk metadata', () => {
394438
it('preserves every source character and offset when bounding oversized chunks', async () => {
395439
const key = 'p'.repeat(80)

‎apps/sim/lib/chunkers/json-yaml-chunker.ts‎

Lines changed: 111 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@ import {
99
normalizeTokenChunkSize,
1010
tokensToChars,
1111
} from '@/lib/chunkers/utils'
12+
import { measureYamlExpansion, type YamlExpansionLimits } from '@/lib/file-parsers/yaml-limits'
13+
import { FILE_PARSER_YAML_LIMITS } from '@/lib/file-parsers/yaml-parser'
1214

1315
const logger = createLogger('JsonYamlChunker')
1416

@@ -20,40 +22,130 @@ type BoundedChunkMetadataMode = 'text-offsets' | 'preserve-range'
2022

2123
const MAX_DEPTH = 5
2224

25+
/**
26+
* Smallest expansion ceiling this chunker imposes, so a knowledge base
27+
* configured with tiny chunks keeps structural chunking on documents it indexes
28+
* perfectly well today.
29+
*/
30+
const MIN_EXPANSION_BYTES = 4 * 1024 * 1024
31+
32+
/**
33+
* How large an expanded document this chunker will materialize.
34+
*
35+
* Structural chunking re-serializes what it parsed, so its cost follows the
36+
* document's *expanded* size rather than its source size, and `yaml.load`
37+
* resolves aliases into shared references — a sub-kilobyte source can carry tens
38+
* of megabytes of expansion. `ChunkBudget` cannot bound that: it counts emitted
39+
* chunks, and every parse and serialization happens before the first is emitted.
40+
*
41+
* The ceiling is the most text this chunker could ever emit, one output budget's
42+
* worth, because a larger document cannot be indexed whole by any chunker and
43+
* paying to expand it buys nothing. Transient allocation therefore stays the
44+
* same order as the output, and every document that fits the budget is chunked
45+
* exactly as before.
46+
*/
47+
function resolveExpansionLimits(
48+
maxChunks: number | undefined,
49+
chunkSize: number
50+
): YamlExpansionLimits {
51+
const emittable =
52+
maxChunks === undefined
53+
? FILE_PARSER_YAML_LIMITS.maxSerializedBytes
54+
: maxChunks * tokensToChars(chunkSize)
55+
56+
return {
57+
/** Bytes bind here; every reached node charges some, so a self-referential anchor still terminates. */
58+
maxNodes: Number.MAX_SAFE_INTEGER,
59+
maxSerializedBytes: Math.min(
60+
FILE_PARSER_YAML_LIMITS.maxSerializedBytes,
61+
Math.max(MIN_EXPANSION_BYTES, emittable)
62+
),
63+
maxDepth: FILE_PARSER_YAML_LIMITS.maxDepth,
64+
}
65+
}
66+
2367
export class JsonYamlChunker {
2468
private chunkSize: number
2569
private minCharactersPerChunk: number
2670
private maxChunks?: number
71+
private readonly expansionLimits: YamlExpansionLimits
2772

2873
constructor(options: ChunkerOptions = {}) {
2974
this.chunkSize = normalizeTokenChunkSize(options.chunkSize ?? 1024, 'JSON/YAML chunk size')
3075
this.minCharactersPerChunk = options.minCharactersPerChunk ?? 100
3176
this.maxChunks = options.maxChunks
77+
this.expansionLimits = resolveExpansionLimits(this.maxChunks, this.chunkSize)
3278
}
3379

34-
static isStructuredData(content: string): boolean {
80+
/**
81+
* Read `content` as JSON, falling back to YAML, and measure what the parsed
82+
* value expands to before anything materializes it.
83+
*
84+
* The source-length check comes first so oversized content is never parsed at
85+
* all; the expansion measurement then catches what length alone cannot — alias
86+
* expansion, and the indentation a pretty-printed re-serialization adds.
87+
*/
88+
private parseWithinLimits(content: string): JsonValue | undefined {
89+
if (content.length > this.expansionLimits.maxSerializedBytes) {
90+
return this.reject(
91+
`source of ${content.length} characters exceeds the ${this.expansionLimits.maxSerializedBytes}-byte ceiling`
92+
)
93+
}
94+
95+
let parsed: unknown
3596
try {
36-
const parsed = JSON.parse(content)
37-
return typeof parsed === 'object' && parsed !== null
97+
parsed = JSON.parse(content)
3898
} catch {
3999
try {
40-
const parsed = yaml.load(content)
41-
return typeof parsed === 'object' && parsed !== null
100+
parsed = yaml.load(content)
42101
} catch {
43-
return false
102+
return undefined
44103
}
45104
}
105+
106+
if (parsed === undefined) return undefined
107+
108+
const measured = measureYamlExpansion(parsed, this.expansionLimits)
109+
if (!measured.within) return this.reject(measured.reason)
110+
111+
return parsed as JsonValue
46112
}
47113

48-
async chunk(content: string): Promise<Chunk[]> {
49-
try {
50-
let data: JsonValue
51-
try {
52-
data = JSON.parse(content) as JsonValue
53-
} catch {
54-
data = yaml.load(content) as JsonValue
114+
private reject(reason: string): undefined {
115+
logger.warn(
116+
'Structured content exceeds the chunking expansion limits, declining to expand it',
117+
{
118+
reason,
55119
}
120+
)
121+
return undefined
122+
}
123+
124+
/**
125+
* Chunk `content` as a structured object or array, or return `null` when it is
126+
* neither — including when its expanded form outgrows the ceiling above. The
127+
* caller then chooses another chunker for it.
128+
*/
129+
static async chunkStructured(
130+
content: string,
131+
options: ChunkerOptions = {}
132+
): Promise<Chunk[] | null> {
133+
const chunker = new JsonYamlChunker(options)
134+
const data = chunker.parseWithinLimits(content)
135+
if (data === null || typeof data !== 'object') return null
136+
137+
return chunker.chunkParsed(data, content)
138+
}
139+
140+
async chunk(content: string): Promise<Chunk[]> {
141+
const data = this.parseWithinLimits(content)
142+
if (data === undefined) return this.chunkAsText(content)
56143

144+
return this.chunkParsed(data, content)
145+
}
146+
147+
private chunkParsed(data: JsonValue, content: string): Chunk[] {
148+
try {
57149
const chunks: Chunk[] = []
58150
this.chunkStructuredData(data, [], 0, chunks, new ChunkBudget(this.maxChunks))
59151

@@ -64,7 +156,7 @@ export class JsonYamlChunker {
64156
} catch (error) {
65157
if (error instanceof ChunkLimitExceededError) throw error
66158
logger.info('Structured data chunking failed, falling back to text chunking')
67-
return this.chunkAsText(content, new ChunkBudget(this.maxChunks))
159+
return this.chunkAsText(content)
68160
}
69161
}
70162

@@ -299,7 +391,11 @@ export class JsonYamlChunker {
299391
}
300392
}
301393

302-
private chunkAsText(content: string, budget: ChunkBudget, chunks: Chunk[] = []): Chunk[] {
394+
private chunkAsText(
395+
content: string,
396+
budget: ChunkBudget = new ChunkBudget(this.maxChunks),
397+
chunks: Chunk[] = []
398+
): Chunk[] {
303399
let currentChunk = ''
304400
let currentTokens = 0
305401
let startIndex = 0
@@ -362,9 +458,4 @@ export class JsonYamlChunker {
362458

363459
return chunks
364460
}
365-
366-
static async chunkJsonYaml(content: string, options: ChunkerOptions = {}): Promise<Chunk[]> {
367-
const chunker = new JsonYamlChunker(options)
368-
return chunker.chunk(content)
369-
}
370461
}

‎apps/sim/lib/file-parsers/yaml-parser.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import { measureYamlExpansion, type YamlExpansionLimits } from '@/lib/file-parse
1010
* the byte cap bounds output a sub-1 KB input can inflate to hundreds of MB;
1111
* the depth cap bounds the traversal's own working set.
1212
*/
13-
const FILE_PARSER_YAML_LIMITS: YamlExpansionLimits = {
13+
export const FILE_PARSER_YAML_LIMITS: YamlExpansionLimits = {
1414
maxNodes: 5_000_000,
1515
maxSerializedBytes: 64 * 1024 * 1024,
1616
maxDepth: 500,

‎apps/sim/lib/knowledge/documents/document-processor.ts‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -268,13 +268,17 @@ export async function processDocument(
268268
mimeType.includes('json') ||
269269
mimeType.includes('yaml')
270270

271-
if (isJsonYaml && JsonYamlChunker.isStructuredData(content)) {
271+
const jsonYamlChunks = isJsonYaml
272+
? await JsonYamlChunker.chunkStructured(content, {
273+
chunkSize,
274+
minCharactersPerChunk,
275+
maxChunks: MAX_DOCUMENT_CHUNKS,
276+
})
277+
: null
278+
279+
if (jsonYamlChunks !== null) {
272280
logger.info('Using JSON/YAML chunker for structured data')
273-
chunks = await JsonYamlChunker.chunkJsonYaml(content, {
274-
chunkSize,
275-
minCharactersPerChunk,
276-
maxChunks: MAX_DOCUMENT_CHUNKS,
277-
})
281+
chunks = jsonYamlChunks
278282
} else if (StructuredDataChunker.isStructuredData(content, mimeType)) {
279283
logger.info('Using structured data chunker for spreadsheet/CSV content')
280284
const rowCount = metadata.totalRows ?? metadata.rowCount

0 commit comments

Comments
 (0)