From 2fcc34904bd4f4184e62ebbeb147e388e9505bc4 Mon Sep 17 00:00:00 2001 From: Jack Date: Wed, 5 Aug 2026 18:53:59 +0800 Subject: [PATCH] fix(chunker): keep oversize text/markdown unit whole (OVER_CAP alignment) (#17854) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Go `TokenChunker` text/markdown path (`mergeByTokenSize`) unconditionally called `splitOversizedUnit` on any unit that exceeded `chunk_token_size`, emitting Go-only sub-chunks. Python's `naive_merge` (`_merge_paragraph_groups`, `rag/nlp/__init__.py`) never atom-splits an oversize unit under either `OVER_CAP` or `UNDER_CAP`: a paragraph larger than the budget becomes its own standalone chunk and the model layer truncates it later. This aligns the text/markdown path with the **structured JSON path** (`invokeJSONPayload` → `mergeByTokenSizeFromJSON(..., subSplitOversize=false)`, #17739). It completes the OVER_CAP alignment started in #17835. --- internal/ingestion/component/chunker/token.go | 34 ++--- .../chunker/token_oversize_whole_test.go | 139 ++++++++++++++++++ .../chunker/token_strict_cap_test.go | 44 ------ 3 files changed, 153 insertions(+), 64 deletions(-) create mode 100644 internal/ingestion/component/chunker/token_oversize_whole_test.go diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index 5d4167fac3..b089cf233f 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -547,10 +547,13 @@ func newChunkText(prevText, incoming string, target int, overlapPct float64, inc // mergeByTokenSize implements exact token-based chunk merging that mirrors // Python's naive_merge (rag/nlp/__init__.py) after the strict chunk_token_num // hard-cap fix. It uses tokenizeStr for precise token counting, treats the -// payload as a single section, splits oversized sections on production sentence -// delimiters, hard-caps atomic oversize units via splitOversizedUnit, and merges -// only when the projected total stays within chunk_token_size. Overlap is -// applied only when the resulting chunk still fits the budget. +// payload as a single section, and splits oversized sections on production +// sentence delimiters. An oversize unit (a single paragraph larger than the +// token budget) is kept whole as a standalone chunk — matching Python OVER_CAP, +// where the model layer truncates it later — instead of being atom-split. +// Sections are merged only when the projected total stays within +// chunk_token_size. Overlap is applied only when the resulting chunk still +// fits the budget. func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *regexp.Regexp) map[string]any { target := c.param.ChunkTokenSize overlapPct := c.param.OverlappedPercent @@ -608,18 +611,6 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r } } - addUnit := func(unit string) { - if tokenizeStr(unit) <= target { - addChunk(unit) - return - } - slog.Debug("TokenChunker: splitting oversized unit via splitOversizedUnit", - "len", len(unit), "tokens", tokenizeStr(unit), "chunk_token_size", target) - for _, piece := range splitOversizedUnit(unit, target) { - addChunk(piece) - } - } - for _, sec := range sections { sec = strings.TrimSpace(sec) if sec == "" { @@ -630,8 +621,11 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r addChunk(t) continue } - // Oversized section: split on production sentence delimiters, then - // hard-cap any unit that still exceeds the budget (unbroken atoms). + // Oversized section: split on production sentence delimiters into + // units. An oversize unit (still exceeds the budget) is passed through + // addChunk and kept whole — no atom-split, matching Python + // naive_merge. mergeDecision forces an oversize incoming unit to + // startNewChunk, so it stands alone as its own chunk. parts := sentenceDelimiter.Split(sec, -1) hadPart := false for _, part := range parts { @@ -651,10 +645,10 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r continue } hadPart = true - addUnit("\n" + part) + addChunk("\n" + part) } if !hadPart { - addUnit(t) + addChunk(t) } } diff --git a/internal/ingestion/component/chunker/token_oversize_whole_test.go b/internal/ingestion/component/chunker/token_oversize_whole_test.go new file mode 100644 index 0000000000..a7065cfef7 --- /dev/null +++ b/internal/ingestion/component/chunker/token_oversize_whole_test.go @@ -0,0 +1,139 @@ +// Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package chunker + +import ( + "context" + "strings" + "testing" +) + +// TestTokenChunker_OversizeUnitKeptWhole pins the Python OVER_CAP contract for +// the text/markdown path: a single paragraph that exceeds chunk_token_size is +// kept as one standalone chunk. Python's naive_merge (_merge_paragraph_groups, +// rag/nlp/__init__.py) never atom-splits an oversize unit; it is kept whole and +// the model layer truncates it later. An unbroken input line with no delimiter +// forces the oversize path while isolating it from the delimiter-splitting logic. +func TestTokenChunker_OversizeUnitKeptWhole(t *testing.T) { + var longLine = strings.Repeat("word ", 400) // ~400 tokens, far above the 32 budget + + cases := []struct { + name string + conf map[string]any + input map[string]any + }{ + { + name: "text path", + conf: map[string]any{"chunk_token_size": 32, "delimiters": []string{"\n"}}, + input: map[string]any{ + "name": "t", "output_format": "text", "text": longLine, + }, + }, + { + name: "markdown path", + conf: map[string]any{"chunk_token_size": 32, "delimiters": []string{"\n"}}, + input: map[string]any{ + "name": "t", "output_format": "markdown", "markdown": longLine, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c, err := NewTokenChunker(tc.conf) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out, err := c.Invoke(context.Background(), nil, tc.input) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, ok := out["chunks"].([]map[string]any) + if !ok { + t.Fatalf("chunks missing or wrong type: %T", out["chunks"]) + } + if len(chunks) != 1 { + t.Fatalf("oversize unit: want 1 standalone chunk, got %d", len(chunks)) + } + if got, _ := chunks[0]["text"].(string); strings.TrimSpace(got) != strings.TrimSpace(longLine) { + t.Fatalf("oversize unit content not preserved: got %q", got) + } + }) + } +} + +// TestTokenChunker_OversizeUnitStandsAloneAfterInBudgetUnit exercises the +// mergeDecision oversize branch (incomingTokens > target -> startNewChunk), +// which TestTokenChunker_OversizeUnitKeptWhole never reaches because its lone +// oversize unit goes through the len(cks)==0 path of addChunk. A short +// in-budget sentence precedes the oversize paragraph; after the sentence +// delimiter split the oversize unit must stand alone as its own chunk +// (matching Python OVER_CAP), not be merged into or atom-split across the +// previous chunk. +func TestTokenChunker_OversizeUnitStandsAloneAfterInBudgetUnit(t *testing.T) { + var longLine = strings.Repeat("word ", 400) // ~400 tokens, far above the 32 budget + inBudget := "Hello world." // ASCII period is not a sentence delimiter; fits 32 + + cases := []struct { + name string + conf map[string]any + input map[string]any + }{ + { + name: "text path", + conf: map[string]any{"chunk_token_size": 32, "delimiters": []string{"\n"}}, + input: map[string]any{ + "name": "t", "output_format": "text", + "text": inBudget + "\n" + longLine, + }, + }, + { + name: "markdown path", + conf: map[string]any{"chunk_token_size": 32, "delimiters": []string{"\n"}}, + input: map[string]any{ + "name": "t", "output_format": "markdown", + "markdown": inBudget + "\n" + longLine, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + c, err := NewTokenChunker(tc.conf) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out, err := c.Invoke(context.Background(), nil, tc.input) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, ok := out["chunks"].([]map[string]any) + if !ok { + t.Fatalf("chunks missing or wrong type: %T", out["chunks"]) + } + if len(chunks) != 2 { + t.Fatalf("oversize after in-budget: want 2 chunks (in-budget + standalone oversize), got %d", len(chunks)) + } + first, _ := chunks[0]["text"].(string) + if first != inBudget { + t.Fatalf("first chunk should be only the in-budget sentence %q, got %q", inBudget, first) + } + second, _ := chunks[1]["text"].(string) + if second != strings.TrimSpace(longLine) { + t.Fatalf("oversize chunk should be the whole long line kept whole, got %q", second) + } + }) + } +} diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index 57e9800aa6..3255d5ff55 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -268,50 +268,6 @@ func TestMergeByTokenSize_UnderCapNoOverflow(t *testing.T) { } } -func TestMergeByTokenSize_UnbrokenAtomStrictCap(t *testing.T) { - // Unbroken dense string (no whitespace / sentence delim) must still - // hard-cap via the character-window fallback inside splitOversizedUnit. - const budget = 20 - // Use many distinct ASCII letters so cl100k does not collapse the whole - // run into a handful of tokens. - var b strings.Builder - for i := 0; i < 400; i++ { - b.WriteByte(byte('a' + i%26)) - } - text := b.String() - if tokenizeStr(text) <= budget { - t.Skipf("tokenizer collapsed unbroken atom to %d tokens (<= budget)", tokenizeStr(text)) - } - comp, err := NewTokenChunker(map[string]any{ - "delimiter_mode": "token_size", - "chunk_token_size": budget, - }) - if err != nil { - t.Fatalf("NewTokenChunker: %v", err) - } - tc := comp.(*TokenChunkerComponent) - out := tc.mergeByTokenSize(text, nil) - chunks, _ := out["chunks"].([]map[string]any) - if len(chunks) < 2 { - t.Fatalf("want multiple chunks for unbroken atom, got %d (total_tokens=%d)", len(chunks), tokenizeStr(text)) - } - var joined strings.Builder - for i, ck := range chunks { - s, _ := ck["text"].(string) - joined.WriteString(s) - // Sub-split pieces are <= budget+1; OVER_CAP merges at most two before - // closing, so a chunk can reach 2*(budget+1). - if n := tokenizeStr(s); n > 2*(budget+1) { - t.Errorf("chunk %d exceeds 2*(budget+1): tokens=%d text=%q", i, n, s) - } - } - // mergeByTokenSize prefixes "\n" on sections; stripping newlines recovers - // the original unbroken atom. - if strings.ReplaceAll(joined.String(), "\n", "") != text { - t.Errorf("content not preserved after stripping newlines: got %q", joined.String()) - } -} - func TestInvokeTextPayload_StrictCapEndToEnd(t *testing.T) { const budget = 32 unit := tokenizeStr(strings.TrimSpace(strings.Repeat("alpha ", 12)))