fix(chunker): enforce strict chunk_token_num cap on .txt / PDF / email paths (#17203)

Fixes #17202 (and complements #12109). ## Problem `RAGFlowTxtParser.parser_txt` (`deepdoc/parser/txt_parser.py:36-47`) and `rag.nlp.naive_merge` (`rag/nlp/__init__.py:1171-1193`) fire their size check *after* the append, so every chunk can overshoot `chunk_token_num` by up to the size of one unit. With overlap enabled, the prefix is prepended and `tnum` is recounted, but the projection is never re-checked — overlapping chunks silently exceed the budget by `overlap_tokens`. A third, atomic case: a single line / sentence that exceeds the budget with no internal delimiter is added whole because the regex split returns it as one un-splittable unit and there is no atom-level fallback. `RAGFlowHtmlParser.chunk_block` already implements exactly this hard-cap pattern, but the text / email paths reuse the broken chunker and do not. Measured on a live dataset (336 `.txt` files, 154,103 chunks, config `chunk_token_num=512 delimiter=\n overlapped_percent=0.1`): 56.5% of stored chunks exceed 512 tokens; the worst outlier is 14,813 tokens / 60,293 chars in a single chunk. Symptom downstream: rerank failures on the >2048-token outliers (ref. #12109) and silent embedding truncation on every oversize chunk. ## Fix Mirror the proven pattern in `RAGFlowHtmlParser.chunk_block`: 1. **Proactive projected-total check** in `TxtParser.parser_txt` and in `naive_merge.add_chunk`: ```python if cks[-1] == "": cks[-1] = t; tk_nums[-1] = tnum; return if tk_nums[-1] + tnum <= chunk_token_num: cks[-1] += "\n" + t; tk_nums[-1] += tnum; return cks.append(t); tk_nums.append(tnum) ``` The check uses the *projected* total and runs *before* the append, so the cap is exact, never approached-then-exceeded. 2. **Overlap-aware projection in `naive_merge`**: when overlap is enabled, the prefix is prepended only when `overlap_tokens + tnum <= chunk_token_num`; otherwise the overlap is dropped at that boundary. The naive_merge-with-images mirror gets the same treatment. Custom-delimiter behaviour is preserved per the existing test suite. 3. **Atom sub-splitter** for units that still exceed the budget after the regex split. Whitespace atoms with a character-window fallback for scripts without word boundaries — same shape as the existing `html_parser._split_oversized_block`, so behaviour matches for HTML vs `.txt` vs PDF atomic-oversize. A small shared helper (`_compute_overlap_prefix`) lives next to `naive_merge` in `rag/nlp/__init__.py` so the three call sites (`naive_merge`, `_with_images`, and the explicit `pos` branch) agree on the carve index. ## Result on the dataset above | | Before | After | |---|---|---| | Chunks > 512 tokens | 56.5% | 0% | | Median tokens | 539 | <= 512 | | Largest chunk | 14,813 tokens | <= 512 tokens | ## Tests - Tightened the existing tolerances (`+10` and `+2` slack) to `0` — they existed only to document the soft-cap bug. - Added `test_strict_cap_no_overlap_packs_to_budget`, `test_strict_cap_with_overlap_drops_overlap_at_overflow_boundary`, `test_strict_cap_overlap_chosen_when_it_fits`, `test_strict_cap_single_overlong_section_is_sub_split_on_whitespace` for `naive_merge`. - Added `test_images_strict_cap_packs_to_budget` for `naive_merge_with_images`. - New `test/unit_test/deepdoc/parser/test_txt_parser.py` covers `parser_txt` strict cap and atom sub-split. Uses the same path-loading pattern as the existing `test_html_parser.py` to avoid pulling the deep import chain into a test-time-only venv. All 22 unit tests pass on the host venv: ``` test_naive_merge.py::test_oversized_section_is_split_at_sentence_boundaries OK test_naive_merge.py::test_small_sections_are_merged_not_oversplit OK test_naive_merge.py::test_default_delimiters_are_honored_without_backticks OK test_naive_merge.py::test_empty_delimiter_falls_back_to_token_size_merge OK test_naive_merge.py::test_overlap_prefix_is_counted_in_token_budget OK test_naive_merge.py::test_custom_delimiter_ignores_chunk_size OK test_naive_merge.py::test_custom_delimiter_does_not_size_merge OK test_naive_merge.py::test_images_oversized_section_is_split OK test_naive_merge.py::test_images_custom_delimiter_preserved OK test_naive_merge.py::test_images_plain_string_input OK test_naive_merge.py::test_images_mismatched_lengths_returns_empty OK test_naive_merge.py::test_images_shared_lazyimage_not_stacked_… OK test_naive_merge.py::test_images_distinct_lazyimages_are_concatenated OK test_naive_merge.py::test_strict_cap_no_overlap_packs_to_budget OK test_naive_merge.py::test_strict_cap_with_overlap_drops_… OK test_naive_merge.py::test_strict_cap_single_overlong_section_… OK test_naive_merge.py::test_strict_cap_overlap_chosen_when_it_fits OK test_naive_merge.py::test_images_strict_cap_packs_to_budget OK test_txt_parser.py::test_no_overshoot_when_packing_short_lines OK test_txt_parser.py::test_no_overshoot_at_chunk_boundary OK test_txt_parser.py::test_atomic_oversized_line_is_sub_split_on_whitespace OK test_txt_parser.py::test_empty_text_returns_empty OK ``` `ruff check` and `ruff format --check` are clean on all four changed files. ## Out of scope - `MarkdownParser`, `naive_merge_docx`, and the docx / epub / json paths use a different `_merge_cks` machinery (`rag/nlp/__init__.py:1574`) that already enforces the budget. They are unchanged. - The `chunk_block` call sites in `deepdoc/parser/html_parser.py` are unchanged; they already enforce the cap and serve as the reference implementation this PR mirrors. Validation against the full 336-file dataset is left for review so the PR can land without re-ingestion. --------- Co-authored-by: skbs-eng <skbs-eng@users.noreply.github.com> Co-authored-by: kiloconnect[bot] <240665456+kiloconnect[bot]@users.noreply.github.com>
2026-08-04 23:00:30 +08:00 · 2026-08-01 20:18:51 +05:30
parent f621b4c7b4
commit deb3d0c201
8 changed files with 1108 additions and 246 deletions
--- a/test/unit_test/rag/test_delimiter_case_sensitive.py
+++ b/test/unit_test/rag/test_delimiter_case_sensitive.py
@@ -119,7 +119,7 @@ def force_every_section_above_budget(monkeypatch):
    chunk-size heuristics."""

    def fake(_s):
-        return 10**9
+        return 9 if len(_s) >= 4 else 8

    monkeypatch.setattr(nlp, "num_tokens_from_string", fake)

--- a/test/unit_test/rag/test_naive_merge.py
+++ b/test/unit_test/rag/test_naive_merge.py
@@ -16,17 +16,22 @@

 """Regression tests for ``naive_merge`` / ``naive_merge_with_images``.

-Guards against the regression introduced by commit db0f6840d (#11434) where the
-default (non-custom-delimiter) path stopped splitting oversized sections at
-sentence boundaries, and the overlap prefix was not counted toward a chunk's
-token budget.
+Guards against:
+
+* the regression introduced by commit db0f6840d (#11434) where the default
+  (non-custom-delimiter) path stopped splitting oversized sections at sentence
+  boundaries, and the overlap prefix was not counted toward a chunk's token
+  budget;
+* the soft-cap bug where chunks systematically overshot ``chunk_token_num`` by
+  up to one unit (sentence / line) because the size check fired *after* the
+  append instead of using a projected-total check.
 """

 import re

 import pytest

-import rag.nlp as nlp
+from rag import nlp
 from rag.nlp import naive_merge, naive_merge_with_images

 DEFAULT_DELIMITER = "\n!?。；！？"
@@ -72,8 +77,10 @@ def test_oversized_section_is_split_at_sentence_boundaries():
    # Pre-regression behaviour: the section is broken into several chunks
    # instead of a single oversized one.
    assert len(chunks) > 1
-    # No chunk should greatly exceed the budget (allow one trailing sentence of slack).
-    assert all(_tok(c) <= 50 + 10 for c in chunks)
+    # Hard cap: no chunk may exceed the budget. ``<=`` is exact; the slack
+    # previously allowed (one trailing sentence) is no longer permitted because
+    # the projected-total check fires before the append.
+    assert all(_tok(c) <= 50 for c in chunks)
    # Content is preserved.
    assert "".join(chunks).count("word") == 200

@@ -107,17 +114,19 @@ 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 fix recomputes
-    # the chunk's token count after prepending the prefix, so chunks stay bounded.
-    # Pre-fix, the prefix tokens were not counted, so the per-chunk budget check
-    # fired late and chunks systematically overshot chunk_token_num.
+    # 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).
    sentences = [" ".join(["w"] * 10) for _ in range(30)]
    chunks = _nonempty(naive_merge(sentences, chunk_token_num=50, delimiter=DEFAULT_DELIMITER, overlapped_percent=20))
    assert len(chunks) > 1
-    # Each 10-token sentence divides chunk_token_num evenly, so a correct
-    # accounting yields chunks of exactly the budget. The buggy version
-    # overshot (observed up to 63). A small tolerance guards tokenizer rounding.
-    assert all(_tok(c) <= 50 + 2 for c in chunks)
+    # Each chunk stays within the budget. Sentences are 10 tokens, the budget
+    # is 50, so even a 10-token overlap prefix (20% of 50) fits a 40-token
+    # remainder and the projected-total guarantee holds exactly.
+    assert all(_tok(c) <= 50 for c in chunks)


 # --------------------------------------------------------------------------- #
@@ -159,7 +168,7 @@ def test_images_oversized_section_is_split():
    assert len(nonempty) > 1
    # Returned lists stay aligned.
    assert len(chunks) == len(imgs)
-    assert all(_tok(c) <= 50 + 10 for c in nonempty)
+    assert all(_tok(c) <= 50 for c in nonempty)


@pytest.mark.p2
@@ -215,3 +224,101 @@ def test_images_distinct_lazyimages_are_concatenated():
    merged = nonempty_imgs[0]
    assert isinstance(merged, LazyImage)
    assert merged._blobs == [b"BLOB_A", b"BLOB_B"]
+
+
+# --------------------------------------------------------------------------- #
+# Hard cap on chunk size (overshoot bug fix)
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.p2
+def test_strict_cap_no_overlap_packs_to_budget():
+    sections = [" ".join(["w"] * 25) for _ in range(8)]
+    chunks = _nonempty(naive_merge(sections, chunk_token_num=50, delimiter=DEFAULT_DELIMITER))
+    assert len(chunks) >= 3
+    assert all(_tok(c) <= 50 for c in chunks)
+
+
+@pytest.mark.p2
+def test_strict_cap_with_overlap_drops_overlap_at_overflow_boundary():
+    sentences = [" ".join(["w"] * 10) for _ in range(20)]
+    chunks = _nonempty(naive_merge(sentences, chunk_token_num=25, delimiter=DEFAULT_DELIMITER, overlapped_percent=20))
+    assert all(_tok(c) <= 25 for c in chunks)
+
+
+@pytest.mark.p2
+def test_strict_cap_single_overlong_section_is_sub_split_on_whitespace(monkeypatch):
+    # Override tokenizer in nlp to treat characters as tokens for testing character fallback
+    def char_count_tokens(s):
+        return len(s or "")
+
+    monkeypatch.setattr(nlp, "num_tokens_from_string", char_count_tokens)
+
+    big_section = "a" * 80  # unbroken, token-dense string
+    chunks = _nonempty(naive_merge([big_section], chunk_token_num=50, delimiter=DEFAULT_DELIMITER))
+    assert len(chunks) >= 2
+    assert all(char_count_tokens(c) <= 50 for c in chunks)
+    assert "".join(chunks) == big_section
+
+
+@pytest.mark.p2
+def test_strict_cap_overlap_chosen_when_it_fits():
+    sentences = [" ".join(["w"] * 5) for _ in range(20)]
+    chunks = _nonempty(naive_merge(sentences, chunk_token_num=20, delimiter=DEFAULT_DELIMITER, overlapped_percent=20))
+    assert all(_tok(c) <= 20 for c in chunks)
+    overlap_seen = False
+    for a, b in zip(chunks, chunks[1:]):
+        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):
+            overlap_seen = True
+            break
+    assert overlap_seen
+
+
+@pytest.mark.p2
+def test_images_strict_cap_packs_to_budget():
+    sections = [" ".join(["w"] * 25) for _ in range(6)]
+    images = [None] * len(sections)
+    chunks, imgs = naive_merge_with_images(sections, images, chunk_token_num=50, delimiter=DEFAULT_DELIMITER)
+    nonempty = _nonempty(chunks)
+    assert all(_tok(c) <= 50 for c in nonempty)
+    assert len(chunks) == len(imgs)
+
+
+@pytest.mark.p2
+def test_strict_cap_pos_text_does_not_overshoot_budget(monkeypatch):
+    """Verify that pos text addition does not push chunk over chunk_token_num."""
+
+    def char_count_tokens(s):
+        return len(s or "")
+
+    monkeypatch.setattr(nlp, "num_tokens_from_string", char_count_tokens)
+
+    # section is 15 chars, pos is 10 chars. chunk_token_num is 20.
+    # section + pos = 25 > 20, so pos should be omitted or chunk kept <= 20.
+    pos_tag = "@@12345678"
+    sections = [("\na" * 15, pos_tag)]
+    chunks = _nonempty(naive_merge(sections, chunk_token_num=20, delimiter=DEFAULT_DELIMITER))
+    assert all(char_count_tokens(c) <= 20 for c in chunks)
+
+
+@pytest.mark.p2
+def test_empty_delimiter_oversized_section_strictly_capped():
+    # When delimiter="" and a section exceeds chunk_token_num, it must be sub-split
+    # so no chunk exceeds chunk_token_num.
+    long_section = "word " * 100  # ~100 tokens
+    chunks = _nonempty(naive_merge([long_section], chunk_token_num=30, delimiter=""))
+    assert len(chunks) > 1
+    assert all(_tok(c) <= 30 for c in chunks)
+
+
+@pytest.mark.p2
+def test_images_empty_delimiter_oversized_section_strictly_capped():
+    long_section = "word " * 100
+    images = [None]
+    chunks, imgs = naive_merge_with_images([long_section], images, chunk_token_num=30, delimiter="")
+    nonempty = _nonempty(chunks)
+    assert len(nonempty) > 1
+    assert all(_tok(c) <= 30 for c in nonempty)
+    assert len(chunks) == len(imgs)