From 2403988b6d812577a6be3406f877238e76649e4e Mon Sep 17 00:00:00 2001 From: bigfish-emily Date: Wed, 5 Aug 2026 19:58:51 +0800 Subject: [PATCH] fix(chunker): drop dead atom-split path in mergeByTokenSizeFromJSON (#17873) --- internal/ingestion/component/chunker/token.go | 39 +++++-------------- .../component/chunker/token_batch1_test.go | 20 +++++----- .../component/chunker/token_pdfpos_test.go | 6 +-- .../chunker/token_strict_cap_test.go | 31 +++++---------- 4 files changed, 33 insertions(+), 63 deletions(-) diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index b089cf233f..b74dd8de6a 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.MergeStrategy()) + merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent, c.param.MergeStrategy()) return chunkOutputs(flatten(merged)) } @@ -719,7 +719,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.MergeStrategy()) + attached = mergeByTokenSizeFromJSON([][]schema.ChunkDoc{flatten(attached)}, c.param.ChunkTokenSize, c.param.OverlappedPercent, c.param.MergeStrategy()) } flat := flatten(attached) @@ -956,15 +956,16 @@ func takeFromStart(text string, tokens int) string { // mergeByTokenSizeFromJSON mirrors Python naive_merge's projected-total // hard cap (rag/nlp/__init__.py after the strict chunk_token_num fix). -// Oversized text units are sub-split via splitOversizedUnit before merge; -// overlap is applied only when overlap+segment still fits the budget. +// Over-budget units are never atom-split: each one stands alone as its own +// chunk (Python naive_merge behavior, #17808 OVER_CAP contract). Overlap is +// applied only when overlap+segment still fits the budget. // // 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 { +func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64, 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 { @@ -1059,30 +1060,10 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over addTextChunk(ck) continue } - // Over-budget unit. - if !subSplitOversize { - // JSON path: Python keeps each over-budget item whole — it does - // not sub-split a single item, so emit it as one chunk. - addTextChunk(ck) - continue - } - // Text path: hard-cap atomic oversize units before merge, matching - // Python's _split_oversized_unit. - slog.Debug("TokenChunker: splitting oversized JSON unit via splitOversizedUnit", - "len", len(ck.Text), "tokens", tk, "chunk_token_size", chunkTokens) - for _, piece := range splitOversizedUnit(ck.Text, chunkTokens) { - if strings.TrimSpace(piece) == "" { - continue - } - cp := cloneChunkDoc(ck) - cp.Text = piece - cp.TKNums = intPtr(tokenizeStr(piece)) - // Coordinates stay on the first piece only to avoid duplicating - // PDF bboxes across atom slices. - addTextChunk(cp) - ck.PDFPositions = nil - ck.Positions = nil - } + // Over-budget unit: keep it whole. Python's naive_merge never + // sub-splits a single item, so emit it as one chunk and let the + // model layer truncate it later. + addTextChunk(ck) } perItem[idx] = merged } diff --git a/internal/ingestion/component/chunker/token_batch1_test.go b/internal/ingestion/component/chunker/token_batch1_test.go index 8e5b506065..8df8139fa5 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, schema.MergeOverCap) + got := mergeByTokenSizeFromJSON(items, budget, 30.0, 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, schema.MergeOverCap) + got := mergeByTokenSizeFromJSON(items, budget, 0.0, schema.MergeOverCap) merged := got[0] // Expect: chunk0 (T1+T2, overflow-closed), N (non-text), chunk1 (T3+T4 merged). if len(merged) != 3 { @@ -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, schema.MergeUnderCap) + got := mergeByTokenSizeFromJSON(items, budget, 0.0, 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, schema.MergeOverCap) + at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100, schema.MergeOverCap) if at100 == nil || len(at100) == 0 { t.Fatalf("overlappedPct=100: nil/empty result") } - at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, true, schema.MergeOverCap) - atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, true, schema.MergeOverCap) + at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, schema.MergeOverCap) + atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, 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, schema.MergeOverCap) + at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0, schema.MergeOverCap) if at0 == nil || len(at0) == 0 { t.Fatalf("overlappedPct=0: nil/empty result") } - atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, true, schema.MergeOverCap) - atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, true, schema.MergeOverCap) + atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, schema.MergeOverCap) + atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, 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, schema.MergeOverCap) + got := mergeByTokenSizeFromJSON(items, 128, 0, 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 71945bec53..d0ab422567 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, schema.MergeOverCap) + got := mergeByTokenSizeFromJSON(items, 128, 0, 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, schema.MergeOverCap) + got := mergeByTokenSizeFromJSON(items, 128, 0, 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, schema.MergeOverCap) + got := mergeByTokenSizeFromJSON(items, 128, 0, 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 3255d5ff55..209cfc5884 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, schema.MergeOverCap) + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0, 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, schema.MergeOverCap) + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20, schema.MergeOverCap) unit := tokenizeStr(sections[0].Text) for i, ck := range got[0] { // OVER_CAP allows one boundary overflow (prev + one unit). The JSON @@ -150,31 +150,20 @@ func TestMergeByTokenSizeFromJSON_OverlapDroppedAtOverflow(t *testing.T) { } } -func TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit(t *testing.T) { - // A single unit larger than the budget must be atom-split before merge. +func TestMergeByTokenSizeFromJSON_OversizedUnitStaysWhole(t *testing.T) { + // Per the #17808 contract an over-budget unit is never atom-split: it + // stands alone as one chunk and the model layer truncates it later. const budget = 30 long := strings.TrimSpace(strings.Repeat("word ", 100)) items := [][]schema.ChunkDoc{{ {Text: long, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(long))}, }} - 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])) + got := mergeByTokenSizeFromJSON(items, budget, 0, schema.MergeOverCap) + if len(got) != 1 || len(got[0]) != 1 { + t.Fatalf("over-budget unit must stay whole, got %d chunk(s)", len(got[0])) } - // cl100k is not additive across whitespace joins: token(a)+token(b) can be - // one less than token(a+b), so the running-sum flush used by both Python's - // _split_oversized_unit and the aligned Go port can leave a piece exactly - // one token over the nominal budget (each sub-split piece <= budget+1). - // OVER_CAP then merges at most two such pieces into one chunk before - // closing it, so the invariant we defend is that no chunk exceeds - // 2*(budget+1): the oversized unit is sub-split (not collapsed into one - // chunk) and at most one boundary overflow is allowed — matching the - // Python reference. - const slack = 2 * (budget + 1) - for i, ck := range got[0] { - if n := tokenizeStr(ck.Text); n > slack { - t.Errorf("chunk %d exceeds 2*(budget+1): tokens=%d (cap=%d)", i, n, budget) - } + if got[0][0].Text != long { + t.Errorf("over-budget chunk text changed: got %q", got[0][0].Text) } }