diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index 10fa83510a..2531ebf9eb 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -398,7 +398,10 @@ const ( // 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. +// 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: @@ -412,17 +415,40 @@ const ( // - 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. -func mergeDecision(prevText, incoming, joinSep string, target int, overlapPct float64, strategy schema.MergeStrategy) (string, mergeAction) { - incomingTokens := tokenizeStr(incoming) +// +// 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 - if tokenizeStr(joined) <= target { + // 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 { @@ -503,11 +529,13 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r tkns = append(tkns, tokenizeStr(out)) return } - out, act := mergeDecision(cks[len(cks)-1], segment, "", target, overlapPct, c.param.MergeStrategy()) + 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 - tkns[len(tkns)-1] = tokenizeStr(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) @@ -930,22 +958,33 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over 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) + out, act := mergeDecision(prev.Text, ck.Text, "\n", chunkTokens, overlappedPct, strategy, intValue(prev.TKNums), tk) switch act { case mergeIntoPrev, mergeThenClose: prev.Text = out - prev.TKNums = intPtr(tokenizeStr(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) } diff --git a/internal/ingestion/component/chunker/token_batch1_test.go b/internal/ingestion/component/chunker/token_batch1_test.go index 8df8139fa5..3dabff167f 100644 --- a/internal/ingestion/component/chunker/token_batch1_test.go +++ b/internal/ingestion/component/chunker/token_batch1_test.go @@ -70,19 +70,21 @@ func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) { bText := strings.Repeat("word ", 18) cText := strings.Repeat("word ", 6) aN, bN, cN := tokenizeStr(aText), tokenizeStr(bText), tokenizeStr(cText) - joinedAB := tokenizeStr(aText + "\n" + bText) - // Budget just below the a+b join so a and b cannot merge without - // overflowing, but a alone and c alone fit, and an overlap prefix carved - // from chunk0 fits ahead of c (overlap path is exercised on chunk 1). - budget := joinedAB - 1 + // The merge decision is now the running sum of per-unit token counts + // (aN + bN), faithful to Python's _merge_text_chunks_by_token_size, NOT the + // re-tokenized a+b join. Budget just below the a+b running sum so a and b + // cannot merge without overflowing (forcing mergeThenClose on chunk0), but + // a alone and c alone fit, and an overlap prefix carved from chunk0 fits + // ahead of c (overlap path is exercised on chunk 1). + budget := aN + bN - 1 if budget < aN { budget = aN } if budget < cN { budget = cN } - if joinedAB <= budget { - t.Fatalf("could not derive tight budget (a=%d b=%d joined=%d budget=%d)", aN, bN, joinedAB, budget) + if aN+bN <= budget { + t.Fatalf("could not derive tight budget (a=%d b=%d sum=%d budget=%d)", aN, bN, aN+bN, budget) } items := [][]schema.ChunkDoc{ { diff --git a/internal/ingestion/component/chunker/token_json_overlap_test.go b/internal/ingestion/component/chunker/token_json_overlap_test.go new file mode 100644 index 0000000000..89f4c59c25 --- /dev/null +++ b/internal/ingestion/component/chunker/token_json_overlap_test.go @@ -0,0 +1,110 @@ +// +// 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" +) + +// TestMergeByTokenSizeFromJSON_TKNumsConsistency exercises the JSON merge path +// with overlap > 0, which previously had NO coverage, and pins the TKNums +// accounting so the running-sum merge decision cannot silently regress to the +// old re-tokenized-join count. +// +// Fixture: three text units a/b/c with a and b small enough to each fit the +// budget but a+b's running sum overflows it (so a+b merge-then-close into +// chunk0), and c starts a fresh chunk (prevClosed) carrying an overlap prefix +// from chunk0. +// +// TKNums accounting (see token.go mergeByTokenSizeFromJSON): +// - on the merge path (chunk0) TKNums is the RUNNING SUM of per-unit counts +// (aN+bN), never tokenizeStr(chunk0.Text); +// - on the boundary path (chunk1, via prevClosed) TKNums is reset to +// tokenizeStr(chunk1.Text) — i.e. the two paths use different caliber. +// This mixed accounting is documented in token.go; the assertions below +// lock the current behavior so any future unification is a deliberate +// change, not a silent drift. +func TestMergeByTokenSizeFromJSON_TKNumsConsistency(t *testing.T) { + for _, overlapPct := range []float64{0, 30} { + t.Run(overlapName(overlapPct), func(t *testing.T) { + aText := strings.Repeat("word ", 18) + bText := strings.Repeat("word ", 18) + cText := strings.Repeat("word ", 6) + aN, bN, cN := tokenizeStr(aText), tokenizeStr(bText), tokenizeStr(cText) + // Budget just below the a+b running sum so a and b cannot merge + // without overflowing (forcing mergeThenClose on chunk0), while a + // alone and c alone fit. + budget := aN + bN - 1 + if budget < aN { + budget = aN + } + if budget < cN { + budget = cN + } + if aN+bN <= budget { + t.Fatalf("could not derive tight budget (a=%d b=%d sum=%d budget=%d)", aN, bN, aN+bN, budget) + } + items := [][]schema.ChunkDoc{ + { + {Text: aText, DocType: "text", CKType: "text", TKNums: intPtr(aN)}, + {Text: bText, DocType: "text", CKType: "text", TKNums: intPtr(bN)}, + {Text: cText, DocType: "text", CKType: "text", TKNums: intPtr(cN)}, + }, + } + got := mergeByTokenSizeFromJSON(items, budget, overlapPct, schema.MergeOverCap) + merged := got[0] + if len(merged) != 2 { + t.Fatalf("want 2 chunks (overflow-closed + overlap/fresh chunk), got %d (a=%d b=%d c=%d budget=%d)", len(merged), aN, bN, cN, budget) + } + + // chunk0 is the merge-then-close of a and b. + if merged[0].Text != aText+"\n"+bText { + t.Errorf("chunk0 text mismatch:\n got=%q\nwant=%q", merged[0].Text, aText+"\n"+bText) + } + // Merge path: TKNums is the running sum, not the re-tokenized text. + if got0 := intValue(merged[0].TKNums); got0 != aN+bN { + t.Errorf("chunk0 TKNums: running sum want %d, got %d (tokenizeStr(chunk0.Text)=%d)", aN+bN, got0, tokenizeStr(merged[0].Text)) + } + + // chunk1 is c (fresh, via prevClosed). With overlap>0 it carries a + // prefix carved from chunk0; verify the overlap path actually ran. + if overlapPct > 0 { + overlap, _ := computeOverlapPrefix(merged[0].Text, overlapPct) + if overlap == "" { + t.Fatal("expected a non-empty overlap prefix from chunk0") + } + if !strings.HasPrefix(merged[1].Text, overlap) { + t.Errorf("chunk1 should start with overlap prefix %q, got %q", overlap, merged[1].Text) + } + } + // Boundary path: TKNums is the re-tokenized chunk1 text count. + if got1 := intValue(merged[1].TKNums); got1 != tokenizeStr(merged[1].Text) { + t.Errorf("chunk1 TKNums: want tokenizeStr(chunk1.Text)=%d, got %d", tokenizeStr(merged[1].Text), got1) + } + }) + } +} + +func overlapName(p float64) string { + if p == 0 { + return "overlap0" + } + return "overlap30" +} diff --git a/internal/ingestion/component/chunker/token_merge_parity_test.go b/internal/ingestion/component/chunker/token_merge_parity_test.go new file mode 100644 index 0000000000..185414f154 --- /dev/null +++ b/internal/ingestion/component/chunker/token_merge_parity_test.go @@ -0,0 +1,285 @@ +// 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" + + "ragflow/internal/agent/runtime" +) + +// mergeSourceLines is a fixed 29-line plain-text payload. The exact wording is +// the RAGFlow introduction paragraph used by the golden-parity fixtures; the +// chunk boundaries below were verified against Python's naive_merge +// (rag/nlp/__init__.py:_merge_paragraph_groups, OVER_CAP) on cl100k_base. +var mergeSourceLines = []string{ + "RAGFlow is an open-source retrieval-augmented generation engine.", + "It ingests documents of many types and splits them into chunks.", + "Each chunk is embedded and stored in a vector database for retrieval.", + "The chunker honours document structure such as headings and tables.", + "Long paragraphs are divided so each chunk fits a token budget.", + "Overlap between chunks preserves context across boundaries.", + "Retrieval returns the most relevant chunks for a user question.", + "The generator then composes an answer from those chunks.", + "Evaluation measures faithfulness and answer relevance.", + "Users tune chunk size to balance recall and precision.", + "Smaller chunks improve recall but increase storage cost.", + "Larger chunks keep more context but may dilute relevance.", + "The system supports many languages and encodings.", + "Deep parsing extracts text from PDF images and scans.", + "Tables are recognised and preserved as structured chunks.", + "Images can be described and attached to neighbouring text.", + "The API exposes dataset, document, and chat endpoints.", + "A web UI lets users manage knowledge bases visually.", + "Batch ingestion processes large corpora efficiently.", + "Concurrency limits protect the embedding service.", + "Caching avoids re-embedding unchanged content.", + "Logging records ingestion progress and errors.", + "Metrics help operators spot slow components.", + "Configuration controls parsers and chunking strategies.", + "Templates encode common ingestion pipelines.", + "The engine scales horizontally behind a load balancer.", + "Security isolates tenant data by document id.", + "Auditing tracks who accessed which knowledge base.", + "Plugins extend parsing for proprietary formats.", +} + +// pythonMergeGroups is a faithful transcription of the Python reference +// chunker merge rag/nlp/__init__.py:_merge_paragraph_groups under +// MergeStrategy.OVER_CAP. It is used here as the oracle for +// TestTokenChunkerMergeMatchesPython. Because it is itself Go code, it only +// proves that the Go TokenChunker reproduces THIS transcription; its fidelity +// to the real Python algorithm is validated by hand against the source lines +// noted below. Keep those line references in sync if the Python side moves. +// +// Mapping to rag/nlp/__init__.py:_merge_paragraph_groups (OVER_CAP): +// - each paragraph is built as "\n" + sub_sec (naive_merge ~:1405), so the +// per-paragraph token count INCLUDES the leading "\n"; we mirror that with +// pt := tokenizeStr("\n" + p). +// - cur_t is the running sum of per-paragraph size("\n" + sub_sec); a +// paragraph merges when cur_t + pt <= cap. +// - an oversized paragraph (pt > cap) is emitted alone as its own chunk +// (Python: never combined with the previous chunk). +// - when cur_t + pt > cap, OVER_CAP merges the overflowing paragraph into the +// current group and then CLOSES the group (merge-then-close), so the next +// paragraph starts a fresh group. +// - empty/whitespace-only paragraphs are DROPPED before merging, mirroring +// the Python caller's `if not sub_sec` filter (rag/nlp/__init__.py:1403) +// and the Go chunker's TrimSpace skip (token.go:703). The oracle must model +// this so it stays a faithful reference for mergeByTokenSize. +func pythonMergeGroups(paragraphs []string, cap int) [][]string { + groups := [][]string{} + cur, curT := []string{}, 0 + for _, p := range paragraphs { + // Empty fragments are dropped, exactly as the Python caller and the + // Go chunker do. Skipping here keeps the oracle faithful. + if strings.TrimSpace(p) == "" { + continue + } + // Python's naive_merge builds each paragraph as "\n" + sub_sec + // (rag/nlp/__init__.py:1405), so the per-paragraph token count + // INCLUDES the leading "\n". Count it the same way so the running + // sum matches Python's size("\n" + sub_sec). + pt := tokenizeStr("\n" + p) + if pt > cap { + if len(cur) > 0 { + groups = append(groups, cur) + } + groups = append(groups, []string{p}) + cur, curT = []string{}, 0 + continue + } + if len(cur) == 0 { + cur, curT = []string{p}, pt + continue + } + if curT+pt <= cap { + cur = append(cur, p) + curT += pt + } else { + // OVER_CAP: merge the overflowing paragraph, then close. + cur = append(cur, p) + curT += pt + groups = append(groups, cur) + cur, curT = []string{}, 0 + } + } + if len(cur) > 0 { + groups = append(groups, cur) + } + return groups +} + +// TestPythonMergeGroupsOracleSanity pins pythonMergeGroups itself, so the +// oracle cannot silently drift in lock-step with the chunker under test +// (which would make TestTokenChunkerMergeMatchesPython a no-op check). The +// assertions are tokenizer-independent INVARIANTS of the naive_merge +// (OVER_CAP) algorithm rather than hard-coded token magnitudes, so the test +// is stable across tokenizer revisions: +// 1. every input paragraph appears exactly once, in order (no loss/dup/reorder); +// 2. every group is VALID: either its running sum is within cap (a normal +// group), or it is a maximal merge-then-close overflow (running sum > +// cap but dropping its last paragraph is within cap), or it is a single +// oversized paragraph that stands alone. The group's position in the slice +// (trailing or not) is irrelevant: an overflow group may be last when the +// input ends right after an overflow. +func TestPythonMergeGroupsOracleSanity(t *testing.T) { + const cap = 8 + inputs := [][]string{ + {"word word word word word word word word word word"}, // clearly oversized -> own group + {"word", "word", "word", "word", "word", "word", "word", "word", "word", "word"}, // overflow path + {"", "x", "y", "z"}, // empty + small + {"alpha", "beta", "gamma", "delta"}, // plain small + } + for ci, in := range inputs { + groups := pythonMergeGroups(in, cap) + // Python (rag/nlp:1403) and the Go chunker (token.go:703) drop + // empty/whitespace-only fragments, so the oracle does too. Compare + // against the non-empty partition of the input rather than `in`. + var nonEmpty []string + for _, p := range in { + if strings.TrimSpace(p) == "" { + continue + } + nonEmpty = append(nonEmpty, p) + } + var flat []string + for _, g := range groups { + flat = append(flat, g...) + } + if len(flat) != len(nonEmpty) { + t.Fatalf("case %d: paragraph count mismatch got=%d want=%d", ci, len(flat), len(nonEmpty)) + } + for i := range nonEmpty { + if nonEmpty[i] != flat[i] { + t.Fatalf("case %d elem %d: order/loss mismatch got=%q want=%q", ci, i, flat[i], nonEmpty[i]) + } + } + // An oversized paragraph (its own token count > cap) must stand + // alone: the oracle emits it as a single-element group. + for gi, g := range groups { + for _, p := range g { + if tokenizeStr("\n"+p) > cap { + if len(g) != 1 { + t.Errorf("case %d group %d: oversized paragraph not isolated (len=%d)", ci, gi, len(g)) + } + break + } + } + } + for gi, g := range groups { + if len(g) == 0 { + t.Fatalf("case %d group %d is empty", ci, gi) + } + sum := 0 + for _, p := range g { + sum += tokenizeStr("\n" + p) + } + prefix := 0 + for _, p := range g[:len(g)-1] { + prefix += tokenizeStr("\n" + p) + } + valid := sum <= cap || (sum > cap && prefix <= cap) + if !valid { + t.Errorf("case %d group %d invalid: sum=%d prefix=%d (cap=%d)", ci, gi, sum, prefix, cap) + } + } + } +} + +// expectedPythonChunks returns the chunk texts Python's naive_merge produces +// for mergeSourceLines at the given chunk_token_size. +func expectedPythonChunks(t *testing.T, chunkTokenSize int) []string { + t.Helper() + groups := pythonMergeGroups(mergeSourceLines, chunkTokenSize) + want := make([]string, len(groups)) + for i, g := range groups { + want[i] = strings.Join(g, "\n") + } + return want +} + +// 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 { + 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"}, + } + input := map[string]any{ + "output_format": "text", + "name": "merge-parity", + "text": text, + } + comp, err := factory("TokenChunker", param) + if err != nil { + t.Fatalf("construct TokenChunker: %v", err) + } + out, err := comp.Invoke(context.Background(), nil, input) + if err != nil { + t.Fatalf("invoke TokenChunker: %v", err) + } + if msg, ok := out["_ERROR"].(string); ok && msg != "" { + t.Fatalf("TokenChunker returned _ERROR: %s", msg) + } + raw, _ := out["chunks"].([]map[string]any) + texts := make([]string, 0, len(raw)) + for _, c := range raw { + if s, ok := c["text"].(string); ok { + texts = append(texts, s) + } + } + return texts +} + +// TestTokenChunkerMergeMatchesPython is the red test for the merge-boundary +// divergence (ragflow chunker parity): +// +// The Go merge decision used tokenizeStr(prev + "\n" + incoming) — the BPE +// token count of the JOINED string — while Python's naive_merge accumulates a +// running sum of per-paragraph token counts (cur_t + size(p)). Because BPE +// tokenization is not additive across the "\n" boundary, the two merge +// decisions disagree: Go can merge one paragraph more (or fewer) than Python, +// shifting every downstream boundary by a line and, once a budget is small +// enough, changing the chunk count outright. +// +// This test asserts the Go TokenChunker output equals Python's naive_merge +// output (verified chunk count AND per-chunk text) for two budgets: +// - chunk_token_size=32: Python emits 9 chunks (Go over-merged to 8 before +// the running-sum fix). +// - chunk_token_size=128: both emit 3 chunks, but Python closes chunk 0 one +// line earlier than Go (boundary offset before the fix). +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) + if len(got) != len(want) { + t.Fatalf("chunk_token_size=%d: chunk count go=%d, python=%d", cap, len(got), len(want)) + } + for i := range got { + if got[i] != want[i] { + t.Errorf("chunk_token_size=%d chunk[%d] mismatch:\n got=%q\nwant=%q", cap, i, got[i], want[i]) + } + } + } +} diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index 16f2c33664..eaf68acb74 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -56,11 +56,18 @@ func TestMergeByTokenSizeFromJSON_StrictCapNoOvershoot(t *testing.T) { } // OVER_CAP (Python's canonical default) permits a chunk to exceed budget // by at most one incoming unit: a chunk is closed right after the - // overflowing merge, so it can hold prev (<= budget) + one unit (<= budget). + // overflowing merge, so its running-sum token count is <= budget+unit. + // The reconstructed chunk text carries the "\n" joins between units, and + // cl100k is non-additive across joins, so the actual token count can land + // a little above budget+unit; allow one extra unit of slack (matching the + // sibling TestMergeByTokenSizeFromJSON_OverlapDroppedAtOverflow tolerance + // for the same cl100k delta). Python's naive_merge produces the same 3 + // chunks here, so this pins the faithful boundary, not the old re-tokenized + // (under-packed) one. unit := tokenizeStr(sections[0].Text) for i, ck := range merged { n := tokenizeStr(ck.Text) - if n > budget+unit { + if n > budget+2*unit { t.Errorf("chunk %d exceeds budget by more than one unit: tokens=%d (cap=%d unit=%d)", i, n, budget, unit) } } @@ -212,19 +219,25 @@ func TestMergeByTokenSize_UnderCapNoOverflow(t *testing.T) { } } - // Control: OVER_CAP (default) must overflow on the same input, proving the - // toggle changes behavior rather than being a no-op. + // Control: OVER_CAP (default) must follow its one-boundary-overflow + // contract on the same input, proving the toggle changes behavior rather + // than being a no-op. With the running-sum merge decision (faithful to + // Python's naive_merge), a chunk may hold prev (<= budget) + one unit, so + // its running-sum token count is <= budget+unit; the reconstructed text + // carries the "\n" joins and cl100k is non-additive across them, so allow + // the same one-extra-unit slack. For this particular input Python's + // OVER_CAP does not actually exceed budget, but it still permits the + // one-unit overflow that UNDER_CAP forbids, so the two strategies produce + // different chunk counts — proving the toggle is live. over := run(false) - overflowed := false - for _, ck := range over { - text, _ := ck["text"].(string) - if tokenizeStr(text) > budget { - overflowed = true - break - } + if len(over) == len(respect) { + t.Errorf("OVER_CAP produced the same chunk count as UNDER_CAP (%d); toggle may be a no-op", len(over)) } - if !overflowed { - t.Errorf("OVER_CAP control produced no overflow on input that UNDER_CAP keeps within budget; toggle may be a no-op") + for i, ck := range over { + text, _ := ck["text"].(string) + if tokenizeStr(text) > budget+sentenceN { + t.Errorf("OVER_CAP chunk %d exceeds one-unit overflow contract: tokens=%d (cap=%d unit=%d)", i, tokenizeStr(text), budget, sentenceN) + } } } @@ -339,3 +352,44 @@ func TestInvokeJSONPayload_UnderCapEndToEnd(t *testing.T) { } } } + +// TestMergeDecisionOverCapVsUnderCapBoundary pins the semantic difference +// between the two merge strategies at the exact boundary +// (prevTokens+incomingTokens > target while incomingTokens <= target), +// independent of BPE text-token-count fluctuations. OVER_CAP must +// merge-then-close (a chunk may exceed target by at most one unit); UNDER_CAP +// must start a fresh chunk (strict no-overflow). This restores the strong +// constraint that the under_cap toggle is live — the integration test in +// TestMergeByTokenSize_UnderCapNoOverflow can only assert that the two +// strategies produce a different chunk COUNT, because on repetitive text the +// re-tokenized merged chunk can fall back under budget (cl100k is +// non-additive across the join). A regression that turns OVER_CAP into a +// no-op would not be caught there, but is caught here. +func TestMergeDecisionOverCapVsUnderCapBoundary(t *testing.T) { + const prevT, incT = 10, 10 + target := prevT + incT - 1 // boundary: running sum > target, unit fits + if incT > target { + 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) + } + + _, 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) + } + + // 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) + } + _, 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) + } +}