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
1315const logger = createLogger ( 'JsonYamlChunker' )
1416
@@ -20,40 +22,148 @@ type BoundedChunkMetadataMode = 'text-offsets' | 'preserve-range'
2022
2123const MAX_DEPTH = 5
2224
25+ /**
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.
29+ */
30+ const MIN_SOURCE_BYTES = 4 * 1024 * 1024
31+
32+ /**
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.
59+ *
60+ * Structural chunking re-serializes what it parsed, so its cost follows the
61+ * document's *expanded* size rather than its source size, and `yaml.load`
62+ * resolves aliases into shared references — a sub-kilobyte source can carry tens
63+ * of megabytes of expansion. `ChunkBudget` cannot bound that: it counts emitted
64+ * chunks, and every parse and serialization happens before the first is emitted.
65+ *
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.
71+ */
72+ function resolveExpansionLimits ( sourceBytes : number , maxSourceBytes : number ) : YamlExpansionLimits {
73+ return {
74+ /** Bytes bind here; every reached node charges some, so a self-referential anchor still terminates. */
75+ maxNodes : Number . MAX_SAFE_INTEGER ,
76+ maxSerializedBytes : Math . min (
77+ FILE_PARSER_YAML_LIMITS . maxSerializedBytes ,
78+ Math . max ( maxSourceBytes , sourceBytes * MAX_EXPANSION_RATIO )
79+ ) ,
80+ maxDepth : FILE_PARSER_YAML_LIMITS . maxDepth ,
81+ }
82+ }
83+
2384export class JsonYamlChunker {
2485 private chunkSize : number
2586 private minCharactersPerChunk : number
2687 private maxChunks ?: number
88+ private readonly maxSourceBytes : number
2789
2890 constructor ( options : ChunkerOptions = { } ) {
2991 this . chunkSize = normalizeTokenChunkSize ( options . chunkSize ?? 1024 , 'JSON/YAML chunk size' )
3092 this . minCharactersPerChunk = options . minCharactersPerChunk ?? 100
3193 this . maxChunks = options . maxChunks
94+ this . maxSourceBytes = resolveMaxSourceBytes ( this . maxChunks , this . chunkSize )
3295 }
3396
34- static isStructuredData ( content : string ) : boolean {
97+ /**
98+ * Read `content` as JSON, falling back to YAML, and measure what the parsed
99+ * value expands to before anything materializes it.
100+ *
101+ * The source-length check comes first so oversized content is never parsed at
102+ * all; the expansion measurement then catches what length alone cannot — alias
103+ * expansion, and the indentation a pretty-printed re-serialization adds.
104+ */
105+ private parseWithinLimits ( content : string ) : JsonValue | undefined {
106+ if ( content . length > this . maxSourceBytes ) {
107+ return this . reject (
108+ `source of ${ content . length } ${ this . maxSourceBytes }
109+ )
110+ }
111+
112+ let parsed : unknown
35113 try {
36- const parsed = JSON . parse ( content )
37- return typeof parsed === 'object' && parsed !== null
114+ parsed = JSON . parse ( content )
38115 } catch {
39116 try {
40- const parsed = yaml . load ( content )
41- return typeof parsed === 'object' && parsed !== null
117+ parsed = yaml . load ( content )
42118 } catch {
43- return false
119+ return undefined
44120 }
45121 }
122+
123+ if ( parsed === undefined ) return undefined
124+
125+ const limits = resolveExpansionLimits ( content . length , this . maxSourceBytes )
126+ const measured = measureYamlExpansion ( parsed , limits )
127+ if ( ! measured . within ) return this . reject ( measured . reason )
128+
129+ return parsed as JsonValue
46130 }
47131
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
132+ private reject ( reason : string ) : undefined {
133+ logger . warn (
134+ 'Structured content exceeds the chunking expansion limits, declining to expand it' ,
135+ {
136+ reason,
55137 }
138+ )
139+ return undefined
140+ }
141+
142+ /**
143+ * Chunk `content` as a structured object or array, or return `null` when it is
144+ * neither — including when its expanded form outgrows the ceiling above. The
145+ * caller then chooses another chunker for it.
146+ */
147+ static async chunkStructured (
148+ content : string ,
149+ options : ChunkerOptions = { }
150+ ) : Promise < Chunk [ ] | null > {
151+ const chunker = new JsonYamlChunker ( options )
152+ const data = chunker . parseWithinLimits ( content )
153+ if ( data === null || typeof data !== 'object' ) return null
154+
155+ return chunker . chunkParsed ( data , content )
156+ }
157+
158+ async chunk ( content : string ) : Promise < Chunk [ ] > {
159+ const data = this . parseWithinLimits ( content )
160+ if ( data === undefined ) return this . chunkAsText ( content )
161+
162+ return this . chunkParsed ( data , content )
163+ }
56164
165+ private chunkParsed ( data : JsonValue , content : string ) : Chunk [ ] {
166+ try {
57167 const chunks : Chunk [ ] = [ ]
58168 this . chunkStructuredData ( data , [ ] , 0 , chunks , new ChunkBudget ( this . maxChunks ) )
59169
@@ -64,7 +174,7 @@ export class JsonYamlChunker {
64174 } catch ( error ) {
65175 if ( error instanceof ChunkLimitExceededError ) throw error
66176 logger . info ( 'Structured data chunking failed, falling back to text chunking' )
67- return this . chunkAsText ( content , new ChunkBudget ( this . maxChunks ) )
177+ return this . chunkAsText ( content )
68178 }
69179 }
70180
@@ -299,7 +409,11 @@ export class JsonYamlChunker {
299409 }
300410 }
301411
302- private chunkAsText ( content : string , budget : ChunkBudget , chunks : Chunk [ ] = [ ] ) : Chunk [ ] {
412+ private chunkAsText (
413+ content : string ,
414+ budget : ChunkBudget = new ChunkBudget ( this . maxChunks ) ,
415+ chunks : Chunk [ ] = [ ]
416+ ) : Chunk [ ] {
303417 let currentChunk = ''
304418 let currentTokens = 0
305419 let startIndex = 0
@@ -362,9 +476,4 @@ export class JsonYamlChunker {
362476
363477 return chunks
364478 }
365-
366- static async chunkJsonYaml ( content : string , options : ChunkerOptions = { } ) : Promise < Chunk [ ] > {
367- const chunker = new JsonYamlChunker ( options )
368- return chunker . chunk ( content )
369- }
370479}
0 commit comments