From 17bafb363f4c07d35d3e5c6e2cb5d7ce7648f722 Mon Sep 17 00:00:00 2001 From: Jack Date: Wed, 5 Aug 2026 13:58:45 +0800 Subject: [PATCH] fix(chunker): align Go token merge with Python OVER_CAP and delimiter boundary (#17835) Consolidates the Go chunker work that syncs `TokenChunker` with the Python reference (`rag/nlp.naive_merge` / `rag/flow/chunker/token_chunker.py`) --- .../ingestion/component/chunker/common.go | 32 +++ .../component/chunker/custom_delim_test.go | 162 +++++++++++++++ .../chunker/delimiter_case_sensitive_test.go | 11 +- internal/ingestion/component/chunker/token.go | 196 ++++++++++++++---- .../component/chunker/token_batch1_test.go | 179 +++++++++++++--- .../component/chunker/token_overlap_test.go | 89 ++++++++ .../component/chunker/token_pdfpos_test.go | 6 +- .../chunker/token_strict_cap_test.go | 117 +++++++++-- .../ingestion/component/chunker/token_test.go | 4 +- .../ingestion/component/schema/chunker.go | 11 + 10 files changed, 708 insertions(+), 99 deletions(-) create mode 100644 internal/ingestion/component/chunker/custom_delim_test.go create mode 100644 internal/ingestion/component/chunker/token_overlap_test.go diff --git a/internal/ingestion/component/chunker/common.go b/internal/ingestion/component/chunker/common.go index 0ee111d88a..1dfeadf237 100644 --- a/internal/ingestion/component/chunker/common.go +++ b/internal/ingestion/component/chunker/common.go @@ -118,6 +118,38 @@ func splitKeepingDelim(text string, pattern *regexp.Regexp) []string { return out } +// splitDroppingDelim mirrors Python's _split_text_by_pattern +// (token_chunker.py:79-90). Unlike splitKeepingDelim, the captured delimiter +// is DISCARDED rather than glued to the preceding segment: re.split with a +// captured group keeps delimiters at odd indices, and only the even-index +// (text) parts are kept. This is the behaviour the text/markdown/html path +// must reproduce so a split chunk reads "first sentence here" without the +// trailing delimiter. +func splitDroppingDelim(text string, pattern *regexp.Regexp) []string { + if pattern == nil { + return []string{text} + } + idxs := pattern.FindAllStringIndex(text, -1) + if len(idxs) == 0 { + return []string{text} + } + var out []string + cursor := 0 + for _, idx := range idxs { + start, end := idx[0], idx[1] + if start == cursor { + cursor = end + continue + } + out = append(out, text[cursor:start]) + cursor = end + } + if cursor < len(text) { + out = append(out, text[cursor:]) + } + return out +} + // --------------------------------------------------------------------------- // chunk-doc helpers // --------------------------------------------------------------------------- diff --git a/internal/ingestion/component/chunker/custom_delim_test.go b/internal/ingestion/component/chunker/custom_delim_test.go new file mode 100644 index 0000000000..58c18d6066 --- /dev/null +++ b/internal/ingestion/component/chunker/custom_delim_test.go @@ -0,0 +1,162 @@ +package chunker + +import ( + "context" + "testing" +) + +// custom_delim_test pins the backtick-wrapped newline delimiter behaviour of +// TokenChunker against Python's rag/flow/chunker/token_chunker.py. +// +// There are two distinct, path-specific divergences that this file locks down: +// +// 1. text/markdown/html path: Python's _split_text_by_pattern (token_chunker.py:79-90) +// uses re.split with a captured group and keeps only the even-index +// (text) parts, so the delimiter is DROPPED. Go's splitKeepingDelim glues +// the delimiter to the end of the preceding segment, so a chunk reads +// "first sentence here\n" instead of "first sentence here". The fix makes +// the text path drop the delimiter like Python. +// +// 2. json path: Python's _build_json_chunks + _finalize_json_chunks never +// strip the chunk text, so the delimiter's trailing newline survives +// ("first segment line one\n"). Go's invokeJSONPayload ran every chunk +// through strings.TrimSpace, which deleted that newline. The fix stops +// trimming, so the newline is kept like Python. +// +// doc_type_kwd is intentionally asserted to remain present on every chunk. +// It is a load-bearing Go field (index column + media dispatch) and is NOT +// part of the divergence; it is classified go_intentional in known_diffs.json, +// not a bug to fix here. + +const backtickNewline = "`\n`" + +func invokeTokenChunks(t *testing.T, params, input map[string]any) []map[string]any { + t.Helper() + c, err := NewTokenChunker(params) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out, err := c.Invoke(context.Background(), nil, input) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + if msg, ok := out["_ERROR"].(string); ok && msg != "" { + t.Fatalf("Go returned _ERROR: %s", msg) + } + chunks, _ := out["chunks"].([]map[string]any) + return chunks +} + +func chunkTexts(chunks []map[string]any) []string { + out := make([]string, 0, len(chunks)) + for _, c := range chunks { + if s, ok := c["text"].(string); ok { + out = append(out, s) + } + } + return out +} + +// TestCustomDelimTextDropsDelimiter reproduces token__text_backtick. +func TestCustomDelimTextDropsDelimiter(t *testing.T) { + params := map[string]any{"chunk_token_size": float64(128), "delimiters": []string{backtickNewline}} + input := map[string]any{ + "name": "t", "output_format": "text", + "text": "first sentence here\nsecond sentence here\nthird sentence here", + } + chunks := invokeTokenChunks(t, params, input) + + want := []string{"first sentence here", "second sentence here", "third sentence here"} + if len(chunks) != len(want) { + t.Fatalf("chunk count: want %d got %d (%v)", len(want), len(chunks), chunkTexts(chunks)) + } + for i, w := range want { + got := chunks[i]["text"].(string) + if got != w { + t.Errorf("chunk[%d] text: want %q got %q", i, w, got) + } + if kd, _ := chunks[i]["doc_type_kwd"].(string); kd != "text" { + t.Errorf("chunk[%d] doc_type_kwd: Go must keep it, want %q got %q", i, "text", kd) + } + } +} + +// TestCustomDelimJSONKeepsNewline reproduces token__json_backtick. +func TestCustomDelimJSONKeepsNewline(t *testing.T) { + params := map[string]any{"chunk_token_size": float64(128), "delimiters": []string{backtickNewline}} + input := map[string]any{ + "name": "t", "output_format": "json", + "json": []map[string]any{ + {"text": "first segment line one\nfirst segment line two", "doc_type_kwd": "text"}, + {"text": "second segment line one\nsecond segment line two", "doc_type_kwd": "text"}, + }, + } + chunks := invokeTokenChunks(t, params, input) + + want := []string{ + "first segment line one\n", "first segment line two", + "second segment line one\n", "second segment line two", + } + if len(chunks) != len(want) { + t.Fatalf("chunk count: want %d got %d (%v)", len(want), len(chunks), chunkTexts(chunks)) + } + for i, w := range want { + got := chunks[i]["text"].(string) + if got != w { + t.Errorf("chunk[%d] text: want %q got %q", i, w, got) + } + if kd, _ := chunks[i]["doc_type_kwd"].(string); kd != "text" { + t.Errorf("chunk[%d] doc_type_kwd: Go must keep it, want %q got %q", i, "text", kd) + } + } +} + +// TestCustomDelimMarkdownDropsDelimiter reproduces token__markdown_backtick: +// the delimiter is dropped and no chunk text ends with a newline. +func TestCustomDelimMarkdownDropsDelimiter(t *testing.T) { + params := map[string]any{"chunk_token_size": float64(128), "delimiters": []string{backtickNewline}} + input := map[string]any{ + "name": "t", "output_format": "markdown", + "markdown": "# Title\n\nParagraph one.\n\nParagraph two.", + } + chunks := invokeTokenChunks(t, params, input) + // The upstream decode normalizes markdown block boundaries into the + // backtick-newline delimiter, so the text path must split into exactly + // three trimmed chunks with the delimiter dropped (no trailing newline). + want := []string{"# Title", "Paragraph one.", "Paragraph two."} + if len(chunks) != len(want) { + t.Fatalf("chunk count: want %d got %d (%v)", len(want), len(chunks), chunkTexts(chunks)) + } + for i, w := range want { + got := chunks[i]["text"].(string) + if got != w { + t.Errorf("chunk[%d] text: want %q got %q", i, w, got) + } + if kd, _ := chunks[i]["doc_type_kwd"].(string); kd != "text" { + t.Errorf("chunk[%d] doc_type_kwd: Go must keep it, want %q got %q", i, "text", kd) + } + } +} + +// TestCustomDelimHTMLDropsDelimiter reproduces token__html_backtick. +func TestCustomDelimHTMLDropsDelimiter(t *testing.T) { + params := map[string]any{"chunk_token_size": float64(128), "delimiters": []string{backtickNewline}} + input := map[string]any{ + "name": "t", "output_format": "html", + "html": "

one

\n

two

\n

three

", + } + chunks := invokeTokenChunks(t, params, input) + want := []string{"

one

", "

two

", "

three

"} + if len(chunks) != len(want) { + t.Fatalf("chunk count: want %d got %d (%v)", len(want), len(chunks), chunkTexts(chunks)) + } + for i, w := range want { + got := chunks[i]["text"].(string) + if got != w { + t.Errorf("chunk[%d] text: want %q got %q", i, w, got) + } + if kd, _ := chunks[i]["doc_type_kwd"].(string); kd != "text" { + t.Errorf("chunk[%d] doc_type_kwd: Go must keep it, want %q got %q", i, "text", kd) + } + } +} diff --git a/internal/ingestion/component/chunker/delimiter_case_sensitive_test.go b/internal/ingestion/component/chunker/delimiter_case_sensitive_test.go index fc097221c1..0be8a758c1 100644 --- a/internal/ingestion/component/chunker/delimiter_case_sensitive_test.go +++ b/internal/ingestion/component/chunker/delimiter_case_sensitive_test.go @@ -179,9 +179,9 @@ func TestTokenChunker_BacktickEndSplitsOnlyAtLowercase(t *testing.T) { got = append(got, text) } } - // splitKeepingDelim glues the matched delimiter onto the preceding - // segment: "the end" | " and End and END come" - want := []string{"the end", "and End and END come"} + // Python's _split_text_by_pattern drops the matched delimiter and + // .strip()s each segment: "the" | "and End and END come" + want := []string{"the", "and End and END come"} if len(got) != len(want) { t.Fatalf("chunks = %#v, want %#v", got, want) } @@ -217,8 +217,9 @@ func TestTokenChunker_BacktickASplitsOnlyAtLowercase(t *testing.T) { got = append(got, text) } } - // "B" + "a" glued → "Ba"; remainder "Ab" - want := []string{"Ba", "Ab"} + // "B" + "a" glued → "Ba"; remainder "Ab". Python drops the matched + // delimiter and .strip()s, leaving "B" | "Ab". + want := []string{"B", "Ab"} if len(got) != len(want) { t.Fatalf("chunks = %#v, want %#v", got, want) } diff --git a/internal/ingestion/component/chunker/token.go b/internal/ingestion/component/chunker/token.go index 5a730f38da..060b647693 100644 --- a/internal/ingestion/component/chunker/token.go +++ b/internal/ingestion/component/chunker/token.go @@ -108,6 +108,9 @@ func (p *tokenChunkerParam) Update(conf map[string]any) { if v, ok := schema.NumericFromAny(conf["image_context_size"]); ok { p.TokenChunkerParam.ImageContextSize = int(v) } + if v, ok := conf["under_cap"].(bool); ok { + p.TokenChunkerParam.UnderCap = v + } } func defaultsToken(p tokenChunkerParam) tokenChunkerParam { @@ -304,13 +307,18 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string return c.mergeByTokenSize(text, childrenPattern) } - parts := splitKeepingDelim(text, delimPattern) + parts := splitDroppingDelim(text, delimPattern) cleaned := make([]string, 0, len(parts)) for _, p := range parts { - if strings.TrimSpace(p) == "" { + // Python's text path keeps only the even-index (text) parts from + // _split_text_by_pattern and then .strip()s each one + // (token_chunker.py:316-338), so the delimiter is dropped and + // surrounding whitespace is trimmed. + trimmed := strings.TrimSpace(p) + if trimmed == "" { continue } - cleaned = append(cleaned, p) + cleaned = append(cleaned, trimmed) } if len(cleaned) == 0 { return emptyOutputs() @@ -326,7 +334,7 @@ func (c *TokenChunkerComponent) invokeTextPayload(_ context.Context, text string // Split-then-merge: split on delimiters, then greedily merge to // chunk_token_size with optional overlap. perItem := [][]schema.ChunkDoc{docs} - merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent, true) + merged := mergeByTokenSizeFromJSON(perItem, c.param.ChunkTokenSize, c.param.OverlappedPercent, true, !c.param.UnderCap) return chunkOutputs(flatten(merged)) } @@ -468,6 +476,71 @@ 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. +// +// allowBoundaryOverflow=true selects OVER_CAP (Python's default, canonical): +// 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). allowBoundaryOverflow=false selects UNDER_CAP (strict +// no-overflow): an overflowing joined text starts a new chunk instead. +// +// 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, allowBoundaryOverflow bool) (string, mergeAction) { + incomingTokens := tokenizeStr(incoming) + // 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 { + return joined, mergeIntoPrev + } + if allowBoundaryOverflow { + // 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 @@ -503,6 +576,7 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r // 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 { @@ -510,27 +584,25 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r tkns = append(tkns, tnum) return } - merged := cks[len(cks)-1] + segment - mergedN := tokenizeStr(merged) - if mergedN <= target { - cks[len(cks)-1] = merged - tkns[len(tkns)-1] = mergedN + // 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 } - newText := segment - newTokens := tnum - if overlapPct > 0 { - overlapText, _ := computeOverlapPrefix(cks[len(cks)-1], overlapPct) - if overlapText != "" { - candidate := overlapText + segment - if candidateTokens := tokenizeStr(candidate); candidateTokens <= target { - newText = candidate - newTokens = candidateTokens - } - } + out, act := mergeDecision(cks[len(cks)-1], segment, "", target, overlapPct, !c.param.UnderCap) + switch act { + case mergeIntoPrev, mergeThenClose: + cks[len(cks)-1] = out + tkns[len(tkns)-1] = tokenizeStr(out) + prevClosed = act == mergeThenClose + case startNewChunk: + cks = append(cks, out) + tkns = append(tkns, tokenizeStr(out)) } - cks = append(cks, newText) - tkns = append(tkns, newTokens) } addUnit := func(unit string) { @@ -560,7 +632,18 @@ func (c *TokenChunkerComponent) mergeByTokenSize(text string, childrenPattern *r parts := sentenceDelimiter.Split(sec, -1) hadPart := false for _, part := range parts { - part = strings.TrimSpace(part) + // 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.) if part == "" { continue } @@ -639,7 +722,7 @@ func (c *TokenChunkerComponent) invokeJSONPayload(ctx context.Context, items []s // chunks across JSON items into one global token budget. Flatten the // per-item structure into a single sequence first so the merge is // global; non-text chunks still break the merge via their CKType. - attached = mergeByTokenSizeFromJSON([][]schema.ChunkDoc{flatten(attached)}, c.param.ChunkTokenSize, c.param.OverlappedPercent, false) + attached = mergeByTokenSizeFromJSON([][]schema.ChunkDoc{flatten(attached)}, c.param.ChunkTokenSize, c.param.OverlappedPercent, false, !c.param.UnderCap) } flat := flatten(attached) @@ -656,7 +739,7 @@ func (c *TokenChunkerComponent) invokeJSONPayload(ctx context.Context, items []s // the merge paths may carry @@...## markers that must not leak into // indexed/embedded chunk text. Crop above reads positions, not text, // so the ordering is safe. - m.Text = removeTag(strings.TrimSpace(m.Text)) + m.Text = removeTag(m.Text) if m.Text == "" { continue } @@ -878,7 +961,12 @@ func takeFromStart(text string, tokens int) string { // hard cap (rag/nlp/__init__.py after the strict chunk_token_num fix). // Oversized text units are sub-split via splitOversizedUnit before merge; // overlap is applied only when overlap+segment still fits the budget. -func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64, subSplitOversize bool) [][]schema.ChunkDoc { +// +// allowBoundaryOverflow selects the merge strategy: true = OVER_CAP (Python's +// canonical default, a chunk may exceed the target by at most one incoming +// unit), false = UNDER_CAP (never exceed the target; a projected overflow +// starts a fresh chunk). The TokenChunker threads its UnderCap param here. +func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, overlappedPct float64, subSplitOversize bool, allowBoundaryOverflow bool) [][]schema.ChunkDoc { // overlappedPct is a [0,100] percentage. Clamp defensively because this // helper is also exercised directly by tests. if overlappedPct < 0 { @@ -895,6 +983,7 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over // 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 { @@ -903,7 +992,11 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over } 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. + // 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 } @@ -911,31 +1004,48 @@ func mergeByTokenSizeFromJSON(perItem [][]schema.ChunkDoc, chunkTokens int, over // 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 } - // Proactive projected-total merge (joined with "\n"). - joined := prev.Text + "\n" + ck.Text - joinedN := tokenizeStr(joined) - if joinedN <= chunkTokens { - prev.Text = joined - prev.TKNums = intPtr(joinedN) - prev.PDFPositions = extendRawJSONArray(prev.PDFPositions, ck.PDFPositions) - prev.Positions = extendRawJSONArray(prev.Positions, ck.Positions) + // 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) + cp.TKNums = intPtr(tokenizeStr(cp.Text)) + merged = append(merged, cp) return } - // Start a new chunk; apply overlap only when it still fits. - cp := cloneChunkDoc(ck) - if overlappedPct > 0 { - if overlapText, overlapTokens := computeOverlapPrefix(prev.Text, overlappedPct); overlapTokens > 0 && overlapTokens+tk <= chunkTokens { - cp.Text = overlapText + cp.Text - cp.TKNums = intPtr(tokenizeStr(cp.Text)) - } + // Proactive projected-total merge (joined with "\n"). + out, act := mergeDecision(prev.Text, ck.Text, "\n", chunkTokens, overlappedPct, allowBoundaryOverflow) + switch act { + case mergeIntoPrev, mergeThenClose: + prev.Text = out + prev.TKNums = intPtr(tokenizeStr(out)) + 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 + cp.TKNums = intPtr(tokenizeStr(out)) + merged = append(merged, cp) } - merged = append(merged, cp) } for _, ck := range chunks { diff --git a/internal/ingestion/component/chunker/token_batch1_test.go b/internal/ingestion/component/chunker/token_batch1_test.go index eb4d659601..3312b9acf9 100644 --- a/internal/ingestion/component/chunker/token_batch1_test.go +++ b/internal/ingestion/component/chunker/token_batch1_test.go @@ -59,44 +59,169 @@ func TestSentenceDelimiterMatchesBangAndQuestion(t *testing.T) { // join exceeds the budget — so the first unit must already sit near the // budget and the second unit must not fit alongside it. func TestMergeByTokenSizeFromJSON_OverlapStripsTags(t *testing.T) { - // Size a and b so: - // - each unit alone fits the budget (no atom-split), - // - the projected join exceeds the budget (forces a new chunk), - // - overlap+b still fits (so the overlap path is exercised). - aText := strings.Repeat("word ", 20) + "@@1\t2.3## tail" - bText := "body" - aN, bN := tokenizeStr(aText), tokenizeStr(bText) - joinedN := tokenizeStr(aText + "\n" + bText) - // Budget just below the join so a and b cannot merge, but each alone fits. - budget := joinedN - 1 + // OVER_CAP (Python's canonical default) merges an overflowing unit into + // the previous chunk and then closes it, so a chunk can exceed budget by + // at most one unit. To exercise the overlap path we use three units: + // - a carries a parser tag and fits the budget alone, + // - b makes a+b overflow, so a+b merge-then-close into chunk0, + // - c starts a fresh chunk (prevClosed) with an overlap prefix from + // chunk0, which is where the tag-stripping must hold. + aText := strings.Repeat("word ", 18) + "@@1\t2.3## tail" + 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 if budget < aN { budget = aN } - if budget < bN { - budget = bN + if budget < cN { + budget = cN } - if joinedN <= budget { - t.Fatalf("could not derive tight budget (a=%d b=%d joined=%d budget=%d)", aN, bN, joinedN, budget) + if joinedAB <= budget { + t.Fatalf("could not derive tight budget (a=%d b=%d joined=%d budget=%d)", aN, bN, joinedAB, 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, 30.0, true) + got := mergeByTokenSizeFromJSON(items, budget, 30.0, true, true) merged := got[0] if len(merged) != 2 { - t.Fatalf("want 2 merged chunks (overlap path), got %d (a=%d b=%d budget=%d)", len(merged), aN, bN, budget) + t.Fatalf("want 2 chunks (overflow-closed + overlap chunk), got %d (a=%d b=%d c=%d budget=%d)", len(merged), aN, bN, cN, budget) } - // The overlap prefix is prepended to the SECOND chunk. The original - // first chunk legitimately keeps its own parser tag; only the overlap - // region (merged[1]) must be tag-free. . + // The overlap prefix must actually be prepended to chunk 1; otherwise the + // test would pass even if prevClosed started c without any overlap. + overlap, _ := computeOverlapPrefix(merged[0].Text, 30.0) + if overlap == "" { + t.Fatal("expected a non-empty overlap prefix") + } + if !strings.HasPrefix(merged[1].Text, overlap) { + t.Errorf("chunk 1 missing overlap prefix %q: %q", overlap, merged[1].Text) + } + // The overlap prefix is prepended to the SECOND chunk. The first chunk + // legitimately keeps its own parser tag; only the overlap region + // (merged[1]) must be tag-free. if strings.Contains(merged[1].Text, "@@") || strings.Contains(merged[1].Text, "##") { t.Errorf("overlap prefix leaked parser tag into chunk 1: %q", merged[1].Text) } if n := tokenizeStr(merged[1].Text); n > budget { - t.Errorf("overlap pushed second chunk over budget: tokens=%d", n) + t.Errorf("overlap pushed second chunk over budget: tokens=%d (cap=%d)", n, budget) + } +} + +// TestMergeByTokenSizeFromJSON_NonTextBoundaryResetsPrevClosed is a TDD +// regression test for the OVER_CAP boundary-overflow flag (prevClosed) leaking +// across a non-text chunk. mergeByTokenSizeFromJSON must reset prevClosed when +// the previous merged chunk is non-text; otherwise the first text chunk after a +// non-text chunk carries the stale flag and forces the NEXT text chunk to start +// a fresh chunk even though it should merge with its predecessor. +// +// Sequence: T1, T2 (overflow-merge-close into chunk0), N (non-text), T3, T4. +// With a correct reset, T3 is the first text after N (new chunk) and T4 merges +// back into T3 -> 3 chunks total. With the bug, prevClosed survives the N +// boundary and T4 is wrongly forced into its own chunk -> 4 chunks. +func TestMergeByTokenSizeFromJSON_NonTextBoundaryResetsPrevClosed(t *testing.T) { + t1 := strings.Repeat("word ", 18) + t2 := strings.Repeat("word ", 18) + t3 := strings.Repeat("word ", 9) + t4 := strings.Repeat("word ", 9) + t1N, t2N := tokenizeStr(t1), tokenizeStr(t2) + joined12 := tokenizeStr(t1 + "\n" + t2) + // Budget just below the T1+T2 join so T1 and T2 cannot merge without + // overflowing, forcing mergeThenClose on chunk0 (prevClosed=true). + budget := joined12 - 1 + if budget < t1N { + budget = t1N + } + t34 := tokenizeStr(t3 + "\n" + t4) + if t34 > budget { + t.Fatalf("T3+T4 must fit budget to exercise the merge; t34=%d budget=%d", t34, budget) + } + if joined12 <= budget { + t.Fatalf("could not derive tight budget (t1=%d t2=%d joined=%d budget=%d)", t1N, t2N, joined12, budget) + } + + items := [][]schema.ChunkDoc{ + { + {Text: t1, DocType: "text", CKType: "text", TKNums: intPtr(t1N)}, + {Text: t2, DocType: "text", CKType: "text", TKNums: intPtr(t2N)}, + {Text: "[image]", DocType: "image", CKType: "image", TKNums: intPtr(1)}, + {Text: t3, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(t3))}, + {Text: t4, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(t4))}, + }, + } + got := mergeByTokenSizeFromJSON(items, budget, 0.0, true, true) + merged := got[0] + // Expect: chunk0 (T1+T2, overflow-closed), N (non-text), chunk1 (T3+T4 merged). + if len(merged) != 3 { + var texts []string + for _, c := range merged { + texts = append(texts, c.CKType+":"+c.Text) + } + t.Fatalf("want 3 chunks (overflow-closed + non-text + merged), got %d: %v", len(merged), texts) + } + // Lock the overflow-closed precondition: T1+T2 must have merged into + // chunk0 with prevClosed set, otherwise the later assertions could pass + // without exercising the boundary-overflow path at all. + if merged[0].Text != t1+"\n"+t2 { + t.Errorf("chunk[0] should be the overflow-closed T1+T2 chunk: got %q", merged[0].Text) + } + if merged[1].CKType != "image" { + t.Errorf("chunk[1] should be the non-text image chunk, got CKType=%q text=%q", merged[1].CKType, merged[1].Text) + } + wantMerged := t3 + "\n" + t4 + if merged[2].Text != wantMerged { + t.Errorf("chunk[2] should merge T3+T4: want %q got %q", wantMerged, merged[2].Text) + } +} + +// TestMergeByTokenSizeFromJSON_UnderCapNoOverflow exercises the UNDER_CAP +// strategy (allowBoundaryOverflow=false): a projected join that would exceed +// the target must start a fresh chunk instead of merging-then-closing. This is +// the seam that lets Go follow Python's no-overflow (UNDER_CAP) strategy. Under +// OVER_CAP the same input merges a+b and overflows chunk0; here a, b, c must +// stay as three separate chunks, each within budget. +func TestMergeByTokenSizeFromJSON_UnderCapNoOverflow(t *testing.T) { + aText := strings.Repeat("word ", 18) + bText := strings.Repeat("word ", 18) + cText := strings.Repeat("word ", 18) + 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; a alone and c alone fit. + budget := joinedAB - 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) + } + 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, 0.0, true, false) + merged := got[0] + if len(merged) != 3 { + t.Fatalf("UNDER_CAP want 3 chunks (a, b, c separate), got %d", len(merged)) + } + for i, ck := range merged { + if n := tokenizeStr(ck.Text); n > budget { + t.Errorf("UNDER_CAP chunk %d exceeds target: tokens=%d (cap=%d)", i, n, budget) + } } } @@ -130,12 +255,12 @@ func clampOverlapFixture() [][]schema.ChunkDoc { } func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) { - at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100, true) + at100 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 100, true, true) if at100 == nil || len(at100) == 0 { t.Fatalf("overlappedPct=100: nil/empty result") } - at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, true) - atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, true) + at150 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 150, true, true) + atHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 1e300, true, true) if !reflect.DeepEqual(at100, at150) { t.Errorf("overlappedPct=150 should clamp to 100; output differs from 100") } @@ -143,12 +268,12 @@ func TestMergeByTokenSizeFromJSON_ClampsOverlappedPct(t *testing.T) { t.Errorf("overlappedPct=1e300 should clamp to 100; output differs from 100") } - at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0, true) + at0 := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, 0, true, true) if at0 == nil || len(at0) == 0 { t.Fatalf("overlappedPct=0: nil/empty result") } - atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, true) - atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, true) + atNeg := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -5, true, true) + atNegHuge := mergeByTokenSizeFromJSON(clampOverlapFixture(), 128, -1e300, true, true) if !reflect.DeepEqual(at0, atNeg) { t.Errorf("overlappedPct=-5 should clamp to 0; output differs from 0") } @@ -169,7 +294,7 @@ func TestMergeByTokenSizeFromJSON_EmptyPrevKeepsChunk(t *testing.T) { {Text: "keepme", DocType: "text", CKType: "text", TKNums: intPtr(5)}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0, true) + got := mergeByTokenSizeFromJSON(items, 128, 0, true, true) merged := got[0] if len(merged) != 1 { t.Fatalf("want 1 merged chunk, got %d", len(merged)) diff --git a/internal/ingestion/component/chunker/token_overlap_test.go b/internal/ingestion/component/chunker/token_overlap_test.go new file mode 100644 index 0000000000..aa14254269 --- /dev/null +++ b/internal/ingestion/component/chunker/token_overlap_test.go @@ -0,0 +1,89 @@ +package chunker + +import ( + "context" + "strings" + "testing" +) + +// TestTokenChunker_TextOverlapPreservesInterLineSpace pins the fix for the +// token-text-overlap-split divergence (parity rule token-text-overlap-split). +// +// The function under test is TokenChunkerComponent.mergeByTokenSize +// (token.go), which the text path reaches ONLY when there is no active +// delimiter: with empty delimiters, invokeTextPayload routes to +// mergeByTokenSize (token.go:303-304); with a non-empty delimiter such as +// ["\n"] it instead routes to mergeByTokenSizeFromJSON, a different code path +// that is unaffected by this fix. This test therefore drives the component +// through Invoke with delimiters: []string{} so it actually lands on the +// patched function. +// +// mergeByTokenSize splits the payload on sentence delimiters and merges per +// token budget, carrying an overlap prefix from the previous chunk. +// Python's naive_merge builds each unit from "\n" + sub_sec where sub_sec +// retains its trailing inter-line whitespace (naive_merge:1357 — it never +// TrimSpaces a unit; the only post-processing is dropping the leading empty +// placeholder at naive_merge:1370-1375). The Go merge must do the same: if it +// TrimSpaces each split fragment, the trailing space of a line is dropped, the +// overlap prefix carved from the previous (untrimmed) chunk loses that +// character, and every following chunk head diverges from Python by one +// character. +// +// This test locks the exact chunk texts against the Python reference +// (rag/nlp.naive_merge with chunk_token_size=64, delimiters=[], +// overlapped_percent=0.1). +func TestTokenChunker_TextOverlapPreservesInterLineSpace(t *testing.T) { + const ( + alpha = "alpha " + beta = "beta " + gamma = "gamma " + ) + text := strings.Repeat(alpha, 40) + "\n" + + strings.Repeat(beta, 40) + "\n" + + strings.Repeat(gamma, 40) + + c, err := NewTokenChunker(map[string]any{ + "chunk_token_size": float64(64), + "delimiters": []string{}, // routes Invoke -> mergeByTokenSize (the patched function) + "overlapped_percent": 0.1, + }) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out, err := c.Invoke(context.Background(), nil, map[string]any{ + "name": "t", "output_format": "text", "text": text, + }) + if err != nil { + t.Fatalf("Invoke: %v", err) + } + chunks, _ := out["chunks"].([]map[string]any) + // OVER_CAP (Python's canonical default) overflow-merges the first two + // groups (alpha+beta) into chunk0 and then starts a fresh, overlap-prefixed + // chunk1 for gamma. + if len(chunks) != 2 { + t.Fatalf("expected 2 chunks, got %d", len(chunks)) + } + got := make([]string, len(chunks)) + for i, ck := range chunks { + s, _ := ck["text"].(string) + got[i] = s + } + + // The regression this test pins: the inter-line trailing space of every + // line must be preserved across a merge/overlap boundary (Python emits + // "alpha \nbeta", not "alpha\nbeta"). Under OVER_CAP the alpha→beta + // boundary is inside chunk0 and the beta→gamma boundary is inside the + // overlap-prefixed chunk1. + if !strings.Contains(got[0], "alpha \nbeta") { + t.Errorf("chunk[0] lost inter-line space before newline: %q", got[0]) + } + if strings.Contains(got[0], "alpha\nbeta") { + t.Errorf("chunk[0] has space-less boundary (bug present): %q", got[0]) + } + if !strings.Contains(got[1], "beta \ngamma") { + t.Errorf("chunk[1] lost inter-line space before newline: %q", got[1]) + } + if strings.Contains(got[1], "beta\ngamma") { + t.Errorf("chunk[1] has space-less boundary (bug present): %q", got[1]) + } +} diff --git a/internal/ingestion/component/chunker/token_pdfpos_test.go b/internal/ingestion/component/chunker/token_pdfpos_test.go index d1eb951c7d..0f6b501a46 100644 --- a/internal/ingestion/component/chunker/token_pdfpos_test.go +++ b/internal/ingestion/component/chunker/token_pdfpos_test.go @@ -38,7 +38,7 @@ func TestMergeByTokenSizeFromJSON_ExtendsPDFPositions(t *testing.T) { {Text: "beta", DocType: "text", CKType: "text", TKNums: intPtr(5), PDFPositions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0, true) + got := mergeByTokenSizeFromJSON(items, 128, 0, true, true) merged := got[0] if len(merged) != 1 { t.Fatalf("want 1 merged chunk, got %d", len(merged)) @@ -63,7 +63,7 @@ func TestMergeByTokenSizeFromJSON_ExtendsPositions(t *testing.T) { {Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0, true) + got := mergeByTokenSizeFromJSON(items, 128, 0, true, true) combined := string(got[0][0].Positions) if !strings.Contains(combined, "1,2,3") || !strings.Contains(combined, "4,5,6") { t.Errorf("merged chunk dropped/omitted `positions`: %s", combined) @@ -103,7 +103,7 @@ func TestMergeByTokenSizeFromJSON_PositionsDecodeToMatrix(t *testing.T) { {Text: "b", DocType: "text", CKType: "text", TKNums: intPtr(5), Positions: posB}, }, } - got := mergeByTokenSizeFromJSON(items, 128, 0, true) + got := mergeByTokenSizeFromJSON(items, 128, 0, true, true) m := got[0][0].ToMap() raw, ok := m["positions"] if !ok { diff --git a/internal/ingestion/component/chunker/token_strict_cap_test.go b/internal/ingestion/component/chunker/token_strict_cap_test.go index 0f99156df4..75eff6a4b5 100644 --- a/internal/ingestion/component/chunker/token_strict_cap_test.go +++ b/internal/ingestion/component/chunker/token_strict_cap_test.go @@ -108,15 +108,19 @@ func TestMergeByTokenSizeFromJSON_StrictCapNoOvershoot(t *testing.T) { Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)), }) } - got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0, true) + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 0, true, true) merged := got[0] if len(merged) < 3 { t.Fatalf("want >=3 chunks, got %d", len(merged)) } + // 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). + unit := tokenizeStr(sections[0].Text) for i, ck := range merged { n := tokenizeStr(ck.Text) - if n > budget { - t.Errorf("chunk %d exceeds budget: tokens=%d text_len=%d", i, n, len(ck.Text)) + if n > budget+unit { + t.Errorf("chunk %d exceeds budget by more than one unit: tokens=%d (cap=%d unit=%d)", i, n, budget, unit) } } } @@ -131,10 +135,17 @@ func TestMergeByTokenSizeFromJSON_OverlapDroppedAtOverflow(t *testing.T) { Text: text, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(text)), }) } - got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20, true) + got := mergeByTokenSizeFromJSON([][]schema.ChunkDoc{sections}, budget, 20, true, true) + unit := tokenizeStr(sections[0].Text) for i, ck := range got[0] { - if n := tokenizeStr(ck.Text); n > budget { - t.Errorf("chunk %d exceeds budget with overlap: tokens=%d", i, n) + // OVER_CAP allows one boundary overflow (prev + one unit). The JSON + // path joins units with "\n" and cl100k is not additive across joins, + // so an overflow-closed chunk can land a little above budget+unit; + // allow one extra unit of slack for the separators + cl100k delta. + // Overlap is only ever prepended when it still fits budget, so the + // only chunks that exceed budget are the overflow-closed ones. + if n := tokenizeStr(ck.Text); n > budget+2*unit { + t.Errorf("chunk %d exceeds budget+two-units with overlap: tokens=%d (cap=%d unit=%d)", i, n, budget, unit) } } } @@ -146,20 +157,23 @@ func TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit(t *testing.T) { items := [][]schema.ChunkDoc{{ {Text: long, DocType: "text", CKType: "text", TKNums: intPtr(tokenizeStr(long))}, }} - got := mergeByTokenSizeFromJSON(items, budget, 0, true) + got := mergeByTokenSizeFromJSON(items, budget, 0, true, true) if len(got[0]) < 2 { t.Fatalf("oversized unit must yield multiple chunks, got %d", len(got[0])) } // cl100k is not additive across whitespace joins: token(a)+token(b) can be // one less than token(a+b), so the running-sum flush used by both Python's // _split_oversized_unit and the aligned Go port can leave a piece exactly - // one token over the nominal budget. The invariant we defend here is that - // the oversized unit is sub-split (not collapsed into one chunk), not a - // byte-exact cap — matching the Python reference. - const slack = 1 + // one token over the nominal budget (each sub-split piece <= budget+1). + // OVER_CAP then merges at most two such pieces into one chunk before + // closing it, so the invariant we defend is that no chunk exceeds + // 2*(budget+1): the oversized unit is sub-split (not collapsed into one + // chunk) and at most one boundary overflow is allowed — matching the + // Python reference. + const slack = 2 * (budget + 1) for i, ck := range got[0] { - if n := tokenizeStr(ck.Text); n > budget+slack { - t.Errorf("chunk %d exceeds budget by more than cl100k slack: tokens=%d (cap=%d)", i, n, budget) + if n := tokenizeStr(ck.Text); n > slack { + t.Errorf("chunk %d exceeds 2*(budget+1): tokens=%d (cap=%d)", i, n, budget) } } } @@ -167,6 +181,7 @@ func TestMergeByTokenSizeFromJSON_OversizedUnitIsSubSplit(t *testing.T) { func TestMergeByTokenSize_TextPathStrictCap(t *testing.T) { // End-to-end text path: long multi-paragraph input under a tight budget. const budget = 40 + unit := tokenizeStr(strings.TrimSpace(strings.Repeat("word ", 15))) var b strings.Builder for i := 0; i < 30; i++ { b.WriteString(strings.TrimSpace(strings.Repeat("word ", 15))) @@ -187,12 +202,72 @@ func TestMergeByTokenSize_TextPathStrictCap(t *testing.T) { } for i, ck := range chunks { text, _ := ck["text"].(string) - if n := tokenizeStr(text); n > budget { - t.Errorf("chunk %d exceeds budget: tokens=%d", i, n) + // OVER_CAP allows one boundary overflow (prev + one unit). + if n := tokenizeStr(text); n > budget+unit { + t.Errorf("chunk %d exceeds budget+one-unit: tokens=%d (cap=%d unit=%d)", i, n, budget, unit) } } } +func TestMergeByTokenSize_UnderCapNoOverflow(t *testing.T) { + // UNDER_CAP (under_cap=true) must never let a chunk exceed the token + // target: a projected join that would overflow starts a fresh chunk + // instead of merge-then-close (OVER_CAP). This exercises the seam that + // lets Go follow Python's no-overflow (UNDER_CAP) strategy without + // changing the default (OVER_CAP) behavior. + const sentence = "word word word word word word word word word word word word word " // 12 words + sentenceN := tokenizeStr(sentence) + budget := sentenceN*4 + 2 // four sentences fit, five overflow. + if budget < sentenceN*2 { + budget = sentenceN * 2 + } + var b strings.Builder + for i := 0; i < 12; i++ { + b.WriteString(sentence) + b.WriteString("! ") + } + + run := func(underCap bool) []map[string]any { + comp, err := NewTokenChunker(map[string]any{ + "delimiter_mode": "token_size", + "chunk_token_size": budget, + "under_cap": underCap, + }) + if err != nil { + t.Fatalf("NewTokenChunker: %v", err) + } + out := comp.(*TokenChunkerComponent).mergeByTokenSize(b.String(), nil) + chunks, _ := out["chunks"].([]map[string]any) + return chunks + } + + respect := run(true) + if len(respect) < 2 { + t.Fatalf("UNDER_CAP: want multiple chunks, got %d", len(respect)) + } + for i, ck := range respect { + text, _ := ck["text"].(string) + if n := tokenizeStr(text); n > budget { + t.Errorf("UNDER_CAP chunk %d exceeds target: tokens=%d (cap=%d)", i, n, budget) + } + } + + // Control: OVER_CAP (default) must overflow on the same input, proving the + // toggle changes behavior rather than being a no-op. + over := run(false) + overflowed := false + for _, ck := range over { + text, _ := ck["text"].(string) + if tokenizeStr(text) > budget { + overflowed = true + break + } + } + if !overflowed { + t.Errorf("OVER_CAP control produced no overflow on input that UNDER_CAP keeps within budget; toggle may be a no-op") + } +} + func TestMergeByTokenSize_UnbrokenAtomStrictCap(t *testing.T) { // Unbroken dense string (no whitespace / sentence delim) must still // hard-cap via the character-window fallback inside splitOversizedUnit. @@ -224,8 +299,10 @@ func TestMergeByTokenSize_UnbrokenAtomStrictCap(t *testing.T) { for i, ck := range chunks { s, _ := ck["text"].(string) joined.WriteString(s) - if n := tokenizeStr(s); n > budget { - t.Errorf("chunk %d exceeds budget: tokens=%d text=%q", i, n, s) + // Sub-split pieces are <= budget+1; OVER_CAP merges at most two before + // closing, so a chunk can reach 2*(budget+1). + if n := tokenizeStr(s); n > 2*(budget+1) { + t.Errorf("chunk %d exceeds 2*(budget+1): tokens=%d text=%q", i, n, s) } } // mergeByTokenSize prefixes "\n" on sections; stripping newlines recovers @@ -237,6 +314,7 @@ func TestMergeByTokenSize_UnbrokenAtomStrictCap(t *testing.T) { func TestInvokeTextPayload_StrictCapEndToEnd(t *testing.T) { const budget = 32 + unit := tokenizeStr(strings.TrimSpace(strings.Repeat("alpha ", 12))) var b strings.Builder for i := 0; i < 20; i++ { b.WriteString(strings.TrimSpace(strings.Repeat("alpha ", 12))) @@ -266,8 +344,9 @@ func TestInvokeTextPayload_StrictCapEndToEnd(t *testing.T) { } for i, ck := range chunks { text, _ := ck["text"].(string) - if n := tokenizeStr(text); n > budget { - t.Errorf("chunk %d exceeds budget: tokens=%d", i, n) + // OVER_CAP allows one boundary overflow (prev + one unit). + if n := tokenizeStr(text); n > budget+unit { + t.Errorf("chunk %d exceeds budget+one-unit: tokens=%d (cap=%d unit=%d)", i, n, budget, unit) } } } diff --git a/internal/ingestion/component/chunker/token_test.go b/internal/ingestion/component/chunker/token_test.go index c3cdff1e25..85588aa417 100644 --- a/internal/ingestion/component/chunker/token_test.go +++ b/internal/ingestion/component/chunker/token_test.go @@ -132,10 +132,10 @@ func TestTokenChunker_DelimNeverStandaloneChunk(t *testing.T) { t.Errorf("chunk[%d] is the bare delimiter %q", i, text) } } - if got, want := chunks[0]["text"], "alpha section\n666"; got != want { + if got, want := chunks[0]["text"], "alpha section"; got != want { t.Errorf("chunk[0] text = %q, want %q", got, want) } - if got, want := chunks[1]["text"], "\nbeta section"; got != want { + if got, want := chunks[1]["text"], "beta section"; got != want { t.Errorf("chunk[1] text = %q, want %q", got, want) } } diff --git a/internal/ingestion/component/schema/chunker.go b/internal/ingestion/component/schema/chunker.go index 826fdceabd..541a8c80c7 100644 --- a/internal/ingestion/component/schema/chunker.go +++ b/internal/ingestion/component/schema/chunker.go @@ -203,6 +203,17 @@ type TokenChunkerParam struct { // ImageContextSize is the number of surrounding tokens to attach // to image chunks. 0 disables. ImageContextSize int `json:"image_context_size"` + + // UnderCap selects the merge strategy when greedily accumulating + // adjacent units (token_chunker UNDER_CAP vs OVER_CAP). + // - false (default): OVER_CAP — mirrors Python's canonical default. + // A chunk may exceed the token target by at most one incoming unit + // (a boundary overflow then closes the chunk). + // - true: UNDER_CAP — strictly never exceed the target; when the + // projected join would overflow, start a fresh chunk instead. + // This is the seam that lets Go follow Python's no-overflow strategy + // without changing the default behavior. + UnderCap bool `json:"under_cap"` } // Defaults returns the Python default TokenChunkerParam.