Skip to content

Commit 1e3e3eb

Browse files
committed
fix(knowledge): allow proportionate expansion in the chunker guard
Comparing the expansion estimate against the output budget mixed two units: measureYamlExpansion charges a flat per-node allowance, so a document of many small values is charged several times its pretty-printed size and was rejected even though its chunks fit the budget — a flat array of a million booleans is charged ~22MB against ~8MB of real output. Allow the larger of two admissible expansions: one that fits the output budget, and one proportionate to the source. Alias expansion overshoots its source by orders of magnitude, so it is still rejected, while an ordinary large document is no longer charged for the estimator's conservatism. Neither allowance ever exceeds the file parser's own cap.
1 parent aba6fe2 commit 1e3e3eb

2 files changed

Lines changed: 57 additions & 25 deletions

File tree

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,20 @@ describe('JsonYamlChunker', () => {
9090
expect(chunks).toBeNull()
9191
})
9292

93+
it('keeps structure for a many-small-node document that fits the budget', async () => {
94+
const flags = JSON.stringify(Array.from({ length: 250_000 }, (_, i) => i % 2 === 0))
95+
96+
const chunks = await JsonYamlChunker.chunkStructured(flags, {
97+
chunkSize: 1024,
98+
minCharactersPerChunk: 1,
99+
maxChunks: 1000,
100+
})
101+
102+
expect(chunks).not.toBeNull()
103+
expect(chunks?.length).toBeGreaterThan(1)
104+
expect(chunks?.[0].text).toContain('true')
105+
})
106+
93107
it('never parses source larger than one output budget', async () => {
94108
const oversized = JSON.stringify({ value: 'x'.repeat(5 * 1024 * 1024) })
95109
const parse = vi.spyOn(JSON, 'parse')

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

Lines changed: 43 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -23,42 +23,59 @@ type BoundedChunkMetadataMode = 'text-offsets' | 'preserve-range'
2323
const MAX_DEPTH = 5
2424

2525
/**
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.
26+
* Smallest source ceiling this chunker imposes, so a knowledge base configured
27+
* with tiny chunks keeps structural chunking on documents it indexes perfectly
28+
* well today.
2929
*/
30-
const MIN_EXPANSION_BYTES = 4 * 1024 * 1024
30+
const MIN_SOURCE_BYTES = 4 * 1024 * 1024
3131

3232
/**
33-
* How large an expanded document this chunker will materialize.
33+
* How far a document may legitimately expand past its own source.
34+
*
35+
* `measureYamlExpansion` charges a flat per-node allowance, so a compact source
36+
* of small values is charged well above its own length — `[1,1,1]` costs about
37+
* 22 estimated bytes per element against two in source. An order of magnitude of
38+
* headroom therefore covers ordinary document shape, while alias expansion
39+
* overshoots it by several orders.
40+
*/
41+
const MAX_EXPANSION_RATIO = 16
42+
43+
/**
44+
* Longest source this chunker will parse: the most text it could ever emit, one
45+
* output budget's worth. A larger document cannot be indexed whole by any
46+
* chunker — `ChunkBudget` stops it either way — so parsing it buys nothing.
47+
*/
48+
function resolveMaxSourceBytes(maxChunks: number | undefined, chunkSize: number): number {
49+
if (maxChunks === undefined) return FILE_PARSER_YAML_LIMITS.maxSerializedBytes
50+
51+
return Math.min(
52+
FILE_PARSER_YAML_LIMITS.maxSerializedBytes,
53+
Math.max(MIN_SOURCE_BYTES, maxChunks * tokensToChars(chunkSize))
54+
)
55+
}
56+
57+
/**
58+
* What the document is allowed to expand to once it is walked as a tree.
3459
*
3560
* Structural chunking re-serializes what it parsed, so its cost follows the
3661
* document's *expanded* size rather than its source size, and `yaml.load`
3762
* resolves aliases into shared references — a sub-kilobyte source can carry tens
3863
* of megabytes of expansion. `ChunkBudget` cannot bound that: it counts emitted
3964
* chunks, and every parse and serialization happens before the first is emitted.
4065
*
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.
66+
* Two expansions are admissible: one that stays within the output budget, and
67+
* one that stays proportionate to the source. Taking the larger of the two keeps
68+
* transient allocation tied to work the chunker would have done anyway, without
69+
* charging an ordinary large document for the estimator's per-node conservatism.
70+
* Neither is ever allowed past what the file parser itself would hand over.
4671
*/
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-
72+
function resolveExpansionLimits(sourceBytes: number, maxSourceBytes: number): YamlExpansionLimits {
5673
return {
5774
/** Bytes bind here; every reached node charges some, so a self-referential anchor still terminates. */
5875
maxNodes: Number.MAX_SAFE_INTEGER,
5976
maxSerializedBytes: Math.min(
6077
FILE_PARSER_YAML_LIMITS.maxSerializedBytes,
61-
Math.max(MIN_EXPANSION_BYTES, emittable)
78+
Math.max(maxSourceBytes, sourceBytes * MAX_EXPANSION_RATIO)
6279
),
6380
maxDepth: FILE_PARSER_YAML_LIMITS.maxDepth,
6481
}
@@ -68,13 +85,13 @@ export class JsonYamlChunker {
6885
private chunkSize: number
6986
private minCharactersPerChunk: number
7087
private maxChunks?: number
71-
private readonly expansionLimits: YamlExpansionLimits
88+
private readonly maxSourceBytes: number
7289

7390
constructor(options: ChunkerOptions = {}) {
7491
this.chunkSize = normalizeTokenChunkSize(options.chunkSize ?? 1024, 'JSON/YAML chunk size')
7592
this.minCharactersPerChunk = options.minCharactersPerChunk ?? 100
7693
this.maxChunks = options.maxChunks
77-
this.expansionLimits = resolveExpansionLimits(this.maxChunks, this.chunkSize)
94+
this.maxSourceBytes = resolveMaxSourceBytes(this.maxChunks, this.chunkSize)
7895
}
7996

8097
/**
@@ -86,9 +103,9 @@ export class JsonYamlChunker {
86103
* expansion, and the indentation a pretty-printed re-serialization adds.
87104
*/
88105
private parseWithinLimits(content: string): JsonValue | undefined {
89-
if (content.length > this.expansionLimits.maxSerializedBytes) {
106+
if (content.length > this.maxSourceBytes) {
90107
return this.reject(
91-
`source of ${content.length} characters exceeds the ${this.expansionLimits.maxSerializedBytes}-byte ceiling`
108+
`source of ${content.length} characters exceeds the ${this.maxSourceBytes}-byte ceiling`
92109
)
93110
}
94111

@@ -105,7 +122,8 @@ export class JsonYamlChunker {
105122

106123
if (parsed === undefined) return undefined
107124

108-
const measured = measureYamlExpansion(parsed, this.expansionLimits)
125+
const limits = resolveExpansionLimits(content.length, this.maxSourceBytes)
126+
const measured = measureYamlExpansion(parsed, limits)
109127
if (!measured.within) return this.reject(measured.reason)
110128

111129
return parsed as JsonValue

0 commit comments

Comments
 (0)