diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index 060b647693..f6b1c83db3 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -334,7 +334,7 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string // Split-then-merge: split on delimiters, then greedily merge to // chunk_token_size with optional overlap. perItem := [][]schema.ChunkDoc{docs} - merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent, true, !c.param.UnderCap) + merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent, true, c.param.MergeStrategy()) return chunkOutputs(flatten(merged)) } @@ -496,18 +496,21 @@ const ( // is the separator used to project the joined text ("" for the text path, "\n" // for the JSON path). target is the token cap. // -// allowBoundaryOverflow=true selects OVER_CAP (Python's default, canonical): -// when the joined text exceeds target but the incoming unit still fits target, -// it is merged into the previous chunk and that chunk is then closed -// (mergeThenClose), forcing the next unit to start a new chunk. An incoming -// unit that already exceeds target is never merged — it stands alone as its own -// chunk (Python OVER_CAP: an oversized paragraph is never combined with the -// previous chunk). allowBoundaryOverflow=false selects UNDER_CAP (strict -// no-overflow): an overflowing joined text starts a new chunk instead. +// strategy selects the merge strategy (schema.MergeStrategy), mirroring +// Python's MergeStrategy: +// - MergeOverCap (default, Python OVER_CAP): when the joined text exceeds +// target but the incoming unit still fits target, it is merged into the +// previous chunk and that chunk is then closed (mergeThenClose), forcing +// the next unit to start a new chunk. An incoming unit that already exceeds +// target is never merged — it stands alone as its own chunk (Python +// OVER_CAP: an oversized paragraph is never combined with the previous +// chunk). +// - MergeUnderCap (Python UNDER_CAP, strict no-overflow): an overflowing +// joined text starts a new chunk instead. // // JSON-only metadata (PDFPositions/Positions/TKNums) is the caller's // responsibility; this helper only returns the merged/new text and the action. -func mergeDecision(prevText, incoming, joinSep string, target int, overlapPct float64, allowBoundaryOverflow bool) (string, mergeAction) { +func mergeDecision(prevText, incoming, joinSep string, target int, overlapPct float64, strategy schema.MergeStrategy) (string, mergeAction) { incomingTokens := tokenizeStr(incoming) // An incoming unit that already exceeds target can never be merged; it // stands alone as its own chunk. @@ -518,7 +521,7 @@ func mergeDecision(prevText, incoming, joinSep string, target int, overlapPct fl if tokenizeStr(joined) <= target { return joined, mergeIntoPrev } - if allowBoundaryOverflow { + if strategy == schema.MergeOverCap { // OVER_CAP: merge the overflowing unit but close the chunk so the // next unit starts fresh. return joined, mergeThenClose @@ -593,7 +596,7 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r tkns = append(tkns, tokenizeStr(out)) return } - out, act := mergeDecision(cks[len(cks)-1], segment, "", target, overlapPct, !c.param.UnderCap) + out, act := mergeDecision(cks[len(cks)-1], segment, "", target, overlapPct, c.param.MergeStrategy()) switch act { case mergeIntoPrev, mergeThenClose: cks[len(cks)-1] = out @@ -722,7 +725,7 @@ func (c *TokenChunkerComponent) invokeJSONPayload(ctx context.Context, items []s // chunks across JSON items into one global token budget. Flatten the // per-item structure into a single sequence first so the merge is // global; non-text chunks still break the merge via their CKType. - attached = mergeByTokenSizeFromJSON([][]schema.ChunkDoc{flatten(attached)}, c.param.ChunkTokenSize, c.param.OverlappedPercent, false, !c.param.UnderCap) + attached = mergeByTokenSizeFromJSON([][]schema.ChunkDoc{flatten(attached)}, c.param.ChunkTokenSize, c.param.OverlappedPercent, false, c.param.MergeStrategy()) } flat := flatten(attached) @@ -962,11 +965,12 @@ func takeFromStart(text string, tokens int) string { // Oversized text units are sub-split via splitOversizedUnit before merge; // overlap is applied only when overlap+segment still fits the budget. // -// allowBoundaryOverflow selects the merge strategy: true = OVER_CAP (Python's -// canonical default, a chunk may exceed the target by at most one incoming -// unit), false = UNDER_CAP (never exceed the target; a projected overflow -// starts a fresh chunk). The TokenChunker threads its UnderCap param here. -func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64, subSplitOversize bool, allowBoundaryOverflow bool) [][]schema.ChunkDoc { +// strategy selects the merge strategy (schema.MergeStrategy): MergeOverCap = +// OVER_CAP (Python's canonical default, a chunk may exceed the target by at +// most one incoming unit), MergeUnderCap = UNDER_CAP (never exceed the target; +// a projected overflow starts a fresh chunk). The TokenChunker threads its +// MergeStrategy() here. +func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64, subSplitOversize bool, strategy schema.MergeStrategy) [][]schema.ChunkDoc { // overlappedPct is a [0,100] percentage. Clamp defensively because this // helper is also exercised directly by tests. if overlappedPct < 0 { @@ -1032,7 +1036,7 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over return } // Proactive projected-total merge (joined with "\n"). - out, act := mergeDecision(prev.Text, ck.Text, "\n", chunkTokens, overlappedPct, allowBoundaryOverflow) + out, act := mergeDecision(prev.Text, ck.Text, "\n", chunkTokens, overlappedPct, strategy) switch act { case mergeIntoPrev, mergeThenClose: prev.Text = out diff --git a/internal/ingestion/component/chunker/token_batch1_test.go b/internal/ingestion/component/chunker/token_batch1_test.go index 3312b9acf9..8e5b506065 100644 --- a/internal/ingestion/component/chunker/token_batch1_test.go +++ b/internal/ingestion/component/chunker/token_batch1_test.go @@ -91,7 +91,7 @@ func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) { {Text: cText, DocType: "text", CKType: "text", TKNums: intPtr(cN)}, }, } - got := mergeByTokenSizeFromJSON(items, budget, 30.0, true, true) + got := mergeByTokenSizeFromJSON(items, budget, 30.0, true, schema.MergeOverCap) merged := got[0] if len(merged) != 2 { t.Fatalf("want 2 chunks (overflow-closed + overlap chunk), got %d (a=%d b=%d c=%d budget=%d)", len(merged), aN, bN, cN, budget) @@ -157,7 +157,7 @@ func TestMergeByTokenSizeFromJSON_NonTextBoundaryResetsPrevClosed(t *testing.T) {Text: t4, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(t4))}, }, } - got := mergeByTokenSizeFromJSON(items, budget, 0.0, true, true) + got := mergeByTokenSizeFromJSON(items, budget, 0.0, true, schema.MergeOverCap) merged := got[0] // Expect: chunk0 (T1+T2, overflow-closed), N (non-text), chunk1 (T3+T4 merged). if len(merged) != 3 { @@ -183,7 +183,7 @@ func TestMergeByTokenSizeFromJSON_NonTextBoundaryResetsPrevClosed(t *testing.T) } // TestMergeByTokenSizeFromJSON_UnderCapNoOverflow exercises the UNDER_CAP -// strategy (allowBoundaryOverflow=false): a projected join that would exceed +// strategy (schema.MergeUnderCap): a projected join that would exceed // the target must start a fresh chunk instead of merging-then-closing. This is // the seam that lets Go follow Python's no-overflow (UNDER_CAP) strategy. Under // OVER_CAP the same input merges a+b and overflows chunk0; here a, b, c must @@ -213,7 +213,7 @@ func TestMergeByTokenSizeFromJSON_UnderCapNoOverflow(t *testing.T) { {Text: cText, DocType: "text", CKType: "text", TKNums: intPtr(cN)}, }, } - got := mergeByTokenSizeFromJSON(items, budget, 0.0, true, false) + got := mergeByTokenSizeFromJSON(items, budget, 0.0, true, schema.MergeUnderCap) merged := got[0] if len(merged) != 3 { t.Fatalf("UNDER_CAP want 3 chunks (a, b, c separate), got %d", len(merged)) @@ -255,12 +255,12 @@ func clampOverlapFixture() [][]schema.ChunkDoc { } func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) { - at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100, true, true) + at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100, true, schema.MergeOverCap) if at100 == nil || len(at100) == 0 { t.Fatalf("overlappedPct=100: nil/empty result") } - at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, true, true) - atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, true, true) + at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, true, schema.MergeOverCap) + atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, true, schema.MergeOverCap) if !reflect.DeepEqual(at100, at150) { t.Errorf("overlappedPct=150 should clamp to 100; output differs from 100") } @@ -268,12 +268,12 @@ func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) { t.Errorf("overlappedPct=1e300 should clamp to 100; output differs from 100") } - at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0, true, true) + at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0, true, schema.MergeOverCap) if at0 == nil || len(at0) == 0 { t.Fatalf("overlappedPct=0: nil/empty result") } - atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, true, true) - atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, true, true) + atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, true, schema.MergeOverCap) + atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, true, schema.MergeOverCap) if !reflect.DeepEqual(at0, atNeg) { t.Errorf("overlappedPct=-5 should clamp to 0; output differs from 0") } @@ -294,7 +294,7 @@ func TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk(t *testing.T) { {Text: "keepme", DocType: "text", CKType: "text", TKNums: intPtr(5)}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0, true, true) + got := mergeByTokenSizeFromJSON(items, 128, 0, true, schema.MergeOverCap) merged := got[0] if len(merged) != 1 { t.Fatalf("want 1 merged chunk, got %d", len(merged)) diff --git a/internal/ingestion/component/chunker/token_pdfpos_test.go b/internal/ingestion/component/chunker/token_pdfpos_test.go index 0f6b501a46..71945bec53 100644 --- a/internal/ingestion/component/chunker/token_pdfpos_test.go +++ b/internal/ingestion/component/chunker/token_pdfpos_test.go @@ -38,7 +38,7 @@ func TestMergeByTokenSizeFromJSON_ExtendsPDFPositions(t *testing.T) { {Text: "beta", DocType: "text", CKType: "text", TKNums: intPtr(5), PDFPositions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0, true, true) + got := mergeByTokenSizeFromJSON(items, 128, 0, true, schema.MergeOverCap) merged := got[0] if len(merged) != 1 { t.Fatalf("want 1 merged chunk, got %d", len(merged)) @@ -63,7 +63,7 @@ func TestMergeByTokenSizeFromJSON_ExtendsPositions(t *testing.T) { {Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0, true, true) + got := mergeByTokenSizeFromJSON(items, 128, 0, true, schema.MergeOverCap) combined := string(got[0][0].Positions) if !strings.Contains(combined, "1,2,3") || !strings.Contains(combined, "4,5,6") { t.Errorf("merged chunk dropped/omitted `positions`: %s", combined) @@ -103,7 +103,7 @@ func TestMergeByTokenSizeFromJSON_PositionsDecodeToMatrix(t *testing.T) { {Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0, true, true) + got := mergeByTokenSizeFromJSON(items, 128, 0, true, schema.MergeOverCap) m := got[0][0].ToMap() raw, ok := m["positions"] if !ok { diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index 75eff6a4b5..57e9800aa6 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -108,7 +108,7 @@ func TestMergeByTokenSizeFromJSON_StrictCapNoOvershoot(t *testing.T) { Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)), }) } - got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0, true, true) + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0, true, schema.MergeOverCap) merged := got[0] if len(merged) < 3 { t.Fatalf("want >=3 chunks, got %d", len(merged)) @@ -135,7 +135,7 @@ func TestMergeByTokenSizeFromJSON_OverlapDroppedAtOverflow(t *testing.T) { Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)), }) } - got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20, true, true) + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20, true, schema.MergeOverCap) unit := tokenizeStr(sections[0].Text) for i, ck := range got[0] { // OVER_CAP allows one boundary overflow (prev + one unit). The JSON @@ -157,7 +157,7 @@ func TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit(t *testing.T) { items := [][]schema.ChunkDoc{{ {Text: long, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(long))}, }} - got := mergeByTokenSizeFromJSON(items, budget, 0, true, true) + got := mergeByTokenSizeFromJSON(items, budget, 0, true, schema.MergeOverCap) if len(got[0]) < 2 { t.Fatalf("oversized unit must yield multiple chunks, got %d", len(got[0])) } @@ -350,3 +350,76 @@ func TestInvokeTextPayload_StrictCapEndToEnd(t *testing.T) { } } } + +// TestInvokeJSONPayload_UnderCapEndToEnd exercises the UNDER_CAP strategy on +// the JSON path end-to-end (invokeJSONPayload). Many single-line JSON items +// are flattened into one global merge sequence, and under_cap=true must keep +// every chunk within chunk_token_size. The OVER_CAP control is asserted to +// pack fewer chunks than UNDER_CAP, proving the toggle is live on the JSON +// path — not just the text path covered by TestInvokeTextPayload_StrictCapEndToEnd. +func TestInvokeJSONPayload_UnderCapEndToEnd(t *testing.T) { + const budget = 32 + unit := tokenizeStr(strings.TrimSpace(strings.Repeat("alpha ", 12))) + + // Each item is one ~unit-sized token block; 24 items give the global + // merge plenty to accumulate. + var items []map[string]any + for i := 0; i < 24; i++ { + items = append(items, map[string]any{ + "text": strings.TrimSpace(strings.Repeat("alpha ", 12)), + "doc_type_kwd": "text", + }) + } + + run := func(underCap bool) []map[string]any { + comp, err := NewTokenChunker(map[string]any{ + "delimiter_mode": "delimiter", + "delimiters": []string{"\n"}, + "chunk_token_size": budget, + "under_cap": underCap, + }) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out, err := comp.Invoke(context.Background(), nil, map[string]any{ + "name": "doc.json", + "output_format": "json", + "json": items, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if errMsg, _ := out["_ERROR"].(string); errMsg != "" { + t.Fatalf("Invoke error payload: %s", errMsg) + } + chunks, _ := out["chunks"].([]map[string]any) + return chunks + } + + // UNDER_CAP: no chunk may exceed the token target. + under := run(true) + if len(under) < 2 { + t.Fatalf("UNDER_CAP: want multiple chunks, got %d", len(under)) + } + for i, ck := range under { + text, _ := ck["text"].(string) + if n := tokenizeStr(text); n > budget { + t.Errorf("UNDER_CAP chunk %d exceeds target: tokens=%d (cap=%d)", i, n, budget) + } + } + + // OVER_CAP control: the same input packs fewer chunks (it allows one + // boundary overflow per chunk), proving the toggle is live on the JSON + // path. Equal lengths would mean under_cap is a no-op here. + over := run(false) + if len(over) >= len(under) { + t.Errorf("OVER_CAP should pack fewer chunks than UNDER_CAP (over=%d under=%d); toggle may be a no-op on the JSON path", len(over), len(under)) + } + for i, ck := range over { + text, _ := ck["text"].(string) + // OVER_CAP may overflow by at most one unit. + if n := tokenizeStr(text); n > budget+unit { + t.Errorf("OVER_CAP chunk %d overflows by more than one unit: tokens=%d (cap=%d unit=%d)", i, n, budget, unit) + } + } +} diff --git a/internal/ingestion/component/schema/chunker.go b/internal/ingestion/component/schema/chunker.go index 541a8c80c7..ad9b9a421c 100644 --- a/internal/ingestion/component/schema/chunker.go +++ b/internal/ingestion/component/schema/chunker.go @@ -211,11 +211,38 @@ type TokenChunkerParam struct { // (a boundary overflow then closes the chunk). // - true: UNDER_CAP — strictly never exceed the target; when the // projected join would overflow, start a fresh chunk instead. - // This is the seam that lets Go follow Python's no-overflow strategy - // without changing the default behavior. + // This is the wire-facing bool; the active strategy is exposed as + // MergeStrategy so callers never have to invert it. UnderCap bool `json:"under_cap"` } +// MergeStrategy selects how the TokenChunker greedily accumulates adjacent +// units into a chunk. It mirrors Python's rag/nlp/__init__.py MergeStrategy so +// the Go and Python implementations stay on the same vocabulary (OVER_CAP / +// UNDER_CAP) instead of a bare inlined bool. +type MergeStrategy int + +const ( + // MergeOverCap is the canonical default (Python OVER_CAP): a chunk may + // exceed the token target by at most one incoming unit (a boundary + // overflow then closes the chunk). + MergeOverCap MergeStrategy = iota + // MergeUnderCap never exceeds the target; when the projected join would + // overflow, a fresh chunk is started instead (Python UNDER_CAP). + MergeUnderCap +) + +// MergeStrategy reports the active merge strategy for this param. It is derived +// from UnderCap so existing configs that set "under_cap" keep working without a +// schema break, and callers read the strategy directly instead of inverting a +// bool at every call site. +func (p TokenChunkerParam) MergeStrategy() MergeStrategy { + if p.UnderCap { + return MergeUnderCap + } + return MergeOverCap +} + // Defaults returns the Python default TokenChunkerParam. func (TokenChunkerParam) Defaults() TokenChunkerParam { return TokenChunkerParam{ diff --git a/internal/ingestion/component/schema/schema_test.go b/internal/ingestion/component/schema/schema_test.go index cdb3dc4474..517513f951 100644 --- a/internal/ingestion/component/schema/schema_test.go +++ b/internal/ingestion/component/schema/schema_test.go @@ -567,3 +567,15 @@ func TestExtractorOutputsJSONRoundTrip(t *testing.T) { // --------------------------------------------------------------------------- func ptrString(s string) *string { return &s } + +// TestTokenChunkerParamMergeStrategy locks the UnderCap bool -> MergeStrategy +// mapping so a future refactor cannot silently flip OVER_CAP/UNDER_CAP without +// a failing test. +func TestTokenChunkerParamMergeStrategy(t *testing.T) { + if got := (TokenChunkerParam{}).MergeStrategy(); got != MergeOverCap { + t.Errorf("default (UnderCap=false) want MergeOverCap, got %v", got) + } + if got := (TokenChunkerParam{UnderCap: true}).MergeStrategy(); got != MergeUnderCap { + t.Errorf("UnderCap=true want MergeUnderCap, got %v", got) + } +}