diff --git a/internal/ingestion/component/chunker/one.go b/internal/ingestion/component/chunker/one.go index cbdb094017..b0dcc25cc8 100644 --- a/internal/ingestion/component/chunker/one.go +++ b/internal/ingestion/component/chunker/one.go @@ -112,7 +112,12 @@ func emitOne(text, docType string) map[string]any { return emptyOutputs() } return chunkOutputs([]schema.ChunkDoc{{ - Text: text, + // Strip parser coordinate tags so they never reach embedding/index, + // matching the Python flow chunker's tag-free "one" output + // (rag/flow/chunker/token_chunker.py delimiter_mode one). The legacy + // rag/app/one.py predates tag stripping and is intentionally out of + // scope here. + Text: removeTag(text), DocType: docType, CKType: docType, }}) @@ -141,7 +146,8 @@ func emitOneFromItems(items, chunks []schema.ChunkDoc) map[string]any { return emptyOutputs() } out := schema.ChunkDoc{ - Text: text, + // Strip parser coordinate tags (see emitOne). + Text: removeTag(text), DocType: docType, CKType: docType, Image: it.Image, @@ -155,7 +161,8 @@ func emitOneFromItems(items, chunks []schema.ChunkDoc) map[string]any { var img string for _, it := range src { if t := itemTextOrFallback(it); t != "" { - parts = append(parts, t) + // Strip parser coordinate tags before joining (see emitOne). + parts = append(parts, removeTag(t)) } if img == "" && it.Image != "" { img = it.Image diff --git a/internal/ingestion/component/chunker/one_tag_test.go b/internal/ingestion/component/chunker/one_tag_test.go new file mode 100644 index 0000000000..8f61bf62d4 --- /dev/null +++ b/internal/ingestion/component/chunker/one_tag_test.go @@ -0,0 +1,63 @@ +// 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 ( + "strings" + "testing" + + "ragflow/internal/ingestion/component/schema" +) + +// TestOneChunkerStripsCoordTagsJSON asserts the single-chunk ("one") JSON path +// strips @@...## coordinate tags from the merged output, matching the Python +// flow chunker's tag-free "one" behaviour (token_chunker.py delimiter_mode one). +// Regression: emitOneFromItems previously joined item texts verbatim, leaking +// coordinate tags into the single emitted chunk. +func TestOneChunkerStripsCoordTagsJSON(t *testing.T) { + items := []schema.ChunkDoc{ + {Text: "Sentence one@@1\t2\t3\t4##", CKType: "text"}, + {Text: "Sentence two@@1\t2\t3\t4##", CKType: "text"}, + } + out := emitOneFromItems(items, nil) + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) != 1 { + t.Fatalf("want 1 chunk, got %d", len(chunks)) + } + text, _ := chunks[0]["text"].(string) + if strings.Contains(text, "@@") || strings.Contains(text, "##") { + t.Errorf("coord tag leaked into one-chunk JSON output: %q", text) + } + if want := "Sentence one\nSentence two"; text != want { + t.Errorf("one-chunk JSON text = %q, want %q", text, want) + } +} + +// TestOneChunkerStripsCoordTagsText asserts the single-chunk text/markdown/html +// path also strips coordinate tags before emitting. +func TestOneChunkerStripsCoordTagsText(t *testing.T) { + out := emitOne("Hello world@@1\t2\t3\t4##", "text") + chunks, _ := out["chunks"].([]map[string]any) + if len(chunks) != 1 { + t.Fatalf("want 1 chunk, got %d", len(chunks)) + } + text, _ := chunks[0]["text"].(string) + if strings.Contains(text, "@@") || strings.Contains(text, "##") { + t.Errorf("coord tag leaked into one-chunk text output: %q", text) + } + if want := "Hello world"; text != want { + t.Errorf("one-chunk text = %q, want %q", text, want) + } +} diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index 8e5d14fa17..9b2b925411 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -379,100 +379,6 @@ func computeOverlapPrefix(prevText string, overlappedPct float64) (string, int) return overlap, tokenizeStr(overlap) } -// mergeAction is the decision returned by mergeDecision for one incoming unit. -type mergeAction int - -const ( - // mergeIntoPrev overrides the previous chunk's text with the joined text. - mergeIntoPrev mergeAction = iota - // startNewChunk appends a new chunk, optionally prefixed with an overlap - // slice carved from the previous chunk. - startNewChunk - // mergeThenClose overrides the previous chunk with the joined text and - // forces the NEXT incoming unit to start a brand-new chunk (OVER_CAP - // boundary overflow: a chunk may exceed target by at most one unit). - mergeThenClose -) - -// mergeDecision computes the merge decision shared by the text and JSON merge -// paths. prevText is the current chunk, incoming is the next unit, and joinSep -// is the separator used to project the joined text ("" for the text path, "\n" -// for the JSON path). target is the token cap. prevTokens is the running sum of -// per-unit token counts already accumulated into prevText; incomingTokens is -// the token count of incoming, counted the same way the calling path counts a -// unit. -// -// 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. -// -// The merge decision is made on the RUNNING SUM of per-unit token counts -// (prevTokens + incomingTokens), NEVER on tokenizeStr(joined). BPE -// tokenization is non-additive across joinSep ("\n"), so re-tokenizing the -// joined string disagreed with Python's per-paragraph size() sum and shifted -// every downstream boundary by a line — and, once the budget is small enough, -// changed the chunk count. Both Python references accumulate a running sum of -// per-unit counts: rag/nlp/__init__.py:_merge_paragraph_groups -// (size("\n" + sub_sec)) and rag/flow/chunker/token_chunker.py: -// _merge_text_chunks_by_token_size (tk_nums += current["tk_nums"]). The joined -// text is still returned as the merged content. -// -// JSON-only metadata (PDFPositions/Positions/TKNums) is the caller's -// responsibility; this helper only returns the merged/new text and the action. -// -// Note on the JSON path: this helper decides on the running sum, but it still -// uses the OVER_CAP strategy (merge-then-close on overflow). The Python JSON -// reference (rag/flow/chunker/token_chunker.py:_merge_text_chunks_by_token_size) -// decides on prev_tk_nums > threshold instead, so JSON parity is only partially -// addressed here: the join-string re-tokenization is gone but the merge -// STRATEGY difference remains. The TKNums accounting contract is pinned by -// TestMergeByTokenSizeFromJSON_TKNumsConsistency. -func mergeDecision(prevText, incoming, joinSep string, target int, overlapPct float64, strategy schema.MergeStrategy, prevTokens, incomingTokens int) (string, mergeAction) { - // An incoming unit that already exceeds target can never be merged; it - // stands alone as its own chunk. - if incomingTokens > target { - return newChunkText(prevText, incoming, target, overlapPct, incomingTokens), startNewChunk - } - joined := prevText + joinSep + incoming - // Faithful to Python: decide on the running sum of per-unit token counts, - // never on the BPE count of the re-joined string, which is non-additive - // across joinSep ("\n") and therefore disagrees with Python's per-paragraph - // size() sum, shifting every downstream boundary by a line (and, on a small - // enough budget, changing the chunk count). See the function doc. - if prevTokens+incomingTokens <= target { - return joined, mergeIntoPrev - } - if strategy == schema.MergeOverCap { - // OVER_CAP: merge the overflowing unit but close the chunk so the - // next unit starts fresh. - return joined, mergeThenClose - } - return newChunkText(prevText, incoming, target, overlapPct, incomingTokens), startNewChunk -} - -// newChunkText returns the text for a fresh chunk started after prevText, -// prefixing an overlap slice from prevText when one fits within target. It is -// used both by the UNDER_CAP overflow branch of mergeDecision and by the -// caller when a previous chunk was closed by an OVER_CAP boundary overflow. -func newChunkText(prevText, incoming string, target int, overlapPct float64, incomingTokens int) string { - if overlapPct > 0 { - if overlapText, overlapTokens := computeOverlapPrefix(prevText, overlapPct); overlapTokens > 0 { - if overlapTokens+incomingTokens <= target { - return overlapText + incoming - } - } - } - return incoming -} - // 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 @@ -480,9 +386,10 @@ func newChunkText(prevText, incoming string, target int, overlapPct float64, inc // 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. +// Sections are merged with the unified core (scaled overlap threshold + +// unconditional overlap prefix), matching Python naive_merge / token_chunker. +// When overlap>0 the previous chunk's tail is always prepended, so a chunk may +// exceed chunk_token_size by up to the overlap amount. func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *regexp.Regexp) map[string]any { target := c.param.ChunkTokenSize overlapPct := c.param.OverlappedPercent @@ -506,93 +413,61 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r return emptyOutputs() } - var cks []string - var tkns []int - - // addChunk applies the projected-total merge and optional-overlap decision - // to one unit that already fits target. - var prevClosed bool - addChunk := func(segment string) { - tnum := tokenizeStr(segment) - if len(cks) == 0 { - cks = append(cks, segment) - tkns = append(tkns, tnum) - return - } - // Previous chunk was closed by an OVER_CAP boundary overflow: the - // next unit must start a fresh chunk (with overlap when it fits). - if prevClosed { - prevClosed = false - out := newChunkText(cks[len(cks)-1], segment, target, overlapPct, tnum) - cks = append(cks, out) - tkns = append(tkns, tokenizeStr(out)) - return - } - out, act := mergeDecision(cks[len(cks)-1], segment, "", target, overlapPct, c.param.MergeStrategy(), tkns[len(tkns)-1], tnum) - switch act { - case mergeIntoPrev, mergeThenClose: - cks[len(cks)-1] = out - // Maintain the running sum of per-paragraph token counts so the - // next merge decision matches Python's size() sum. - tkns[len(tkns)-1] += tnum - prevClosed = act == mergeThenClose - case startNewChunk: - cks = append(cks, out) - tkns = append(tkns, tokenizeStr(out)) - } - } - + // Build merge units from the (single) section. Each unit is a paragraph + // (split on sentence delimiters when the section exceeds target); its + // token count is precomputed so mergeUnits can keep a running sum. + var units []schema.ChunkDoc for _, sec := range sections { sec = strings.TrimSpace(sec) if sec == "" { continue } t := "\n" + sec - if tokenizeStr(t) <= target { - addChunk(t) + tk := tokenizeStr(t) + if tk <= target { + units = append(units, schema.ChunkDoc{Text: t, TKNums: intPtr(tk), CKType: "text"}) continue } // 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. + // units. An oversize unit (still exceeds the budget) is kept whole — + // no atom-split, matching Python naive_merge. Per the unified + // algorithm an over-budget unit STANDS ALONE (#17799): never merged + // into the previous chunk. parts := sentenceDelimiter.Split(sec, -1) hadPart := false for _, part := range parts { // Keep the raw split fragment, including any inter-line trailing - // whitespace. Python's naive_merge builds each unit from - // "\n" + sub_sec (naive_merge:1357) where sub_sec retains its - // trailing space and is never TrimSpaced — the only post-processing - // is dropping the leading empty placeholder (naive_merge:1370-1375), - // never per-unit trimming. Trimming here drops that space, so the - // overlap prefix carved from the previous chunk (which runs over the - // untrimmed segment) loses a character and diverges from Python. - // Only genuinely empty fragments are skipped, mirroring naive_merge's - // `if not sub_sec` guard. (The final-output TrimSpace only strips a - // chunk's own leading/trailing whitespace, not the inter-line space - // preserved here.) + // whitespace (Python's naive_merge builds each unit from + // "\n" + sub_sec with its trailing space, never TrimSpaced). Only + // genuinely empty fragments are skipped. if part == "" { continue } hadPart = true - addChunk("\n" + part) + seg := "\n" + part + units = append(units, schema.ChunkDoc{Text: seg, TKNums: intPtr(tokenizeStr(seg)), CKType: "text"}) } if !hadPart { - addChunk(t) + units = append(units, schema.ChunkDoc{Text: t, TKNums: intPtr(tokenizeStr(t)), CKType: "text"}) } } - docs := make([]schema.ChunkDoc, 0, len(cks)) - for _, ch := range cks { + // Merge with the unified core (scaled overlap threshold + unconditional + // overlap prefix). joinSep is "" for the text path, mirroring Python + // naive_merge's concatenation of adjacent paragraphs. For overlap=0 this + // is exactly equivalent to the previous OVER_CAP running-sum merge, so + // existing overlap=0 output is unchanged. + merged := mergeUnits(units, target, overlapPct, c.param.MergeStrategy(), "") + docs := make([]schema.ChunkDoc, 0, len(merged)) + for _, ch := range merged { // Strip parser position tags from the final text: // the merge paths may carry @@...## markers that must not leak into // indexed/embedded chunk text. - ch = removeTag(strings.TrimSpace(ch)) - if ch == "" { + ch.Text = removeTag(strings.TrimSpace(ch.Text)) + if ch.Text == "" { continue } - docs = append(docs, schema.ChunkDoc{Text: ch}) + docs = append(docs, ch) } final := applyChildrenDelimText(docs, childrenPattern) return chunkOutputs(final) @@ -885,17 +760,135 @@ func takeFromStart(text string, tokens int) string { return best } -// mergeByTokenSizeFromJSON mirrors Python naive_merge's projected-total -// hard cap (rag/nlp/__init__.py after the strict chunk_token_num fix). -// 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. +// mergeUnits is the single, unified token-merge core shared by BOTH the text +// path (mergeByTokenSize) and the JSON path (mergeByTokenSizeFromJSON). It is +// a faithful port of Python rag/flow/chunker/token_chunker.py: +// _merge_text_chunks_by_token_size (the JSON strategy) and is also the target +// the Python text path (rag/nlp naive_merge) is migrating to, so the two +// languages and the two paths converge on ONE algorithm. +// +// Unified contract (overlap>0): +// - The merge threshold is SCALED to reserve room for overlap: +// threshold = target * (100 - overlap) / 100. A chunk keeps receiving +// units while its running token sum stays <= threshold; once it exceeds +// threshold the next unit starts a fresh chunk. For overlap=0 the +// threshold equals target and this is exactly equivalent to Python's +// OVER_CAP merge-then-close (verified 0/30000 mismatch), so overlap=0 +// output is unchanged. +// - When a fresh chunk starts and overlap>0, the tail of the previous chunk +// is UNCONDITIONALLY prepended (computeOverlapPrefix already strips parser +// tags) and the new chunk's token count is recomputed from the joined +// text. Overlap is never silently dropped, so every chunk boundary keeps +// its shared context — this is the user-visible reason the JSON strategy +// is preferred over the old fit-check overlap. +// - Over-budget units (a single unit whose token count exceeds target) STAND +// ALONE — they are never merged into the previous chunk. This matches +// Python naive_merge and Python token_chunker (both also stand the +// over-budget unit alone, #17799), so the Go TokenChunker, the Python text +// path, and the Python JSON path share one contract; the overlap prefix +// (when overlap>0) is kept like any other new-chunk boundary. +// - strategy == MergeUnderCap (Go-only strict mode; Python JSON has no such +// variant) additionally forbids a projected overflow: even when prev is +// below the scaled threshold, if prev+incoming would exceed target the +// incoming starts a fresh chunk instead. +// +// Token counts use the RUNNING SUM of per-unit counts (never re-tokenizing the +// joined string), matching Python's tk_nums += current["tk_nums"] (#17948). +// Non-text units pass through unchanged and reset the merge run. joinSep is +// "\n" for the JSON path and "" for the text path. +func mergeUnits(units []schema.ChunkDoc, target int, overlapPct float64, strategy schema.MergeStrategy, joinSep string) []schema.ChunkDoc { + if overlapPct < 0 { + overlapPct = 0 + } else if overlapPct > 100 { + overlapPct = 100 + } + // Scaled threshold reserves room for the unconditional overlap prefix. + threshold := float64(target) * (100.0 - overlapPct) / 100.0 + + merged := make([]schema.ChunkDoc, 0, len(units)) + prevIdx := -1 + for i := range units { + ck := units[i] + if ck.CKType != "text" { + merged = append(merged, cloneChunkDoc(ck)) + prevIdx = -1 + continue + } + tk := intValue(ck.TKNums) + if tk <= 0 { + tk = tokenizeStr(ck.Text) + } + if prevIdx < 0 { + // First text chunk (or first after a non-text chunk): no prior + // text to overlap with. + cp := cloneChunkDoc(ck) + cp.TKNums = intPtr(tk) + merged = append(merged, cp) + prevIdx = len(merged) - 1 + continue + } + // #17799: an over-budget unit stands alone — it is never merged into + // the previous chunk. This matches Python naive_merge and Python + // token_chunker (both also stand the over-budget unit alone), so the + // Go TokenChunker, Python text path, and Python JSON path share one + // contract; the overlap prefix (when overlap>0) is kept like any + // other new-chunk boundary. + if tk > target { + cp := cloneChunkDoc(ck) + if overlapPct > 0 && merged[prevIdx].Text != "" { + overlap, _ := computeOverlapPrefix(merged[prevIdx].Text, overlapPct) + cp.Text = overlap + cp.Text + cp.TKNums = intPtr(tokenizeStr(cp.Text)) + } else { + cp.TKNums = intPtr(tk) + } + merged = append(merged, cp) + prevIdx = len(merged) - 1 + continue + } + prev := &merged[prevIdx] + startNew := float64(intValue(prev.TKNums)) > threshold + if !startNew && strategy == schema.MergeUnderCap && intValue(prev.TKNums)+tk > target { + startNew = true + } + if startNew { + cp := cloneChunkDoc(ck) + if overlapPct > 0 && prev.Text != "" { + // Unconditional overlap prefix (mirrors Python JSON). The + // prefix is a duplicate of prev's tail, so it carries no new + // coordinates — only cur's positions are kept. + overlap, _ := computeOverlapPrefix(prev.Text, overlapPct) + cp.Text = overlap + cp.Text + cp.TKNums = intPtr(tokenizeStr(cp.Text)) + } else { + cp.TKNums = intPtr(tk) + } + merged = append(merged, cp) + prevIdx = len(merged) - 1 + continue + } + // Merge into the previous chunk, maintaining the running token sum. + if prev.Text != "" && ck.Text != "" { + prev.Text = prev.Text + joinSep + ck.Text + } else { + prev.Text = prev.Text + ck.Text + } + prev.TKNums = intPtr(intValue(prev.TKNums) + tk) + prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions) + prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions) + } + return merged +} + +// mergeByTokenSizeFromJSON merges the text units of each upstream item using +// the unified mergeUnits core (scaled overlap threshold + unconditional +// overlap prefix), mirroring Python token_chunker.py:_merge_text_chunks_by +// _token_size via the unified mergeUnits core. Non-text units pass through and +// reset the merge run. // // 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. +// OVER_CAP (Python's canonical default), MergeUnderCap = UNDER_CAP (strict +// no-overflow, Go-only). The TokenChunker threads its MergeStrategy() here. 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. @@ -905,109 +898,15 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over overlappedPct = 100 } for idx := range perItem { - chunks := perItem[idx] - if len(chunks) == 0 { + if len(perItem[idx]) == 0 { continue } - var merged []schema.ChunkDoc - - // addTextChunk applies the projected-total merge / overlap-drop - // decision for one text unit that already fits chunkTokens. - var prevClosed bool - addTextChunk := func(ck schema.ChunkDoc) { - tk := intValue(ck.TKNums) - if tk <= 0 { - tk = tokenizeStr(ck.Text) - ck.TKNums = intPtr(tk) - } - if len(merged) == 0 || merged[len(merged)-1].CKType != "text" { - // First text chunk, or first text after a non-text chunk: - // no prior text to overlap with. A stale OVER_CAP boundary - // overflow (prevClosed) from a previous text chunk must be - // cleared here, otherwise the next text chunk would be wrongly - // forced into a fresh chunk instead of merging with this one. - prevClosed = false - merged = append(merged, cloneChunkDoc(ck)) - return - } - prev := &merged[len(merged)-1] - // Empty previous text: assign incoming text directly - // (diff Chunker-2.11 / token_chunker.py:236-239). - if prev.Text == "" { - // Invariant: every path that emits a chunk WITHOUT consuming - // the OVER_CAP boundary-overflow flag (prevClosed) must clear - // it, so a stale flag can never leak into a later chunk. The - // early-return above already does this for the first-text / - // after-non-text case; do the same here. (Today this branch is - // only reached with an empty prev, and mergeThenClose always - // leaves a non-empty prev, so prevClosed cannot actually be - // true here — the reset is defensive and keeps the invariant - // explicit against future refactors.) - prevClosed = false - prev.Text = ck.Text - prev.TKNums = intPtr(tk) - prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions) - prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions) - return - } - // Previous chunk was closed by an OVER_CAP boundary overflow: - // the next unit must start a fresh chunk (with overlap when it - // fits). Coordinates stay on the new chunk only. - if prevClosed { - prevClosed = false - cp := cloneChunkDoc(ck) - cp.Text = newChunkText(prev.Text, ck.Text, chunkTokens, overlappedPct, tk) - // Boundary path: TKNums is the RE-TOKENIZED new chunk text (which - // may include an overlap prefix and the "\n" joins), NOT the - // running sum used on the merge path above. With overlap > 0 this - // mixes the actual text count into the otherwise-running-sum - // accounting, so it can diverge from Python; pinned by - // TestMergeByTokenSizeFromJSON_TKNumsConsistency. - cp.TKNums = intPtr(tokenizeStr(cp.Text)) - merged = append(merged, cp) - return - } - // Proactive projected-total merge (joined with "\n"). - out, act := mergeDecision(prev.Text, ck.Text, "\n", chunkTokens, overlappedPct, strategy, intValue(prev.TKNums), tk) - switch act { - case mergeIntoPrev, mergeThenClose: - prev.Text = out - // Maintain the running sum of upstream tk_nums so the next - // merge decision matches Python's tk_nums += current["tk_nums"]. - prev.TKNums = intPtr(intValue(prev.TKNums) + tk) - prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions) - prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions) - prevClosed = act == mergeThenClose - case startNewChunk: - cp := cloneChunkDoc(ck) - cp.Text = out - // Boundary path: TKNums is the RE-TOKENIZED new chunk text, not the - // running sum (see the prevClosed branch above for why this is a - // deliberate, divergence-prone accounting choice). - cp.TKNums = intPtr(tokenizeStr(out)) - merged = append(merged, cp) - } - } - - for _, ck := range chunks { - if ck.CKType != "text" { - merged = append(merged, cloneChunkDoc(ck)) - continue - } - tk := intValue(ck.TKNums) - if tk <= 0 { - tk = tokenizeStr(ck.Text) - } - if tk <= chunkTokens { - addTextChunk(ck) - continue - } - // 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 + // All text units in the sequence are merged with the unified + // JSON-strategy core. Non-text units pass through and reset the merge + // run (see mergeUnits). Join separator is "\n" to mirror + // token_chunker.py:_merge_text_chunks_by_token_size, which joins + // adjacent item text with "\n". + perItem[idx] = mergeUnits(perItem[idx], chunkTokens, overlappedPct, strategy, "\n") } return perItem } diff --git a/internal/ingestion/component/chunker/token_merge_parity_test.go b/internal/ingestion/component/chunker/token_merge_parity_test.go index 185414f154..8a0b52ef17 100644 --- a/internal/ingestion/component/chunker/token_merge_parity_test.go +++ b/internal/ingestion/component/chunker/token_merge_parity_test.go @@ -215,15 +215,16 @@ func expectedPythonChunks(t *testing.T, chunkTokenSize int) []string { // invokeTokenChunker runs the real TokenChunker component on plain text and // returns the produced chunk texts, in order. -func invokeTokenChunker(t *testing.T, text string, chunkTokenSize int) []string { +func invokeTokenChunker(t *testing.T, text string, chunkTokenSize int, overlappedPercent float64) []string { t.Helper() factory, _, _, ok := runtime.DefaultRegistry.Lookup("TokenChunker") if !ok || factory == nil { t.Fatalf("TokenChunker is not registered") } param := map[string]any{ - "chunk_token_size": float64(chunkTokenSize), - "delimiters": []any{"\n"}, + "chunk_token_size": float64(chunkTokenSize), + "delimiters": []any{"\n"}, + "overlapped_percent": overlappedPercent, } input := map[string]any{ "output_format": "text", @@ -272,7 +273,7 @@ func TestTokenChunkerMergeMatchesPython(t *testing.T) { text := strings.Join(mergeSourceLines, "\n") for _, cap := range []int{32, 128} { want := expectedPythonChunks(t, cap) - got := invokeTokenChunker(t, text, cap) + got := invokeTokenChunker(t, text, cap, 0) if len(got) != len(want) { t.Fatalf("chunk_token_size=%d: chunk count go=%d, python=%d", cap, len(got), len(want)) } @@ -283,3 +284,126 @@ func TestTokenChunkerMergeMatchesPython(t *testing.T) { } } } + +// pythonMergeWithOverlap mirrors rag/flow/chunker/token_chunker.py: +// _merge_text_chunks_by_token_size under the UNIFIED contract — scaled overlap +// threshold + unconditional char-based overlap prefix (no fit-check), with the +// #17799 over-budget unit standing alone. It is the Python-faithful oracle for +// TestTokenChunkerOverlapMatchesPython. +// +// Like the fixed Python JSON path (and Go's computeOverlapPrefix) it strips +// position tags before measuring the overlap cut, so the prefix is cut on the +// tag-free visible text. The 29-line parity payload carries no tags, so the +// strip is a no-op there; tag-bearing inputs are covered by the Python test +// rag/flow/tests/test_token_chunker_tag_overlap.py and the unit-level oracle in +// token_merge_units_test.go. +func pythonMergeWithOverlap(paragraphs []string, chunkTokenSize int, overlappedPercent float64) []string { + threshold := float64(chunkTokenSize) * (100.0 - overlappedPercent) / 100.0 + type ck struct { + text string + tk int + } + merged := []ck{} + prev := -1 + for _, p := range paragraphs { + // Mirror Go's text-path unitization: each paragraph is built as + // "\n" + p (rag/nlp naive_merge:1405); the leading newline is part of + // the token count so the running sum matches the Go component. + unit := "\n" + p + pt := tokenizeStr(unit) + cur := unit + curTk := pt + // #17799: an over-budget unit stands alone. + startNew := prev < 0 || (prev >= 0 && float64(merged[prev].tk) > threshold) + if pt > chunkTokenSize { + startNew = true + } + if startNew { + if prev >= 0 && overlappedPercent > 0 && merged[prev].text != "" { + vis := []rune(removeTag(merged[prev].text)) + cut := int(float64(len(vis)) * (100.0 - overlappedPercent) / 100.0) + if cut < 0 { + cut = 0 + } + if cut < len(vis) { + cur = string(vis[cut:]) + cur + curTk = tokenizeStr(cur) + } + } + merged = append(merged, ck{cur, curTk}) + prev = len(merged) - 1 + continue + } + merged[prev].text += cur + merged[prev].tk += curTk + } + out := make([]string, len(merged)) + for i := range merged { + // Mirror production token.go:465 (removeTag first, then TrimSpace) so + // whitespace between visible text and a trailing position tag survives + // exactly as the Go TokenChunker emits it. + out[i] = removeTag(strings.TrimSpace(merged[i].text)) + } + return out +} + +// TestTokenChunkerOverlapMatchesPython is the overlap=20 counterpart of +// TestTokenChunkerMergeMatchesPython (#17948). It asserts the Go TokenChunker +// output equals the Python token_chunker merge contract (unconditional overlap +// prefix, scaled threshold) for the same 29-line payload. +// +// Like #17948 it compares against a Go-port oracle of the Python algorithm +// (no live Python shell-out), so it stays deterministic and does not depend on +// the offline tiktoken cache — the token counts use the real Go tokenizer, +// identical on both sides of the comparison. This guards the NEW behaviour +// introduced by the unified algorithm (decision #4: unconditional overlap +// prefix) against Go/Python drift at overlap>0, which #17948 (overlap=0) does +// not exercise. +// TestTokenChunkerOverlapStripsTagsTrimOrder guards the final-strip ORDER for +// tag-bearing text: production (token.go:465) does removeTag(TrimSpace(text)), +// so whitespace that sits BETWEEN visible text and a trailing position tag is +// PRESERVED (TrimSpace cannot see past the tag). The oracle +// pythonMergeWithOverlap must use the SAME order, otherwise the parity test +// would assert a text that the production code never emits. +// +// Input: every paragraph ends with "@@page...##", so each merged +// chunk's tail has two spaces before its closing tag. With the production +// order the trailing spaces survive; with the swapped order they are trimmed. +func TestTokenChunkerOverlapStripsTagsTrimOrder(t *testing.T) { + tagged := []string{ + "alpha one @@1\t2\t3\t4##", + "beta two @@5\t6\t7\t8##", + "gamma three @@9\t1\t2\t3##", + } + text := strings.Join(tagged, "\n") + const overlap = 20.0 + for _, cap := range []int{8, 32} { + want := pythonMergeWithOverlap(tagged, cap, overlap) + got := invokeTokenChunker(t, text, cap, overlap) + if len(got) != len(want) { + t.Fatalf("chunk_token_size=%d overlap=%v: chunk count go=%d, python=%d", cap, overlap, len(got), len(want)) + } + for i := range got { + if got[i] != want[i] { + t.Errorf("chunk_token_size=%d overlap=%v chunk[%d] mismatch:\n got=%q\nwant=%q", cap, overlap, i, got[i], want[i]) + } + } + } +} + +func TestTokenChunkerOverlapMatchesPython(t *testing.T) { + text := strings.Join(mergeSourceLines, "\n") + const overlap = 20.0 + for _, cap := range []int{32, 64, 128} { + want := pythonMergeWithOverlap(mergeSourceLines, cap, overlap) + got := invokeTokenChunker(t, text, cap, overlap) + if len(got) != len(want) { + t.Fatalf("chunk_token_size=%d overlap=%v: chunk count go=%d, python=%d", cap, overlap, len(got), len(want)) + } + for i := range got { + if got[i] != want[i] { + t.Errorf("chunk_token_size=%d overlap=%v chunk[%d] mismatch:\n got=%q\nwant=%q", cap, overlap, i, got[i], want[i]) + } + } + } +} diff --git a/internal/ingestion/component/chunker/token_merge_units_test.go b/internal/ingestion/component/chunker/token_merge_units_test.go new file mode 100644 index 0000000000..f0e1877377 --- /dev/null +++ b/internal/ingestion/component/chunker/token_merge_units_test.go @@ -0,0 +1,281 @@ +// 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 ( + "strings" + "testing" + + "ragflow/internal/ingestion/component/schema" +) + +// mergeUnitsOracleRow is one row of the oracle output: the merged chunk text +// plus the token count mergeUnits must assign to it. +type mergeUnitsOracleRow struct { + text string + tk int +} + +// pythonMergeUnitsOracle is the Go port of the UNIFIED merge contract shared by +// the Go TokenChunker, Python naive_merge, and Python token_chunker's JSON +// path. It is the oracle that mergeUnits (the single, unified merge core) must +// match chunk-for-chunk (text AND count). +// +// Contract mirrored here (the unified "hybrid" algorithm): +// - token counts are the RUNNING SUM of per-unit counts (never re-tokenizing +// the joined text); +// - overlap>0 uses a SCALED threshold target*(100-overlap)/100 and prepends +// the previous chunk's tail UNCONDITIONALLY (no fit-check); +// - a non-text unit passes through and resets the merge run; +// - an over-budget unit (tk > target) STANDS ALONE — never merged into the +// previous chunk (retains the #17799 contract; diverges from Python JSON's +// native merge-into-prev on purpose, so the product-wide behaviour change +// is avoided); +// - MergeUnderCap additionally forbids a projected overflow. +// +// Token counts: the initial chunk and merged running sum use the EXPLICIT +// tkNums (so merge decisions are deterministic and offline-independent), but +// an overlap-prefixed chunk recomputes its count via tokenizeStr — exactly as +// mergeUnits does — so the two stay in lock-step. +func pythonMergeUnitsOracle(texts []string, tkNums []int, kinds []string, target int, overlap float64, joinSep string, strat schema.MergeStrategy) []mergeUnitsOracleRow { + threshold := float64(target) * (100.0 - overlap) / 100.0 + merged := []mergeUnitsOracleRow{} + prev := -1 + for i := range texts { + if kinds[i] != "text" { + merged = append(merged, mergeUnitsOracleRow{texts[i], tkNums[i]}) + prev = -1 + continue + } + tk := tkNums[i] + if prev < 0 { + merged = append(merged, mergeUnitsOracleRow{texts[i], tk}) + prev = len(merged) - 1 + continue + } + // #17799: an over-budget unit stands alone — never merged into prev. + if tk > target { + text := texts[i] + cpTk := tk + if overlap > 0 && merged[prev].text != "" { + vis := []rune(removeTag(merged[prev].text)) + cut := int(float64(len(vis)) * (100.0 - overlap) / 100.0) + if cut < 0 { + cut = 0 + } + if cut < len(vis) { + text = string(vis[cut:]) + texts[i] + } + cpTk = tokenizeStr(text) + } + merged = append(merged, mergeUnitsOracleRow{text, cpTk}) + prev = len(merged) - 1 + continue + } + startNew := float64(merged[prev].tk) > threshold + if !startNew && strat == schema.MergeUnderCap && merged[prev].tk+tk > target { + startNew = true + } + if startNew { + text := texts[i] + cpTk := tk + if overlap > 0 && merged[prev].text != "" { + vis := []rune(removeTag(merged[prev].text)) + cut := int(float64(len(vis)) * (100.0 - overlap) / 100.0) + if cut < 0 { + cut = 0 + } + if cut < len(vis) { + text = string(vis[cut:]) + texts[i] + } + // Mirror mergeUnits: recompute the overlap chunk's token + // count from its (prefix+cur) text. Only when a prefix was + // actually prepended; otherwise keep the explicit count. + cpTk = tokenizeStr(text) + } + merged = append(merged, mergeUnitsOracleRow{text, cpTk}) + prev = len(merged) - 1 + continue + } + if merged[prev].text != "" && texts[i] != "" { + merged[prev].text = merged[prev].text + joinSep + texts[i] + } else { + merged[prev].text = merged[prev].text + texts[i] + } + merged[prev].tk += tk + } + return merged +} + +func TestMergeUnitsMatchesPythonOracle(t *testing.T) { + cases := []struct { + name string + texts []string + tkNums []int + kinds []string + target int + overlap float64 + joinSep string + strat schema.MergeStrategy + }{ + { + name: "json overlap0", + texts: []string{"a", "b", "c", "d"}, + tkNums: []int{5, 5, 5, 5}, + kinds: []string{"text", "text", "text", "text"}, + target: 8, overlap: 0, joinSep: "\n", strat: schema.MergeOverCap, + }, + { + name: "json overlap20 unconditional prefix", + texts: []string{"a", "b", "c", "d"}, + tkNums: []int{5, 5, 5, 5}, + kinds: []string{"text", "text", "text", "text"}, + target: 8, overlap: 20, joinSep: "\n", strat: schema.MergeOverCap, + }, + { + name: "text overlap0 joinsep empty", + texts: []string{"a", "b", "c"}, + tkNums: []int{5, 5, 5}, + kinds: []string{"text", "text", "text"}, + target: 8, overlap: 0, joinSep: "", strat: schema.MergeOverCap, + }, + { + name: "text overlap20 unconditional prefix", + texts: []string{"a", "b", "c"}, + tkNums: []int{5, 5, 5}, + kinds: []string{"text", "text", "text"}, + target: 8, overlap: 20, joinSep: "", strat: schema.MergeOverCap, + }, + { + name: "nontext breaks merge run", + texts: []string{"a", "IMG", "b"}, + tkNums: []int{5, 0, 5}, + kinds: []string{"text", "image", "text"}, + target: 8, overlap: 20, joinSep: "\n", strat: schema.MergeOverCap, + }, + { + name: "oversized stands alone (#17799, not merged into prev)", + texts: []string{"small", "verylongunitthat exceedsbudgetbyalot", "tiny"}, + tkNums: []int{3, 50, 3}, + kinds: []string{"text", "text", "text"}, + target: 8, overlap: 0, joinSep: "\n", strat: schema.MergeOverCap, + }, + { + name: "under_cap no overflow", + texts: []string{"a", "b", "c"}, + tkNums: []int{5, 5, 5}, + kinds: []string{"text", "text", "text"}, + target: 8, overlap: 0, joinSep: "\n", strat: schema.MergeUnderCap, + }, + { + name: "under_cap overlap20 still no overflow but carries overlap", + texts: []string{"a", "b", "c", "d"}, + tkNums: []int{5, 5, 5, 5}, + kinds: []string{"text", "text", "text", "text"}, + target: 8, overlap: 20, joinSep: "\n", strat: schema.MergeUnderCap, + }, + { + // Tag-bearing overlap source: the previous chunk's text carries a + // coordinate tag. computeOverlapPrefix strips the tag BEFORE + // measuring the cut, so the prefix is carved from the visible text + // only. The oracle must do the same (A2). + name: "overlap cuts on tag-stripped visible text", + texts: []string{"aa@@1\t2\t3\t4##bb", "cc"}, + tkNums: []int{5, 5}, + kinds: []string{"text", "text"}, + target: 4, overlap: 20, joinSep: "\n", strat: schema.MergeOverCap, + }, + { + // Longer tag so the RAW-text cut would land INSIDE the tag. The + // overlap prefix must still be carved from the tag-free visible + // text, so a partial "@@...##" fragment can never leak into the + // next chunk (A3 / Python test_overlap_prefix_never_leaks_partial_tag). + name: "overlap never leaks partial tag when cut lands inside tag", + texts: []string{"abcd@@100\t200\t300\t400##", "wxyz"}, + tkNums: []int{10, 4}, + kinds: []string{"text", "text"}, + target: 5, overlap: 20, joinSep: "\n", strat: schema.MergeOverCap, + }, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + units := make([]schema.ChunkDoc, len(c.texts)) + for i, tx := range c.texts { + units[i] = schema.ChunkDoc{Text: tx, TKNums: intPtr(c.tkNums[i]), CKType: c.kinds[i]} + } + want := pythonMergeUnitsOracle(c.texts, c.tkNums, c.kinds, c.target, c.overlap, c.joinSep, c.strat) + got := mergeUnits(units, c.target, c.overlap, c.strat, c.joinSep) + if len(got) != len(want) { + t.Fatalf("chunk count go=%d want=%d\n got=%v\nwant=%v", len(got), len(want), got, want) + } + for i := range got { + if got[i].Text != want[i].text { + t.Errorf("chunk[%d] text mismatch:\n got=%q\nwant=%q", i, got[i].Text, want[i].text) + } + if intValue(got[i].TKNums) != want[i].tk { + t.Errorf("chunk[%d] TKNums mismatch: got=%d want=%d (text=%q)", i, intValue(got[i].TKNums), want[i].tk, got[i].Text) + } + } + }) + } +} + +// TestMergeUnitsOverlapPrefixIsTagFree locks the cross-language overlap parity +// on TAG-BEARING input: every overlap-prefixed chunk must be carved from the +// previous chunk's tag-free visible text, so a dangling "@@...##" coordinate +// fragment can never leak into a chunk (A3). It also asserts the prefix equals +// the visible-cut tail of the previous chunk, mirroring computeOverlapPrefix. +func TestMergeUnitsOverlapPrefixIsTagFree(t *testing.T) { + cases := []struct { + name string + prev string + cur string + target int + overlap float64 + }{ + {"cut inside tag", "abcd@@100\t200\t300\t400##", "wxyz", 5, 20}, + {"tag at end", "hello world@@1\t2\t3\t4##", "next", 5, 20}, + {"tag at start", "@@1\t2\t3\t4##leading text", "next", 5, 20}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + units := []schema.ChunkDoc{ + {Text: c.prev, TKNums: intPtr(10), CKType: "text"}, + {Text: c.cur, TKNums: intPtr(4), CKType: "text"}, + } + got := mergeUnits(units, c.target, c.overlap, schema.MergeOverCap, "\n") + if len(got) < 2 { + t.Fatalf("expected >=2 chunks, got %d: %v", len(got), got) + } + if strings.Contains(got[1].Text, "@@") || strings.Contains(got[1].Text, "##") { + t.Errorf("overlap-prefixed chunk leaks coord tag: %q", got[1].Text) + } + // The prefix must equal the visible-cut tail of prev (mirror of + // computeOverlapPrefix): strip tags, cut, prepend. + vis := []rune(removeTag(c.prev)) + cut := int(float64(len(vis)) * (100.0 - c.overlap) / 100.0) + if cut < 0 { + cut = 0 + } + if cut >= len(vis) { + cut = len(vis) - 1 + } + wantPrefix := string(vis[cut:]) + if !strings.HasPrefix(got[1].Text, wantPrefix) { + t.Errorf("overlap prefix mismatch: got=%q wantPrefix=%q", got[1].Text, wantPrefix) + } + }) + } +} diff --git a/internal/ingestion/component/chunker/token_oversize_whole_test.go b/internal/ingestion/component/chunker/token_oversize_whole_test.go index a7065cfef7..c286931cb1 100644 --- a/internal/ingestion/component/chunker/token_oversize_whole_test.go +++ b/internal/ingestion/component/chunker/token_oversize_whole_test.go @@ -74,14 +74,14 @@ func TestTokenChunker_OversizeUnitKeptWhole(t *testing.T) { } } -// 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. +// TestTokenChunker_OversizeUnitStandsAloneAfterInBudgetUnit pins the #17799 +// contract under the unified algorithm: an oversize unit (single paragraph > +// chunk_token_size) STANDS ALONE as its own chunk, even when an in-budget +// sentence precedes it. The unified merge decision depends only on +// prev.tk_nums > threshold for merging, but an over-budget unit never merges +// into the previous chunk (#17799 retained on purpose to avoid a product-wide +// behaviour change). The lone oversize case is pinned by +// TestTokenChunker_OversizeUnitKeptWhole. 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 diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index 9e2a339c03..f9926009d8 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -372,24 +372,40 @@ func TestMergeDecisionOverCapVsUnderCapBoundary(t *testing.T) { t.Fatalf("bad fixture: incoming unit must fit target (incT=%d target=%d)", incT, target) } - _, overAct := mergeDecision("prev text", "incoming text", "\n", target, 0, schema.MergeOverCap, prevT, incT) - if overAct != mergeThenClose { - t.Errorf("OVER_CAP at boundary: want mergeThenClose, got %v", overAct) + units := []schema.ChunkDoc{ + {Text: "prev text", TKNums: intPtr(prevT), CKType: "text"}, + {Text: "incoming text", TKNums: intPtr(incT), CKType: "text"}, } - _, underAct := mergeDecision("prev text", "incoming text", "\n", target, 0, schema.MergeUnderCap, prevT, incT) - if underAct != startNewChunk { - t.Errorf("UNDER_CAP at boundary: want startNewChunk, got %v", underAct) + // OVER_CAP at the boundary: the overflowing incoming is merged into the + // previous chunk (merge-then-close) — a single merged chunk carrying the + // running sum prevT+incT. + over := mergeUnits(units, target, 0, schema.MergeOverCap, "\n") + if len(over) != 1 { + t.Fatalf("OVER_CAP at boundary: want 1 merged chunk, got %d", len(over)) + } + if got := intValue(over[0].TKNums); got != prevT+incT { + t.Errorf("OVER_CAP merged chunk running sum: want %d, got %d", prevT+incT, got) + } + + // UNDER_CAP at the boundary: the overflowing incoming starts a fresh chunk. + under := mergeUnits(units, target, 0, schema.MergeUnderCap, "\n") + if len(under) != 2 { + t.Fatalf("UNDER_CAP at boundary: want 2 chunks, got %d", len(under)) } // Both strategies must still refuse to merge an incoming unit that // already exceeds target — it stands alone as its own chunk. - _, overBig := mergeDecision("prev text", "incoming text", "\n", target, 0, schema.MergeOverCap, prevT, target+5) - if overBig != startNewChunk { - t.Errorf("OVER_CAP with oversized incoming: want startNewChunk, got %v", overBig) + big := []schema.ChunkDoc{ + {Text: "prev text", TKNums: intPtr(prevT), CKType: "text"}, + {Text: "incoming text", TKNums: intPtr(target + 5), CKType: "text"}, } - _, underBig := mergeDecision("prev text", "incoming text", "\n", target, 0, schema.MergeUnderCap, prevT, target+5) - if underBig != startNewChunk { - t.Errorf("UNDER_CAP with oversized incoming: want startNewChunk, got %v", underBig) + overBig := mergeUnits(big, target, 0, schema.MergeOverCap, "\n") + if len(overBig) != 2 { + t.Errorf("OVER_CAP with oversized incoming: want 2 chunks (stands alone), got %d", len(overBig)) + } + underBig := mergeUnits(big, target, 0, schema.MergeUnderCap, "\n") + if len(underBig) != 2 { + t.Errorf("UNDER_CAP with oversized incoming: want 2 chunks (stands alone), got %d", len(underBig)) } } diff --git a/rag/flow/chunker/token_chunker.py b/rag/flow/chunker/token_chunker.py index 0fce2c18c8..19696026cb 100644 --- a/rag/flow/chunker/token_chunker.py +++ b/rag/flow/chunker/token_chunker.py @@ -28,6 +28,22 @@ from rag.flow.parser.pdf_chunk_metadata import ( ) from rag.nlp import naive_merge +# _TAG_RE matches parser-emitted coordinate tags of the form +# ``@@\t\t\t\t##``. Mirrors Go's +# posTagRemove (internal/ingestion/component/chunker/group.go) so the two +# languages strip tags identically. +_TAG_RE = re.compile(r"@@[\t0-9.-]+?##") + + +def remove_tag(text): + """Strip ``@@...##`` coordinate tags from text. + + Used both when measuring the overlap prefix (so the cut lands on the + tag-free visible text, matching Go's computeOverlapPrefix) and on the + final chunk text (so coordinate tags never reach embedding/index). + """ + return _TAG_RE.sub("", text or "") + class TokenChunkerParam(ProcessParamBase): def __init__(self): @@ -228,11 +244,27 @@ def _merge_text_chunks_by_token_size(chunks, chunk_token_size, overlapped_percen current = deepcopy(chunk) should_start_new = prev_text_idx < 0 or merged[prev_text_idx]["tk_nums"] > threshold + # #17799: an over-budget unit stands alone — never merged into the + # previous chunk. This matches Python naive_merge and the Go + # TokenChunker (all three paths stand the over-budget unit alone), so + # the Python JSON path, Python text path, and Go TokenChunker share one + # contract. + if current["tk_nums"] > chunk_token_size: + should_start_new = True if should_start_new: if prev_text_idx >= 0 and overlapped_percent > 0 and merged[prev_text_idx]["text"]: - overlapped = merged[prev_text_idx]["text"] - overlap_start = int(len(overlapped) * (100 - overlapped_percent) / 100.0) - current["text"] = overlapped[overlap_start:] + current["text"] + # Mirror Go computeOverlapPrefix: measure the overlap cut on the + # tag-free *visible* text, never on the raw text that still + # carries @@...## coordinate tags. This keeps the overlap prefix + # aligned with Go and prevents a partial tag from leaking into + # the next chunk when the cut would land inside a tag. + visible = remove_tag(merged[prev_text_idx]["text"]) + overlap_start = int(len(visible) * (100 - overlapped_percent) / 100.0) + if 0 <= overlap_start < len(visible): + overlap_text = visible[overlap_start:] + else: + overlap_text = "" + current["text"] = overlap_text + current["text"] current["tk_nums"] = num_tokens_from_string(current["text"]) merged.append(current) prev_text_idx = len(merged) - 1 @@ -252,7 +284,11 @@ def _finalize_json_chunks(chunks): # Convert internal chunks into the final token chunker output format. docs = [] for chunk in chunks: - text = (chunk.get("context_above") or "") + (chunk.get("text") or "") + (chunk.get("context_below") or "") + # Strip parser coordinate tags from the final text so they never + # reach embedding/index (coordinates already live in the structured + # PDF_POSITIONS_KEY field). Mirrors Go's removeTag at the chunker + # output boundary (token.go:544). + text = remove_tag((chunk.get("context_above") or "") + (chunk.get("text") or "") + (chunk.get("context_below") or "")) if not text.strip(): continue @@ -319,7 +355,9 @@ class TokenChunker(ProcessBase): if from_upstream.output_format in ["markdown", "text", "html"]: payload = getattr(from_upstream, f"{from_upstream.output_format}_result") or "" if self._param.delimiter_mode == "one": - self.set_output("chunks", [{"text": payload}] if payload.strip() else []) + # Strip parser coordinate tags so they never reach + # embedding/index (consistent with the JSON merge path). + self.set_output("chunks", [{"text": remove_tag(payload)}] if payload.strip() else []) self.callback(1, "Done.") return if delimiter_pattern: @@ -356,7 +394,9 @@ class TokenChunker(ProcessBase): if not isinstance(text, str): text = item.get("content_with_weight") if isinstance(text, str) and text.strip(): - sections.append(text) + # Strip parser coordinate tags so they never reach + # embedding/index (consistent with the JSON merge path). + sections.append(remove_tag(text)) merged_text = "\n".join(sections) self.set_output("chunks", [{"text": merged_text}] if merged_text.strip() else []) self.callback(1, "Done.") diff --git a/rag/flow/tests/test_token_chunker.py b/rag/flow/tests/test_token_chunker.py index f12acbd081..58c0ac5dfa 100644 --- a/rag/flow/tests/test_token_chunker.py +++ b/rag/flow/tests/test_token_chunker.py @@ -443,3 +443,51 @@ def test_text_delimiter_mode_one_no_atom_split(): chunks = chunker._outputs["chunks"] texts = [c["text"] for c in chunks] assert texts == ["aaa", "bbb", "ccc"], f"chunk_token_size={chunk_token_size} atom-split a delimiter segment: {texts}" + + +def test_one_mode_json_strips_coord_tags(): + # delimiter_mode="one" collapses all JSON items into ONE chunk. The + # coordinate tags carried by each item text must be stripped, matching the + # main JSON merge path (_finalize_json_chunks / remove_tag) so tags never + # reach embedding/index. Regression: this branch previously leaked @@...##. + with _load_token_chunker_with_stubs() as token_chunker_module: + token_chunker = token_chunker_module.TokenChunker + param = token_chunker_module.TokenChunkerParam() + param.delimiter_mode = "one" + chunker = token_chunker(None, "token_chunker", param) + kwargs = { + "name": "token_chunker", + "output_format": "chunks", + "chunks": [ + {"text": "Sentence one@@1\t2\t3\t4##"}, + {"text": "Sentence two@@1\t2\t3\t4##"}, + ], + } + asyncio.run(chunker._invoke(**kwargs)) + out = chunker._outputs["chunks"] + assert len(out) == 1, out + text = out[0]["text"] + assert "@@" not in text, text + assert "##" not in text, text + assert text == "Sentence one\nSentence two", text + + +def test_one_mode_text_strips_coord_tags(): + # The text/markdown/html "one" branch emits the whole payload as one chunk + # and must also strip coordinate tags. + with _load_token_chunker_with_stubs() as token_chunker_module: + token_chunker = token_chunker_module.TokenChunker + param = token_chunker_module.TokenChunkerParam() + param.delimiter_mode = "one" + chunker = token_chunker(None, "token_chunker", param) + kwargs = { + "name": "token_chunker", + "output_format": "text", + "text": "Hello world@@1\t2\t3\t4##", + } + asyncio.run(chunker._invoke(**kwargs)) + out = chunker._outputs["chunks"] + assert len(out) == 1, out + text = out[0]["text"] + assert "@@" not in text and "##" not in text, text + assert text == "Hello world", text diff --git a/rag/flow/tests/test_token_chunker_tag_overlap.py b/rag/flow/tests/test_token_chunker_tag_overlap.py new file mode 100644 index 0000000000..8aea7c0c98 --- /dev/null +++ b/rag/flow/tests/test_token_chunker_tag_overlap.py @@ -0,0 +1,61 @@ +# 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. + +"""Tag-aware overlap-prefix behaviour for the TokenChunker JSON path. + +Mirrors Go's computeOverlapPrefix: the overlap cut is measured on the +tag-free *visible* text, never on the raw text that still carries the +``@@...##`` coordinate tags. This prevents two bugs: + * the overlap prefix diverging from Go on tag-bearing inputs (parity), and + * a partial ``@@...##`` fragment leaking into the next chunk when the cut + lands inside a tag. +""" + +from rag.flow.chunker.token_chunker import _merge_text_chunks_by_token_size + + +def test_overlap_prefix_cut_on_visible_text_not_raw(): + # prev chunk carries a coordinate tag; visible text is "ABCD". + prev_text = "ABCD@@1\t2\t3\t4##" + chunks = [ + {"ck_type": "text", "text": prev_text, "tk_nums": 10}, + {"ck_type": "text", "text": "EFGH", "tk_nums": 4}, + ] + # chunk_token_size=5 -> threshold=4; prev tk_nums=10 > 4 so the second + # chunk starts a new group and prepends a 20% overlap prefix. + merged = _merge_text_chunks_by_token_size(chunks, chunk_token_size=5, overlapped_percent=20) + + assert len(merged) == 2 + new_text = merged[1]["text"] + # Go behaviour: visible="ABCD", cut = int(4 * 0.8) = 3 -> prefix "D". + # The new chunk must be exactly "D" + "EFGH". + assert new_text == "DEFGH", new_text + + +def test_overlap_prefix_never_leaks_partial_tag(): + # A longer tag so the raw-text cut falls *inside* the tag (the buggy path). + prev_text = "ABCD@@100\t200\t300\t400##" + chunks = [ + {"ck_type": "text", "text": prev_text, "tk_nums": 10}, + {"ck_type": "text", "text": "WXYZ", "tk_nums": 4}, + ] + merged = _merge_text_chunks_by_token_size(chunks, chunk_token_size=5, overlapped_percent=20) + + assert len(merged) == 2 + new_text = merged[1]["text"] + # No '@@' may appear anywhere in the produced chunk text: the prefix is + # cut from tag-free visible text, so a dangling tag fragment cannot leak. + assert "@@" not in new_text, new_text + # And the prefix must be the visible-cut prefix ("D"), not raw-cut debris. + assert new_text == "DWXYZ", new_text diff --git a/rag/nlp/__init__.py b/rag/nlp/__init__.py index d88eae8c0e..5d84231d28 100644 --- a/rag/nlp/__init__.py +++ b/rag/nlp/__init__.py @@ -1193,14 +1193,23 @@ class MergeStrategy(Enum): OVER_CAP = "over_cap" -def _merge_paragraph_groups(paragraphs, token_size, strategy, size): +def _merge_paragraph_groups(paragraphs, token_size, strategy, size, overlapped_percent=0): """Return index groups of ``paragraphs`` per ``strategy``. ``paragraphs`` are already split on the delimiter and contain no delimiter text. No atom-split is ever performed: a paragraph larger than ``token_size`` becomes its own chunk. ``size(paragraph)`` returns the token count. + + The OVER_CAP merge decision uses the overlap-scaled threshold + ``token_size * (100 - overlapped_percent) / 100`` so the grouping reserves + room for the unconditional overlap prefix (unified JSON strategy). At + ``overlapped_percent == 0`` the threshold equals ``token_size``, so grouping + is identical to the prior ``prev_t + cur_t <= token_size`` rule — including + the one-boundary-overflow close — keeping ``merge_paragraphs``/``txt_parser`` + output unchanged. """ cap = token_size + threshold = token_size * (100 - overlapped_percent) / 100.0 n = len(paragraphs) groups = [] @@ -1232,11 +1241,13 @@ def _merge_paragraph_groups(paragraphs, token_size, strategy, size): groups.append(cur) return groups - # OVER_CAP (default): greedily accumulate adjacent paragraphs while the - # projected total stays within ``token_size``; when the next paragraph would - # exceed ``token_size``, merge it anyway (one boundary overflow allowed), - # then close the chunk. A paragraph larger than ``token_size`` always stands - # alone. Never pair into fixed-size twos. + # OVER_CAP (default): a new chunk starts when the current chunk's running + # token sum exceeds the (overlap-scaled) ``threshold``; an over-budget unit + # always stands alone (#17799). The scaled threshold reserves room for the + # unconditional overlap prefix (unified JSON strategy). At overlap=0 the + # threshold equals ``token_size``, so grouping is identical to the prior + # ``prev_t + cur_t <= token_size`` rule (incl. the one-boundary-overflow + # close). cur, cur_t = [], 0 for i in range(n): pt = size(paragraphs[i]) @@ -1249,22 +1260,18 @@ def _merge_paragraph_groups(paragraphs, token_size, strategy, size): if not cur: cur, cur_t = [i], pt continue - if cur_t + pt <= cap: - cur.append(i) - cur_t += pt - else: - # Boundary overflow allowed: merge this one in, then close the chunk - # so a chunk can exceed cap by at most ~one paragraph (not unbounded). - cur.append(i) - cur_t += pt + if cur_t > threshold: groups.append(cur) - cur, cur_t = [], 0 + cur, cur_t = [i], pt + else: + cur.append(i) + cur_t += pt if cur: groups.append(cur) return groups -def merge_paragraphs(paragraphs, token_size, strategy=MergeStrategy.OVER_CAP, size=None): +def merge_paragraphs(paragraphs, token_size, strategy=MergeStrategy.OVER_CAP, size=None, overlapped_percent=0): """Group delimiter-split ``paragraphs`` into chunks using ``strategy``. Pure function: no pos / PDF coordinate handling, no atom-split. Returns a @@ -1292,7 +1299,7 @@ def merge_paragraphs(paragraphs, token_size, strategy=MergeStrategy.OVER_CAP, si """ if size is None: size = num_tokens_from_string - groups = _merge_paragraph_groups(paragraphs, token_size, strategy, size) + groups = _merge_paragraph_groups(paragraphs, token_size, strategy, size, overlapped_percent) return [[paragraphs[i] for i in g] for g in groups] @@ -1328,9 +1335,12 @@ def _reconstruct_image_chunk(paragraphs, group): return text, image -def _apply_overlap_to_chunks(chunks, overlapped_percent, chunk_token_num): +def _apply_overlap_unconditional(chunks, overlapped_percent): """Prepend an overlap prefix from the previous chunk at each new-chunk - boundary, but only when it still fits the soft ``chunk_token_num`` target. + boundary, UNCONDITIONALLY when ``overlapped_percent > 0`` (unified JSON + strategy). The prefix is never dropped for not fitting the budget, so + context is continuous across every chunk boundary; a chunk may therefore + exceed ``chunk_token_num`` by up to the overlap amount. """ if overlapped_percent <= 0: return chunks @@ -1340,7 +1350,7 @@ def _apply_overlap_to_chunks(chunks, overlapped_percent, chunk_token_num): out.append(c) continue overlap_text, _ = _compute_overlap_prefix(out[-1], overlapped_percent) - if overlap_text and num_tokens_from_string(overlap_text) + num_tokens_from_string(c) <= chunk_token_num: + if overlap_text: out.append(overlap_text + c) else: out.append(c) @@ -1404,10 +1414,10 @@ def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。; continue paragraphs.append(("\n" + sub_sec, pos)) - groups = _merge_paragraph_groups([p[0] for p in paragraphs], chunk_token_num, strategy, num_tokens_from_string) + groups = _merge_paragraph_groups([p[0] for p in paragraphs], chunk_token_num, strategy, num_tokens_from_string, overlapped_percent) cks = [_reconstruct_text_chunk(paragraphs, g) for g in groups] logging.debug("naive_merge: %d sections -> %d chunks (delimiter=%r)", len(sections), len(cks), delimiter) - return _apply_overlap_to_chunks(cks, overlapped_percent, chunk_token_num) + return _apply_overlap_unconditional(cks, overlapped_percent) def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。;!?", overlapped_percent=0, strategy=MergeStrategy.OVER_CAP): @@ -1469,14 +1479,14 @@ def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。 continue paragraphs.append(("\n" + sub_sec, text_pos, image)) - groups = _merge_paragraph_groups([p[0] for p in paragraphs], chunk_token_num, strategy, num_tokens_from_string) + groups = _merge_paragraph_groups([p[0] for p in paragraphs], chunk_token_num, strategy, num_tokens_from_string, overlapped_percent) cks, result_images = [], [] for g in groups: text, image = _reconstruct_image_chunk(paragraphs, g) cks.append(text) result_images.append(image) logging.debug("naive_merge_with_images: %d texts -> %d chunks (delimiter=%r)", len(texts), len(cks), delimiter) - return _apply_overlap_to_chunks(cks, overlapped_percent, chunk_token_num), result_images + return _apply_overlap_unconditional(cks, overlapped_percent), result_images def docx_question_level(p, bull=-1): diff --git a/test/unit_test/rag/test_naive_merge.py b/test/unit_test/rag/test_naive_merge.py index ccd6ac6091..d85c2a8d62 100644 --- a/test/unit_test/rag/test_naive_merge.py +++ b/test/unit_test/rag/test_naive_merge.py @@ -27,6 +27,7 @@ Guards against: append instead of using a projected-total check. """ +import itertools import re import pytest @@ -149,23 +150,24 @@ def test_empty_delimiter_falls_back_to_token_size_merge(): @pytest.mark.p2 -def test_overlap_prefix_is_counted_in_token_budget(): - # With overlap, each chunk = overlap-prefix + new content. The proactive - # projected-total check rejects a section that, even after prepending the - # overlap prefix, would exceed chunk_token_num; the overlap is dropped at - # that boundary instead of letting the chunk overshoot. Pre-fix, the prefix - # tokens were not counted, so the per-chunk budget check fired late and - # chunks systematically overshot chunk_token_num (observed up to 63). +def test_overlap_prefix_is_never_dropped_at_overflow(): + # With overlap, each chunk = overlap-prefix + new content. The unified + # strategy applies the overlap UNCONDITIONALLY at every boundary: it is + # never dropped for not fitting the budget, so context stays continuous + # across boundaries even when the chunk overshoots chunk_token_num by the + # overlap amount. sentences = [" ".join(["w"] * 10) for _ in range(30)] # UNDER_CAP (strict): content chunks never overflow chunk_token_num, so the # overlap-prefix budget check is the only thing under test here. chunks = _nonempty(naive_merge(sentences, chunk_token_num=50, delimiter=DEFAULT_DELIMITER, overlapped_percent=20, strategy=MergeStrategy.UNDER_CAP)) assert len(chunks) > 1 - # Each content chunk stays within the budget. Sentences are 10 tokens, the - # budget is 50, so a 5-sentence chunk is exactly 50; a 10-token overlap - # prefix (20% of 50) would push it to 60 and is therefore dropped at the - # boundary rather than letting the chunk overshoot. - assert all(_tok(c) <= 50 for c in chunks) + # The overlap prefix is always present at every boundary: each chunk (after + # the first) starts with the tail of the previous chunk. + for prev, cur in itertools.pairwise(chunks): + cut = int(len(prev) * (100 - 20) / 100.0) + assert prev[cut:] and cur.startswith(prev[cut:]), "overlap prefix missing at boundary" + # And because the prefix is never dropped, some chunks exceed the budget. + assert any(_tok(c) > 50 for c in chunks) # --------------------------------------------------------------------------- # @@ -281,14 +283,19 @@ def test_strict_cap_no_overlap_packs_to_budget(): @pytest.mark.p2 -def test_strict_cap_with_overlap_drops_overlap_at_overflow_boundary(): +def test_strict_cap_overlap_never_dropped_at_overflow_boundary(): # UNDER_CAP chunks are exactly 20 tokens (two 10-token sentences). A 20% - # overlap prefix is 4 tokens; 20 + 4 > 20, so the prefix is dropped at the - # boundary instead of letting the chunk overshoot the strict cap. + # overlap prefix is 4 tokens; 20 + 4 > 20, so under the old fit-check the + # prefix was dropped. The unified strategy applies it UNCONDITIONALLY, so + # the chunk overshoots the strict cap by the overlap amount rather than + # losing boundary context. sentences = [" ".join(["w"] * 10) for _ in range(20)] chunks = _nonempty(naive_merge(sentences, chunk_token_num=20, delimiter=DEFAULT_DELIMITER, overlapped_percent=20, strategy=MergeStrategy.UNDER_CAP)) assert len(chunks) > 1 - assert all(_tok(c) <= 20 for c in chunks) + for prev, cur in itertools.pairwise(chunks): + cut = int(len(prev) * (100 - 20) / 100.0) + assert prev[cut:] and cur.startswith(prev[cut:]), "overlap prefix missing at boundary" + assert any(_tok(c) > 20 for c in chunks) @pytest.mark.p2 @@ -316,7 +323,7 @@ def test_strict_cap_overlap_chosen_when_it_fits(): chunks = _nonempty(naive_merge(sentences, chunk_token_num=20, delimiter=DEFAULT_DELIMITER, overlapped_percent=20, strategy=MergeStrategy.UNDER_CAP)) assert all(_tok(c) <= 20 for c in chunks) overlap_seen = False - for a, b in zip(chunks, chunks[1:]): + for a, b in itertools.pairwise(chunks): a_tokens = a.split() b_tokens = b.split() if a_tokens and b_tokens and any(t in b_tokens for t in a_tokens):