From 109b74e41071afd3c97f7cdb0eb7ea6a7dd321e2 Mon Sep 17 00:00:00 2001 From: Jack Date: Thu, 6 Aug 2026 15:50:52 +0800 Subject: [PATCH] refactor(go): remove dead atom-split helpers from TokenChunker (#17920) ## Summary - Remove `splitOversizedUnit`, `splitAtomByTokenBudget` and `atomRE` from `internal/ingestion/component/chunker/token.go`. - Delete `split_oversized_guard_test.go` (added by #17740), which guarded the removed atom-split behaviour. - Drop the now-unused `wordCount`/`charCount` helpers from `token_strict_cap_test.go`. - Add `TestMergeByTokenSize_OversizedUnitStaysWhole` to pin the #17799 contract invariant (over-budget unit stays whole, never atom-split) on the **text path**. The JSON path is already covered by `TestMergeByTokenSizeFromJSON_OversizedUnitStaysWhole`. ## Why The production merge path (`mergeByTokenSize` / `mergeByTokenSizeFromJSON`) keeps over-budget units whole and relies on the embedding/rerank layer to truncate them, per the TokenChunker contract (#17799: remove atom-split, no hard_cap). The deleted helpers implemented the opposite behaviour and had **no production caller**, so they contradicted the contract and misled readers into thinking atom-split was active. ## Parser vs chunker layering Python's `_split_oversized_unit` lives at the **parser layer** (pre-split before `naive_merge`), not in the chunker. Go's parser backends are currently skeletons, so there is no parser-side equivalent yet; if added later it belongs in `internal/parser/parser/*`, not the chunker. ## Test plan `bash build.sh --test ./internal/ingestion/component/chunker/...` passes; the new text-path test passes and the orphaned atom-split tests are gone. ## Changes - 3 files changed, 32 insertions(+), 250 deletions(-) --- .../chunker/split_oversized_guard_test.go | 82 ------------ internal/ingestion/component/chunker/token.go | 122 ++---------------- .../chunker/token_strict_cap_test.go | 89 +++++-------- 3 files changed, 43 insertions(+), 250 deletions(-) delete mode 100644 internal/ingestion/component/chunker/split_oversized_guard_test.go diff --git a/internal/ingestion/component/chunker/split_oversized_guard_test.go b/internal/ingestion/component/chunker/split_oversized_guard_test.go deleted file mode 100644 index 18001bfa50..0000000000 --- a/internal/ingestion/component/chunker/split_oversized_guard_test.go +++ /dev/null @@ -1,82 +0,0 @@ -// -// 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 ( - "slices" - "testing" -) - -// b1Text is the single long paragraph from parity case -// token__b1_count_sensitive. With no effective delimiter and a small -// chunk_token_size, a live tokenizer splits it into many chunks, while a -// dead (zero-counting) tokenizer collapses it into one. These two tests -// pin both behaviours so the running-sum alignment in splitOversizedUnitWith -// (token.go) stays byte-identical to Python rag/nlp._split_oversized_unit -// and a silently dead tokenizer is caught. -const b1Text = "RAGFlow is a retrieval augmented generation engine that ingests heterogeneous documents and slices them into retrievable passages. The chunker must respect a token budget so that each passage fits the context window of the downstream language model without truncation. Token counting is the linchpin of this budget because every merge decision reads the per segment token length. When the encoder is missing the chunker silently reports zero tokens for every string and the budget is never exceeded so an entire document collapses into one oversized chunk. That degenerate behaviour is invisible to a test that only checks for a non empty result because zero is a number like any other. The only way to notice is to compare the actual chunk count against a reference implementation that counts tokens correctly. This case pins that comparison with a long single paragraph and a deliberately small budget so the correct count yields many chunks and a dead encoder yields exactly one. Retrieval quality depends on passages being coherent and appropriately sized so the guard is not merely cosmetic but protects the core promise of the system." - -// TestSplitOversizedUnitRunningSumMatchesPython pins the whitespace-atom -// sub-split of splitOversizedUnitWith against Python's -// rag/nlp._split_oversized_unit. Go now uses the same running-sum flush -// (token.go: currentTokens+aTokens > budget), so the emitted pieces must -// match Python exactly. This is the positive guard that compensates the -// slack=1 relaxation in token_strict_cap_test.go: the oversized unit is -// sub-split on the same boundaries Python uses, not merely kept under a -// loose cap. -func TestSplitOversizedUnitRunningSumMatchesPython(t *testing.T) { - const budget = 50 - got := splitOversizedUnitWith(b1Text, budget, tokenizeStr) - - want := []string{ - "RAGFlow is a retrieval augmented generation engine that ingests heterogeneous documents and slices them into retrievable passages. The chunker must respect a token budget so that each passage fits the context window of the ", - "downstream language model without truncation. Token counting is the linchpin of this budget because every merge decision reads the per segment token length. When the encoder is missing the chunker silently reports zero tokens for every string and the budget ", - "is never exceeded so an entire document collapses into one oversized chunk. That degenerate behaviour is invisible to a test that only checks for a non empty result because zero is a number like any other. The only way to notice is ", - "to compare the actual chunk count against a reference implementation that counts tokens correctly. This case pins that comparison with a long single paragraph and a deliberately small budget so the correct count yields many chunks and a dead encoder yields exactly ", - "one. Retrieval quality depends on passages being coherent and appropriately sized so the guard is not merely cosmetic but protects the core promise of the system.", - } - - if !slices.Equal(got, want) { - t.Errorf("splitOversizedUnitWith diverges from Python _split_oversized_unit: got %d pieces, want %d", len(got), len(want)) - for i := 0; i < max(len(got), len(want)); i++ { - if i >= len(got) || i >= len(want) || got[i] != want[i] { - gt, wt := "", "" - if i < len(got) { - gt = got[i] - } - if i < len(want) { - wt = want[i] - } - t.Errorf("piece[%d]:\n got: %q\n want: %q", i, gt, wt) - } - } - } -} - -// TestSplitOversizedUnitDeadTokenizerCollapses guards against a silently -// dead tokenizer. With a zero-counting function the running-sum flush never -// triggers, so the whole paragraph collapses into exactly one chunk — -// diverging from the live multi-chunk baseline above. A test that only -// checks for a non-empty result would miss this, so we assert the exact -// collapse count. -func TestSplitOversizedUnitDeadTokenizerCollapses(t *testing.T) { - const budget = 50 - got := splitOversizedUnitWith(b1Text, budget, func(string) int { return 0 }) - if len(got) != 1 { - t.Fatalf("dead tokenizer must collapse paragraph into exactly one chunk, got %d", len(got)) - } -} diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index b74dd8de6a..81dc412c84 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -53,6 +53,19 @@ // generated on demand for text chunks that carry PDF positions: // cropImageChunks crops the text region and writes a preview image, // then imageUploadDecorator uploads it to img_id. See pdfcrop_cgo.go. +// +// - OVER-BUDGET UNITS (contract #17799): a single item that exceeds +// chunk_token_size is KEPT WHOLE as its own chunk and is NOT +// atom-split; the embedding/rerank layer truncates it later. The +// TokenChunker must never sub-split a single item (the naive_merge +// invariant). If oversized-unit handling is ever needed to avoid a +// single mega-chunk, it belongs at the CHUNKER side (or a dedicated +// PreSplitter stage between Parser and Chunker), fed by an EXPLICIT +// token budget + tokenizer — NOT in the parser, and NOT as a +// char-window atom-split. The earlier splitOversizedUnit / +// splitAtomByTokenBudget helpers were a misplaced (parser-layer logic +// wrongly living in the chunker) and unwired vestige; they were +// removed to align with this contract. package chunker import ( @@ -347,115 +360,6 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string // it would diverge from Python's chunk boundaries. var sentenceDelimiter = regexp.MustCompile(`(\n|[!?。；！？])`) -// atomRE matches whitespace runs or non-whitespace runs. Mirrors Python -// `_split_oversized_unit`'s `re.findall(r"\s+|\S+", text)`. -var atomRE = regexp.MustCompile(`\s+|\S+`) - -// splitAtomByTokenBudget splits a single non-whitespace atom into -// substrings that each have <= chunkTokenNum tokens. Mirrors Python -// rag/nlp._split_atom_by_token_budget (binary search on rune prefixes). -func splitAtomByTokenBudget(atom string, chunkTokenNum int, countFn func(string) int) []string { - if atom == "" { - return nil - } - if countFn == nil { - countFn = tokenizeStr - } - if countFn(atom) <= chunkTokenNum { - return []string{atom} - } - runes := []rune(atom) - var pieces []string - start := 0 - n := len(runes) - for start < n { - low := start + 1 - high := n - bestEnd := start + 1 - for low <= high { - mid := (low + high) / 2 - if countFn(string(runes[start:mid])) <= chunkTokenNum { - bestEnd = mid - low = mid + 1 - } else { - high = mid - 1 - } - } - pieces = append(pieces, string(runes[start:bestEnd])) - start = bestEnd - } - return pieces -} - -// splitOversizedUnit splits a unit that exceeds chunkTokenNum tokens into -// pieces that each fit the budget. Whitespace is the primary break (mirrors -// Python rag/nlp._split_oversized_unit / HtmlParser._split_oversized_block); -// a single non-whitespace run longer than the budget falls back to -// token-budget-based character windows. -func splitOversizedUnit(text string, chunkTokenNum int) []string { - return splitOversizedUnitWith(text, chunkTokenNum, tokenizeStr) -} - -func splitOversizedUnitWith(text string, chunkTokenNum int, countFn func(string) int) []string { - if countFn == nil { - countFn = tokenizeStr - } - if countFn(text) <= chunkTokenNum { - return []string{text} - } - var pieces []string - current := "" - // Running sum of per-atom token counts for the current piece. Mirrors - // Python rag/nlp._split_oversized_unit's `current_tokens`. We flush when - // this running sum (not the exact count of the joined string) would - // exceed the budget, because cl100k token counting is not additive across - // whitespace joins: token(a)+token(b) can differ from token(a+b), so the - // joined-string fit check drifts one atom off Python's boundary. - currentTokens := 0 - tokenCache := map[string]int{} - - atomTokens := func(atom string) int { - // Whitespace-only atoms contribute 0 in isolation (mirrors Python - // atom.isspace()), matching the packing heuristic used by - // rag/nlp._split_oversized_unit. - if strings.TrimSpace(atom) == "" { - return 0 - } - if n, ok := tokenCache[atom]; ok { - return n - } - n := countFn(atom) - tokenCache[atom] = n - return n - } - - for _, atom := range atomRE.FindAllString(text, -1) { - aTokens := atomTokens(atom) - if aTokens > chunkTokenNum && strings.TrimSpace(atom) != "" { - if current != "" { - pieces = append(pieces, current) - current = "" - currentTokens = 0 - } - pieces = append(pieces, splitAtomByTokenBudget(atom, chunkTokenNum, countFn)...) - continue - } - // Running-sum fit check, identical to Python's - // `current_tokens + a_tokens > chunk_token_num`. - if current != "" && currentTokens+aTokens > chunkTokenNum { - pieces = append(pieces, current) - current = "" - currentTokens = 0 - } - current += atom - currentTokens += aTokens - } - if current != "" { - pieces = append(pieces, current) - } - return pieces -} - // computeOverlapPrefix returns (overlapText, overlapTokenCount) carved from // the tail of prevText after stripping parser tags. overlappedPct is a // percentage in [0, 100]. Mirrors Python rag/nlp._compute_overlap_prefix. diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index 209cfc5884..16f2c33664 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -20,69 +20,10 @@ import ( "context" "strings" "testing" - "unicode/utf8" "ragflow/internal/ingestion/component/schema" ) -// wordCount is a deterministic tokenizer stand-in used only via -// splitOversizedUnitWith in unit-level helper tests. -func wordCount(s string) int { - s = strings.TrimSpace(s) - if s == "" { - return 0 - } - return len(strings.Fields(s)) -} - -func charCount(s string) int { return utf8.RuneCountInString(s) } - -func TestSplitOversizedUnit_WhitespacePacksToBudget(t *testing.T) { - // 100 words, budget 30 → must yield multiple pieces, each ≤ 30 words. - text := strings.TrimSpace(strings.Repeat("word ", 100)) - pieces := splitOversizedUnitWith(text, 30, wordCount) - if len(pieces) < 2 { - t.Fatalf("want multiple pieces, got %d: %#v", len(pieces), pieces) - } - total := 0 - for _, p := range pieces { - n := wordCount(p) - if n > 30 { - t.Errorf("piece exceeds budget: tokens=%d text=%q", n, p) - } - total += n - } - if total != 100 { - t.Errorf("word count not preserved: got %d want 100", total) - } -} - -func TestSplitOversizedUnit_UnbrokenAtomFallsBackToCharWindows(t *testing.T) { - // Unbroken run with char-as-token counting — must sub-split on runes. - atom := strings.Repeat("a", 80) - pieces := splitOversizedUnitWith(atom, 50, charCount) - if len(pieces) < 2 { - t.Fatalf("want >=2 pieces for unbroken atom, got %d", len(pieces)) - } - joined := strings.Join(pieces, "") - if joined != atom { - t.Errorf("content not preserved: got %q", joined) - } - for _, p := range pieces { - if charCount(p) > 50 { - t.Errorf("piece exceeds budget: %d runes in %q", charCount(p), p) - } - } -} - -func TestSplitOversizedUnit_WithinBudgetUnchanged(t *testing.T) { - text := "hello world" - pieces := splitOversizedUnitWith(text, 100, wordCount) - if len(pieces) != 1 || pieces[0] != text { - t.Fatalf("within-budget text must be returned as-is, got %#v", pieces) - } -} - func TestComputeOverlapPrefix_StripsTagsAndCounts(t *testing.T) { prev := strings.Repeat("word ", 20) + "@@1\t2.3## tail" overlap, n := computeOverlapPrefix(prev, 30) @@ -198,6 +139,36 @@ func TestMergeByTokenSize_TextPathStrictCap(t *testing.T) { } } +// TestMergeByTokenSize_OversizedUnitStaysWhole mirrors +// TestMergeByTokenSizeFromJSON_OversizedUnitStaysWhole on the text path: a +// single block of text that exceeds chunk_token_size and contains no +// sentence delimiter must be emitted as ONE whole chunk — never atom-split. +// This pins the #17799 contract invariant on the text path (the JSON path +// is already covered by TestMergeByTokenSizeFromJSON_OversizedUnitStaysWhole). +func TestMergeByTokenSize_OversizedUnitStaysWhole(t *testing.T) { + const budget = 30 + // One long run with no '\n' / '!?' / '。；！？' delimiter: the text path + // splits oversized sections only on sentenceDelimiter, so this whole + // run is one unit that still exceeds the budget and must stay whole. + long := strings.TrimSpace(strings.Repeat("word ", 100)) + comp, err := NewTokenChunker(map[string]any{ + "delimiter_mode": "token_size", + "chunk_token_size": budget, + }) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out := comp.(*TokenChunkerComponent).mergeByTokenSize(long, nil) + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) != 1 { + t.Fatalf("over-budget unit must stay whole, got %d chunk(s)", len(chunks)) + } + text, _ := chunks[0]["text"].(string) + if text != long { + t.Errorf("over-budget chunk text changed: got %q", text) + } +} + func TestMergeByTokenSize_UnderCapNoOverflow(t *testing.T) { // UNDER_CAP (under_cap=true) must never let a chunk exceed the token // target: a projected join that would overflow starts a fresh chunk