From 9b05e5c67ecad3c528c0675506f886703a1aae6d Mon Sep 17 00:00:00 2001 From: Jack Date: Wed, 5 Aug 2026 11:50:07 +0800 Subject: [PATCH] Fix: delimiter is chunk boundary, drop token_size atom-split (OVER_CAP default) (#17808) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Fixes a regression introduced by #17203 (strict-cap atom-split) and a secondary delimiter-handling bug from #17723. **Root cause:** - #17203 added `_split_oversized_unit` / `_compute_chunk_update`, which split oversize units into ≤ token_size pieces. This collapsed `token_size=1` into 1-token chunks and set the cap at 512, mismatching the model-layer truncation boundary (embedding ~8191 / rerank 500/4096/8192/2048). Atom-split is unnecessary: oversize units stay whole and the model layer truncates. - #17723's delimiter handling dropped consecutive delimiters (`A####B` -> `A##B`), glued JSON items with `"".join`, ignored `children_delimiters`, and stripped whitespace delimiters. ## Changes - New pure helper `merge_paragraphs(paragraphs, token_size, strategy)` with a `MergeStrategy` enum (`UNDER_CAP` / `OVER_CAP`); **default `OVER_CAP`**. `UNDER_CAP` is a strict cap (never overflows `token_size`); `OVER_CAP` greedily accumulates adjacent paragraphs while the projected total stays within `token_size`, merging one boundary-overflow paragraph before closing. Oversize paragraphs stand alone. - `naive_merge` / `naive_merge_with_images` / `RAGFlowTxtParser.parser_txt` now use `merge_paragraphs`; atom-split removed. `naive_merge` / `naive_merge_with_images` always split a section on the delimiter whenever one is present (even when the section already fits `token_size`), so delimiter text never leaks into a chunk. Only the empty-delimiter (size-only) mode skips splitting. - `token_chunker`: delimiter text is dropped (not stripped); JSON flush joins buffered items with `"\n"`; `children_delimiters` and `PDF_POSITIONS_KEY` are preserved on the delimiter path. PDF positions are now attributed **per segment** — each split chunk carries only the positions of the item(s) that contributed to it — fixing a leak where page-N coordinates were attached to page-M chunks and all segments shared one preview image. - `test_txt_parser.py` rewritten to assert the new contract (not the old strict cap); `naive_merge` and delimiter-case-sensitive matrices updated. ## Contract (refs #17799) - user specified delimiter = chunk boundary; user specified delimiter text never enters a chunk. - `token_size` = soft target + merge strategy; no atom-split. - Default strategy = `OVER_CAP`; migration can switch to `UNDER_CAP` (strict cap). - `OVER_CAP` has no hard cap; the model layer truncates oversize units. `UNDER_CAP` enforces a strict cap. ## Notes - Closes the wrong-object revert in #17774 (revert #17723 would re-introduce delimiter-in-chunk and the strict cap). - Go-side alignment (`internal/ingestion/component/chunker/token.go`) is a follow-up PR. --------- Co-authored-by: CodeBuddy --- deepdoc/parser/txt_parser.py | 26 +- rag/flow/chunker/token_chunker.py | 77 +++- rag/flow/tests/test_token_chunker.py | 178 ++++++++ .../tests/test_token_chunker_delimiter.py | 15 +- rag/nlp/__init__.py | 408 +++++++++--------- .../deepdoc/parser/test_txt_parser.py | 108 +++++ .../rag/test_delimiter_case_sensitive.py | 17 +- test/unit_test/rag/test_merge_paragraphs.py | 206 +++++++++ test/unit_test/rag/test_naive_merge.py | 123 ++++-- 9 files changed, 879 insertions(+), 279 deletions(-) create mode 100644 test/unit_test/deepdoc/parser/test_txt_parser.py create mode 100644 test/unit_test/rag/test_merge_paragraphs.py diff --git a/deepdoc/parser/txt_parser.py b/deepdoc/parser/txt_parser.py index f6414bfc97..2f914ffb08 100644 --- a/deepdoc/parser/txt_parser.py +++ b/deepdoc/parser/txt_parser.py @@ -17,8 +17,8 @@ import logging import re -from common.token_utils import num_tokens_from_string from deepdoc.parser.utils import get_text +from rag.nlp import MergeStrategy, merge_paragraphs from rag.nlp.delim import ( compile_delimiter_pattern, normalize_text_newlines, @@ -35,22 +35,6 @@ class RAGFlowTxtParser: def parser_txt(cls, txt, chunk_token_num=128, delimiter="\n!?;。；！？", keep_delimiters=False): if not isinstance(txt, str): raise TypeError("txt type should be str!") - cks = [""] - tk_nums = [0] - - def add_chunk(t): - nonlocal cks, tk_nums - tnum = num_tokens_from_string(t) - if tk_nums[-1] > chunk_token_num: - cks.append(t) - tk_nums.append(tnum) - else: - if cks[-1]: - cks[-1] += "\n" + t - else: - cks[-1] += t - tk_nums[-1] += tnum - txt = normalize_text_newlines(txt) parsed_dels = parse_delimiter_field(delimiter) dels = compile_delimiter_pattern(parsed_dels) @@ -60,6 +44,7 @@ class RAGFlowTxtParser: bool(dels), ) secs = re.split(r"(%s)" % dels, txt) if dels else [txt] + paragraphs = [] for index, sec in enumerate(secs): if dels and re.match(f"^{dels}$", sec): continue @@ -67,7 +52,12 @@ class RAGFlowTxtParser: continue if keep_delimiters and index + 1 < len(secs) and re.match(f"^{dels}$", secs[index + 1]): sec += secs[index + 1] - add_chunk(sec) + paragraphs.append(sec) + # Group delimiter-split paragraphs with the OVER_CAP merge strategy: no + # atom-split, delimiter text never enters a chunk. A paragraph larger + # than chunk_token_num stands alone; the model layer truncates it. + groups = merge_paragraphs(paragraphs, chunk_token_num, MergeStrategy.OVER_CAP) + cks = ["\n".join(g) for g in groups] logging.debug("parser_txt: %d sections -> %d chunks (chunk_token_num=%d)", len(secs), len(cks), chunk_token_num) return [[c, ""] for c in cks] diff --git a/rag/flow/chunker/token_chunker.py b/rag/flow/chunker/token_chunker.py index ab9063f43c..a9a9ff7a95 100644 --- a/rag/flow/chunker/token_chunker.py +++ b/rag/flow/chunker/token_chunker.py @@ -78,6 +78,9 @@ def _compile_delimiter_pattern(delimiters): def _split_text_by_pattern(text, pattern): # Split text by the compiled delimiter pattern and discard delimiters. + # No atom-split is performed; empty segments between consecutive delimiters + # are dropped but whitespace-only segments are preserved (the delimiter is + # the boundary, not stripped away). if not pattern: return [text or ""] @@ -85,7 +88,7 @@ def _split_text_by_pattern(text, pattern): chunks = [] for i in range(0, len(split_texts), 2): chunk = split_texts[i] - if chunk.strip(): + if chunk: chunks.append(chunk) return chunks @@ -359,31 +362,79 @@ class TokenChunker(ProcessBase): text_chunks = _build_json_chunks(json_result, "") chunks = [] text_buffer = [] + text_buffer_pos = [] def flush_text_buffer(): if not text_buffer: return - combined_text = "".join(text_buffer) - split_texts = _split_text_by_pattern(combined_text, delimiter_pattern) - chunks.extend( - { - "text": text, - "doc_type_kwd": "text", - "ck_type": "text", - "tk_nums": num_tokens_from_string(text), - } - for text in split_texts - if text.strip() - ) + # Join buffered text items with "\n" so adjacent item text is not + # glued together (e.g. "hello" + "world" must not become "helloworld"). + # The delimiter is then applied to the combined text; a segment may + # span across item boundaries (the "\n" glue is not itself a + # delimiter), so each segment carries only the PDF positions of the + # buffered item(s) whose text contributed to it -- never the union of + # every item (which previously leaked page-N coordinates into + # page-M chunks and made all segments share one preview image). + parts = [] + item_ranges = [] # (start, end) of each buffered item in combined_text + offset = 0 + for text in text_buffer: + start = offset + parts.append(text) + offset += len(text) + item_ranges.append((start, offset)) + parts.append("\n") + offset += 1 + combined_text = "".join(parts[:-1]) # drop the trailing glue + + if delimiter_pattern: + raw = re.split(r"(%s)" % delimiter_pattern, combined_text, flags=re.DOTALL) + segments = [] # (text, start, end) within combined_text + pos = 0 + for i in range(0, len(raw), 2): + seg = raw[i] + seg_start = pos + seg_end = pos + len(seg) + if seg: + segments.append((seg, seg_start, seg_end)) + pos = seg_end + if i + 1 < len(raw): + pos += len(raw[i + 1]) + else: + segments = [(combined_text, 0, len(combined_text))] + + for text, seg_start, seg_end in segments: + if not text.strip(): + continue + seg_pos = [] + for (istart, iend), item_pos in zip(item_ranges, text_buffer_pos): + # A segment overlaps an item when their character ranges + # intersect; collect that item's coordinates. + if seg_start < iend and istart < seg_end: + seg_pos.extend(item_pos or []) + chunks.append( + { + "text": text, + "doc_type_kwd": "text", + "ck_type": "text", + PDF_POSITIONS_KEY: deepcopy(seg_pos), + "tk_nums": num_tokens_from_string(text), + } + ) text_buffer.clear() + text_buffer_pos.clear() for chunk in text_chunks: if chunk["ck_type"] == "text": text_buffer.append(chunk["text"]) + text_buffer_pos.append(chunk.get(PDF_POSITIONS_KEY)) else: flush_text_buffer() chunks.append(chunk) flush_text_buffer() + # Apply children_delimiters (secondary split) before finalizing. + if custom_pattern: + chunks = _split_chunk_docs_by_children(chunks, custom_pattern) _attach_context_to_media_chunks(chunks, self._param.table_context_size, self._param.image_context_size) await restore_pdf_text_previews(chunks, from_upstream, self._canvas) self.set_output("chunks", _finalize_json_chunks(chunks)) diff --git a/rag/flow/tests/test_token_chunker.py b/rag/flow/tests/test_token_chunker.py index e50dba6fc6..8866b53de2 100644 --- a/rag/flow/tests/test_token_chunker.py +++ b/rag/flow/tests/test_token_chunker.py @@ -5,6 +5,7 @@ import types from contextlib import contextmanager from pathlib import Path + @contextmanager def _load_token_chunker_with_stubs(): root = Path(__file__).resolve().parents[3] @@ -185,3 +186,180 @@ def test_token_chunker_prefers_upstream_chunks_for_json_output_format_chunks(): asyncio.run(chunker._invoke(**kwargs)) assert chunker._outputs["chunks"] == [{"text": "CHAPTER-AWARE"}] + + +def _build_json_chunker(param: dict, monkeypatch_positions=True): + """Build a TokenChunker (bypassing ComponentBase.__init__) wired for the JSON + ``delimiter_mode`` path, with heavy deps stubbed. + + Returns ``(chunker, module)`` so callers can monkeypatch the module-global + ``extract_pdf_positions`` (it is imported as a name, so rebinding the module + attribute reaches the call sites inside ``_build_json_chunks``). + """ + with _load_token_chunker_with_stubs() as token_chunker_module: + token_chunker = token_chunker_module.TokenChunker + param_obj = token_chunker_module.TokenChunkerParam() + for key, value in param.items(): + setattr(param_obj, key, value) + + chunker = token_chunker(None, "token_chunker", param_obj) + chunker._canvas = types.SimpleNamespace(_doc_id=None, _tenant_id="t") + if monkeypatch_positions: + # Echo per-item positions so we can assert PDF coordinates survive. + token_chunker_module.extract_pdf_positions = lambda item: item.get("positions", []) + + yield token_chunker_module, chunker + + +def test_json_delimiter_mode_drop_delimiter_text(): + # The delimiter is a boundary: its text must never appear inside a chunk. + for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": ["`##`"]}): + kwargs = { + "name": "token_chunker", + "output_format": "json", + "json_result": [{"text": "first part##second part##third part", "doc_type_kwd": "text"}], + } + asyncio.run(chunker._invoke(**kwargs)) + chunks = chunker._outputs["chunks"] + texts = [c["text"] for c in chunks] + assert texts == ["first part", "second part", "third part"] + assert all("##" not in t for t in texts) + + +def test_json_delimiter_mode_newline_join_not_glued(): + # Regression for #17723: JSON flush must join buffered text items with "\\n", + # never glue them. Two adjacent items "hello" + "world" must stay + # "hello\\nworld", never become "helloworld". + for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": []}): + kwargs = { + "name": "token_chunker", + "output_format": "json", + "json_result": [ + {"text": "hello", "doc_type_kwd": "text"}, + {"text": "world", "doc_type_kwd": "text"}, + ], + } + asyncio.run(chunker._invoke(**kwargs)) + chunks = chunker._outputs["chunks"] + assert len(chunks) == 1 + assert chunks[0]["text"] == "hello\nworld" + + +def test_json_delimiter_mode_children_delimiters_applied(): + # Regression for #17723: children_delimiters (secondary split) must run before + # finalizing the JSON ``delimiter_mode`` path, or they are silently ignored. + for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": [], "children_delimiters": ["|"]}): + kwargs = { + "name": "token_chunker", + "output_format": "json", + "json_result": [{"text": "alpha|beta", "doc_type_kwd": "text"}], + } + asyncio.run(chunker._invoke(**kwargs)) + chunks = chunker._outputs["chunks"] + texts = [c["text"] for c in chunks] + assert texts == ["alpha", "beta"] + + +def test_json_delimiter_mode_pdf_positions_retained(): + # PDF coordinates carried on the combined chunk must survive into the output. + for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": []}): + kwargs = { + "name": "token_chunker", + "output_format": "json", + "json_result": [ + {"text": "hello", "doc_type_kwd": "text", "positions": [[1, 0, 10, 0, 5]]}, + {"text": "world", "doc_type_kwd": "text", "positions": [[2, 0, 20, 0, 8]]}, + ], + } + asyncio.run(chunker._invoke(**kwargs)) + chunks = chunker._outputs["chunks"] + assert len(chunks) == 1 + assert chunks[0].get("pdf_positions") == [[1, 0, 10, 0, 5], [2, 0, 20, 0, 8]] + + +def test_json_delimiter_mode_pdf_positions_per_segment_not_broadcast(): + # Regression for #3 (PDF coordinate leak): when consecutive text items from + # different pages are buffered and then split by a custom delimiter, each + # output segment must carry only the PDF positions of the item(s) that + # contributed to it -- not the union of every buffered item. The old code + # broadcast ``combined_pos`` to every split chunk, so a page-1 segment also + # claimed page-2 coordinates and all segments shared one PDF preview image. + for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": ["`。`"]}): + # Mirror production's preview-cache behaviour: a chunk's preview image is + # keyed by its position set, so chunks sharing positions share one image. + async def _restore_previews(chunks, from_upstream, canvas): + preview_cache = {} + for chunk in chunks: + positions = chunk.get("pdf_positions") or [] + key = tuple(tuple(p[:5]) for p in positions) + if key in preview_cache: + chunk["img_id"] = preview_cache[key] + else: + new_id = "img-%d" % len(preview_cache) + chunk["img_id"] = new_id + preview_cache[key] = new_id + + module.restore_pdf_text_previews = _restore_previews + + kwargs = { + "name": "doc.pdf", + "output_format": "json", + "json_result": [ + {"text": "第一章。第二段", "doc_type_kwd": "text", "positions": [[1, 0, 10, 0, 5]]}, + {"text": "第三章。第四章", "doc_type_kwd": "text", "positions": [[2, 0, 20, 0, 8]]}, + ], + } + asyncio.run(chunker._invoke(**kwargs)) + chunks = chunker._outputs["chunks"] + texts = [c["text"] for c in chunks] + # Custom "。" splits the buffered text into three segments; the "\n" join + # between the two items is NOT a split point, so the middle segment spans + # both pages. + assert texts == ["第一章", "第二段\n第三章", "第四章"], texts + + positions = [c.get("pdf_positions") for c in chunks] + # Page-1-only segment must NOT carry page-2 coordinates. + assert positions[0] == [[1, 0, 10, 0, 5]], positions + # Spanning segment legitimately carries both pages. + assert positions[1] == [[1, 0, 10, 0, 5], [2, 0, 20, 0, 8]], positions + # Page-2-only segment must NOT carry page-1 coordinates. + assert positions[2] == [[2, 0, 20, 0, 8]], positions + + # Previews must not be shared: distinct position sets -> distinct images. + img_ids = [c.get("img_id") for c in chunks] + assert len(set(img_ids)) == len(img_ids), img_ids + + +def test_json_delimiter_mode_consecutive_delimiter_keeps_boundary(): + # Regression for #17723: "A####B" with pattern "##" must yield ["A", "B"], + # both boundary-adjacent segments preserved (the bug collapsed it to "A##B"). + for module, chunker in _build_json_chunker({"delimiter_mode": "delimiter", "delimiters": ["`##`"]}): + kwargs = { + "name": "token_chunker", + "output_format": "json", + "json_result": [{"text": "A####B", "doc_type_kwd": "text"}], + } + asyncio.run(chunker._invoke(**kwargs)) + chunks = chunker._outputs["chunks"] + texts = [c["text"] for c in chunks] + assert texts == ["A", "B"] + assert all("##" not in t for t in texts) + + +def test_text_delimiter_mode_token_size_zero_or_one_no_atom_split(): + # token_size=0/1 must not atom-split delimiter segments into 1-token chunks; + # the delimiter path produces delimiter-boundary chunks regardless of cap. + for module, chunker in _build_json_chunker({"delimiter_mode": "token_size", "delimiters": ["`|`"]}): + # text path: delimiter_mode is token_size but a custom delimiter is + # present, so the delimiter branch (_split_text_by_pattern) is used. + kwargs = { + "name": "token_chunker", + "output_format": "text", + "text": "aaa|bbb|ccc", + } + for chunk_token_size in (0, 1): + setattr(chunker._param, "chunk_token_size", chunk_token_size) + asyncio.run(chunker._invoke(**kwargs)) + chunks = chunker._outputs["chunks"] + texts = [c["text"] for c in chunks] + assert texts == ["aaa", "bbb", "ccc"], f"token_size={chunk_token_size} atom-split a delimiter segment: {texts}" diff --git a/rag/flow/tests/test_token_chunker_delimiter.py b/rag/flow/tests/test_token_chunker_delimiter.py index c11e2842aa..af5d78010a 100644 --- a/rag/flow/tests/test_token_chunker_delimiter.py +++ b/rag/flow/tests/test_token_chunker_delimiter.py @@ -92,10 +92,12 @@ def test_token_chunker_token_size_mode_does_not_split_sentences(): ) -def test_naive_merge_empty_delimiter_ignores_newline_break(): - """Root-cause check: naive_merge('') can cut mid-sentence; naive_merge('\\n') cannot. +def test_naive_merge_empty_delimiter_keeps_unit_whole(): + """Empty delimiter -> no split; the whole payload is one chunk (no + atom-split). A '\\n' delimiter still honours the newline boundary. - Documents why forwarding the configured delimiters fixes the TokenChunker bug. + Documents the new contract: without a delimiter there is nothing to split + on, so the unit is kept whole and the model layer truncates it. """ if num_tokens_from_string("alive tokenizer probe sentence") <= 0: import pytest @@ -108,10 +110,7 @@ def test_naive_merge_empty_delimiter_ignores_newline_break(): chunks_empty = naive_merge(payload, 128, "") chunks_nl = naive_merge(payload, 128, "\n") - split_empty = [s for s in sentences if not any(s in t for t in chunks_empty)] + # Empty delimiter no longer cuts (no atom-split): everything stays in one chunk. + assert len(chunks_empty) == 1 split_nl = [s for s in sentences if not any(s in t for t in chunks_nl)] - - # The empty-delimiter call is the one that cuts sentences; the '\\n' call - # pre-splits on newline and keeps each sentence whole. - assert split_empty, "expected naive_merge('') to cut at least one sentence mid-stream" assert not split_nl, "naive_merge('\\n') should preserve every sentence boundary" diff --git a/rag/nlp/__init__.py b/rag/nlp/__init__.py index b6edcfcd8b..d88eae8c0e 100644 --- a/rag/nlp/__init__.py +++ b/rag/nlp/__init__.py @@ -19,6 +19,7 @@ import logging import random import re from collections import Counter, defaultdict +from enum import Enum import chardet import roman_numbers as r @@ -1177,121 +1178,177 @@ def _compute_overlap_prefix(prev_text, overlapped_percent): return overlap_text, num_tokens_from_string(overlap_text) -def _split_atom_by_token_budget(atom, chunk_token_num, token_count_fn=None): - """Split a single non-whitespace string `atom` into substrings that each - have <= chunk_token_num tokens. +class MergeStrategy(Enum): + """How ``merge_paragraphs`` groups delimiter-split paragraphs into chunks. + + ``OVER_CAP`` (default) greedily accumulates adjacent paragraphs while the + projected total stays within ``token_size``; when the next paragraph would + exceed ``token_size``, it is still merged (one boundary overflow is allowed), + then the chunk is closed. ``UNDER_CAP`` only merges when the projected total + still fits the soft ``token_size`` target and never overflows. Switching + strategy is a single enum value — no logic change elsewhere. """ - if token_count_fn is None: - token_count_fn = num_tokens_from_string - if not atom: - return [] - if token_count_fn(atom) <= chunk_token_num: - return [atom] - pieces = [] - start = 0 - n = len(atom) - while start < n: - low = start + 1 - high = n - best_end = start + 1 - while low <= high: - mid = (low + high) // 2 - substring = atom[start:mid] - if token_count_fn(substring) <= chunk_token_num: - best_end = mid - low = mid + 1 + + UNDER_CAP = "under_cap" + OVER_CAP = "over_cap" + + +def _merge_paragraph_groups(paragraphs, token_size, strategy, size): + """Return index groups of ``paragraphs`` per ``strategy``. + + ``paragraphs`` are already split on the delimiter and contain no delimiter + text. No atom-split is ever performed: a paragraph larger than ``token_size`` + becomes its own chunk. ``size(paragraph)`` returns the token count. + """ + cap = token_size + n = len(paragraphs) + groups = [] + + if strategy == MergeStrategy.UNDER_CAP: + cur = [] + cur_tokens = 0 + for i in range(n): + p = paragraphs[i] + if not cur: + cur = [i] + cur_tokens = size(p) + if cur_tokens > cap: + groups.append(cur) + cur = [] + cur_tokens = 0 + continue + if cur_tokens + size(p) <= cap: + cur.append(i) + cur_tokens += size(p) else: - high = mid - 1 - pieces.append(atom[start:best_end]) - start = best_end - return pieces + groups.append(cur) + cur = [i] + cur_tokens = size(p) + if cur_tokens > cap: + groups.append(cur) + cur = [] + cur_tokens = 0 + if cur: + groups.append(cur) + return groups - -def _split_oversized_unit(text, chunk_token_num, token_count_fn=None): - """Split a single unit that exceeds ``chunk_token_num`` tokens into pieces - that each fit the budget. Whitespace is used as the primary break (mirrors - ``RAGFlowHtmlParser._split_oversized_block``); a single run of non-whitespace - longer than the budget falls back to token-budget-based character windows. - """ - if token_count_fn is None: - token_count_fn = num_tokens_from_string - if token_count_fn(text or "") <= chunk_token_num: - return [text] - pieces = [] - current = "" - current_tokens = 0 - token_cache = {} - - def atom_tokens(atom): - if atom.isspace(): - return 0 - if atom not in token_cache: - token_cache[atom] = token_count_fn(atom) - return token_cache[atom] - - # Match whitespace runs OR non-whitespace runs (i.e. individual words/tokens). - for atom in re.findall(r"\s+|\S+", text or ""): - a_tokens = atom_tokens(atom) - if a_tokens > chunk_token_num and not atom.isspace(): - # An atom longer than the budget: flush current buffer, then carve - # token-budget-based slices out of the atom itself. - if current: - pieces.append(current) - current = "" - current_tokens = 0 - for sub_piece in _split_atom_by_token_budget(atom, chunk_token_num, token_count_fn): - pieces.append(sub_piece) + # OVER_CAP (default): greedily accumulate adjacent paragraphs while the + # projected total stays within ``token_size``; when the next paragraph would + # exceed ``token_size``, merge it anyway (one boundary overflow allowed), + # then close the chunk. A paragraph larger than ``token_size`` always stands + # alone. Never pair into fixed-size twos. + cur, cur_t = [], 0 + for i in range(n): + pt = size(paragraphs[i]) + if pt > cap: + if cur: + groups.append(cur) + groups.append([i]) + cur, cur_t = [], 0 continue - if current and current_tokens + a_tokens > chunk_token_num: - pieces.append(current) - current = "" - current_tokens = 0 - current += atom - current_tokens += a_tokens - if current: - pieces.append(current) - return pieces + if not cur: + cur, cur_t = [i], pt + continue + if cur_t + pt <= cap: + cur.append(i) + cur_t += pt + else: + # Boundary overflow allowed: merge this one in, then close the chunk + # so a chunk can exceed cap by at most ~one paragraph (not unbounded). + cur.append(i) + cur_t += pt + groups.append(cur) + cur, cur_t = [], 0 + if cur: + groups.append(cur) + return groups -def _compute_chunk_update(last_ck: str, t: str, pos: str, chunk_token_num: int, overlapped_percent: float): - tnum = num_tokens_from_string(t) - if not pos or tnum < 8: - pos = "" +def merge_paragraphs(paragraphs, token_size, strategy=MergeStrategy.OVER_CAP, size=None): + """Group delimiter-split ``paragraphs`` into chunks using ``strategy``. - # First chunk ever — no previous content to overlap with. - if last_ck == "": - new_t = t + pos if t.find(pos) < 0 else t - final_t = new_t if num_tokens_from_string(new_t) <= chunk_token_num else t - return "first", final_t, num_tokens_from_string(final_t) + Pure function: no pos / PDF coordinate handling, no atom-split. Returns a + list of chunks, each a list of the original paragraph strings (order and + identity preserved). ``token_size`` is a soft target; see ``MergeStrategy``. - # Proactive merge: append only if the *projected* total still fits. - merged = last_ck + t - merged_pos = merged + pos if last_ck.find(pos) < 0 else merged - if num_tokens_from_string(merged_pos) <= chunk_token_num: - return "merge", merged_pos, num_tokens_from_string(merged_pos) - elif num_tokens_from_string(merged) <= chunk_token_num: - return "merge", merged, num_tokens_from_string(merged) + ``size`` defaults to ``num_tokens_from_string`` and is resolved at call + time (not captured at definition) so tests can monkeypatch the tokenizer + deterministically via ``rag.nlp.num_tokens_from_string``. - # Need a new chunk. Apply overlap prefix from the previous chunk — - # but only when the projected size (overlap + t) fits — otherwise drop - # the overlap for this boundary so the chunk stays within budget. - new_t = t - new_tnum = tnum - if overlapped_percent > 0: - overlap_text, overlap_tokens = _compute_overlap_prefix(last_ck, overlapped_percent) - if overlap_tokens + new_tnum <= chunk_token_num: - new_t = overlap_text + t - new_tnum = num_tokens_from_string(new_t) - if t.find(pos) < 0: - new_t_with_pos = new_t + pos - new_tnum_with_pos = num_tokens_from_string(new_t_with_pos) - if new_tnum_with_pos <= chunk_token_num: - new_t = new_t_with_pos - new_tnum = new_tnum_with_pos - return "append", new_t, new_tnum + Chunking contract (refs #17799) + -------------------------------- + * **Delimiter is a chunk boundary.** The delimiter text specified by the + user never enters a chunk. ``naive_merge`` / ``naive_merge_with_images`` + split every section on the delimiter (except the empty-delimiter + size-only mode) so boundary text cannot leak into a chunk. + * **``token_size`` is a soft target + merge strategy.** There is no + atom-split: a paragraph larger than ``token_size`` stands alone as its own + chunk and is truncated later by the model layer. + * **Default strategy is ``OVER_CAP``.** A migration that needs the old + strict behaviour can opt into ``UNDER_CAP``. + * **``OVER_CAP`` has no hard cap** (the model layer truncates oversize + units); **``UNDER_CAP`` enforces a strict cap** and never overflows + ``token_size``. + """ + if size is None: + size = num_tokens_from_string + groups = _merge_paragraph_groups(paragraphs, token_size, strategy, size) + return [[paragraphs[i] for i in g] for g in groups] -def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。；！？", overlapped_percent=0): +def _reconstruct_text_chunk(paragraphs, group): + """Rebuild a chunk string from a ``merge_paragraphs`` group, re-attaching + ``pos`` (PDF coordinate tag) per the historical caller convention: append + ``pos`` to a paragraph when it is not already present in the running text. + """ + text = "" + for idx in group: + ptext, ppos = paragraphs[idx] + new_text = text + ptext + if ppos and ptext.find(ppos) < 0 and new_text.find(ppos) < 0: + new_text += ppos + text = new_text + return text + + +def _reconstruct_image_chunk(paragraphs, group): + """Like ``_reconstruct_text_chunk`` but also concatenates the image of every + merged paragraph (mirrors the previous ``concat_img`` dedupe behaviour). + """ + text = "" + image = None + for idx in group: + ptext, ppos, pimg = paragraphs[idx] + new_text = text + ptext + if ppos and ptext.find(ppos) < 0 and new_text.find(ppos) < 0: + new_text += ppos + text = new_text + if pimg is not None: + image = pimg if image is None else concat_img(image, pimg) + return text, image + + +def _apply_overlap_to_chunks(chunks, overlapped_percent, chunk_token_num): + """Prepend an overlap prefix from the previous chunk at each new-chunk + boundary, but only when it still fits the soft ``chunk_token_num`` target. + """ + if overlapped_percent <= 0: + return chunks + out = [] + for i, c in enumerate(chunks): + if i == 0: + out.append(c) + continue + overlap_text, _ = _compute_overlap_prefix(out[-1], overlapped_percent) + if overlap_text and num_tokens_from_string(overlap_text) + num_tokens_from_string(c) <= chunk_token_num: + out.append(overlap_text + c) + else: + out.append(c) + return out + + +def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。；！？", overlapped_percent=0, strategy=MergeStrategy.OVER_CAP): + """Split sections into chunks. Chunking contract: see ``merge_paragraphs`` (refs #17799).""" if not sections: return [] if isinstance(sections, str): @@ -1300,29 +1357,17 @@ def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。； sections = [(s, "") for s in sections] # Normalize line endings so delimiter ``\n`` matches ``\r\n`` and standalone ``\r``. sections = [(normalize_text_newlines(s), pos) for s, pos in sections] - cks = [""] - tk_nums = [0] - - def add_chunk(t, pos): - nonlocal cks, tk_nums - action, text, tk_num = _compute_chunk_update(cks[-1], t, pos, chunk_token_num, overlapped_percent) - if action in ("first", "merge"): - cks[-1] = text - tk_nums[-1] = tk_num - else: - cks.append(text) - tk_nums.append(tk_num) # Parse the delimiter field once, via the canonical helper (#17383). # `has_custom` means the field contains a backtick-wrapped token — the - # historical signal that chunk_token_num should be bypassed. Splitting - # itself uses every parsed delimiter (bare and wrapped). + # historical signal that chunk_token_num should be bypassed: each segment is + # its own chunk. parsed_dels = parse_delimiter_field(delimiter) has_custom = has_wrapped_delimiter(delimiter) if has_custom: # Custom delimiters ignore chunk_token_num: each segment is its own chunk. custom_pattern = compile_delimiter_pattern(parsed_dels) - cks, tk_nums = [], [] + cks = [] for sec, pos in sections: split_sec = re.split(r"(%s)" % custom_pattern, sec, flags=re.DOTALL) if custom_pattern else [sec] for sub_sec in split_sec: @@ -1337,79 +1382,47 @@ def naive_merge(sections: str | list, chunk_token_num=128, delimiter="\n。； if local_pos and text.find(local_pos) < 0: text += local_pos cks.append(text) - tk_nums.append(num_tokens_from_string(text)) return cks - # Split oversized sections at sentence delimiters; add_chunk re-merges to size. - # Units that exceed the budget after the regex split (a single long line with - # no delimiter, e.g. PDF / .txt runs of unbroken text) are sub-split on - # whitespace atoms with a character-window fallback, mirroring the html path. + # Default path: split every section on the delimiter into paragraphs (no + # delimiter text), then group paragraphs with the chosen merge strategy. + # No atom-split is performed: a paragraph larger than ``chunk_token_num`` + # becomes its own chunk; the model layer truncates oversize units. + # + # A section is split on the delimiter whenever one is present -- even when + # the whole section already fits ``chunk_token_num``. The delimiter is a + # chunk boundary and its text must never leak into a chunk; only the + # empty-delimiter (size-only) mode below skips splitting. dels = compile_delimiter_pattern(parsed_dels) + paragraphs = [] # list of (text, pos) for sec, pos in sections: - sec_text = "\n" + sec - if num_tokens_from_string(sec_text) <= chunk_token_num: - add_chunk(sec_text, pos) + if not dels: + paragraphs.append(("\n" + sec, pos)) continue - if dels: - for sub_sec in re.split(r"(%s)" % dels, sec, flags=re.DOTALL): - if not sub_sec or re.fullmatch(dels, sub_sec): - continue - text = "\n" + sub_sec - if num_tokens_from_string(text) <= chunk_token_num: - add_chunk(text, pos) - else: - logging.debug("Splitting oversized unit (len=%d, tokens=%d) via _split_oversized_unit", len(text), num_tokens_from_string(text)) - for piece in _split_oversized_unit(text, chunk_token_num): - add_chunk(piece, pos) - else: - logging.debug("Splitting oversized unit (len=%d, tokens=%d) via _split_oversized_unit (no delimiters)", len(sec_text), num_tokens_from_string(sec_text)) - for piece in _split_oversized_unit(sec_text, chunk_token_num): - add_chunk(piece, pos) + for sub_sec in re.split(r"(%s)" % dels, sec, flags=re.DOTALL): + if not sub_sec or re.fullmatch(dels, sub_sec): + continue + paragraphs.append(("\n" + sub_sec, pos)) + groups = _merge_paragraph_groups([p[0] for p in paragraphs], chunk_token_num, strategy, num_tokens_from_string) + cks = [_reconstruct_text_chunk(paragraphs, g) for g in groups] logging.debug("naive_merge: %d sections -> %d chunks (delimiter=%r)", len(sections), len(cks), delimiter) - # Drop the leading empty placeholder that exists only so ``add_chunk`` could - # detect "first chunk ever" without an extra flag. - if cks and cks[0] == "": - cks = cks[1:] - tk_nums = tk_nums[1:] - return cks + return _apply_overlap_to_chunks(cks, overlapped_percent, chunk_token_num) -def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。；！？", overlapped_percent=0): +def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。；！？", overlapped_percent=0, strategy=MergeStrategy.OVER_CAP): + """Split texts (with images) into chunks. Chunking contract: see ``merge_paragraphs`` (refs #17799).""" if not texts or len(texts) != len(images): return [], [] - cks = [""] - result_images = [None] - tk_nums = [0] - - def add_chunk(t, image, pos=""): - nonlocal cks, result_images, tk_nums - action, text, tk_num = _compute_chunk_update(cks[-1], t, pos, chunk_token_num, overlapped_percent) - if action == "first": - cks[-1] = text - tk_nums[-1] = tk_num - result_images[-1] = image - elif action == "merge": - cks[-1] = text - tk_nums[-1] = tk_num - if result_images[-1] is None: - result_images[-1] = image - else: - result_images[-1] = concat_img(result_images[-1], image) - else: - cks.append(text) - result_images.append(image) - tk_nums.append(tk_num) # Parse the delimiter field once, via the canonical helper (#17383). - # See the matching block in ``naive_merge`` for the rationale on - # `has_custom` (backtick-wrapped tokens opt into chunk-token-num bypass). + # See ``naive_merge`` for the ``has_custom`` rationale. parsed_dels = parse_delimiter_field(delimiter) has_custom = has_wrapped_delimiter(delimiter) if has_custom: # Custom delimiters ignore chunk_token_num: each segment is its own chunk. custom_pattern = compile_delimiter_pattern(parsed_dels) - cks, result_images, tk_nums = [], [], [] + cks, result_images = [], [] for text, image in zip(texts, images): text_str = text[0] if isinstance(text, tuple) else text if text_str is None: @@ -1430,14 +1443,15 @@ def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。 text_seg += local_pos cks.append(text_seg) result_images.append(image) - tk_nums.append(num_tokens_from_string(text_seg)) return cks, result_images - # Split oversized sections at sentence delimiters; the section's image rides - # along on every piece (concat_img dedupes when pieces re-merge into a chunk). - # Units still exceeding the budget after the regex split are sub-split on - # whitespace atoms so they cannot blow past the token cap. + # Default path: split every text on the delimiter into paragraphs (no + # delimiter text) carrying its image, then group with the merge strategy. + # Images of merged paragraphs are concatenated; no atom-split is performed. + # As in ``naive_merge``, a small text is still split on the delimiter so + # the boundary text never leaks into a chunk; only empty-delimiter skips. dels = compile_delimiter_pattern(parsed_dels) + paragraphs = [] # list of (text, pos, image) for text, image in zip(texts, images): # if text is tuple, unpack it if isinstance(text, tuple): @@ -1447,32 +1461,22 @@ def naive_merge_with_images(texts, images, chunk_token_num=128, delimiter="\n。 text_str = text or "" text_pos = "" text_str = normalize_text_newlines(text_str) - text_seg = "\n" + text_str - if num_tokens_from_string(text_seg) <= chunk_token_num: - add_chunk(text_seg, image, text_pos) + if not dels: + paragraphs.append(("\n" + text_str, text_pos, image)) continue - if dels: - for sub_sec in re.split(r"(%s)" % dels, text_str, flags=re.DOTALL): - if not sub_sec or re.fullmatch(dels, sub_sec): - continue - sub_text = "\n" + sub_sec - if num_tokens_from_string(sub_text) <= chunk_token_num: - add_chunk(sub_text, image, text_pos) - else: - logging.debug("Splitting oversized unit (len=%d, tokens=%d) via _split_oversized_unit", len(sub_text), num_tokens_from_string(sub_text)) - for piece in _split_oversized_unit(sub_text, chunk_token_num): - add_chunk(piece, image, text_pos) - else: - logging.debug("Splitting oversized unit (len=%d, tokens=%d) via _split_oversized_unit (no delimiters)", len(text_seg), num_tokens_from_string(text_seg)) - for piece in _split_oversized_unit(text_seg, chunk_token_num): - add_chunk(piece, image, text_pos) + for sub_sec in re.split(r"(%s)" % dels, text_str, flags=re.DOTALL): + if not sub_sec or re.fullmatch(dels, sub_sec): + continue + paragraphs.append(("\n" + sub_sec, text_pos, image)) + groups = _merge_paragraph_groups([p[0] for p in paragraphs], chunk_token_num, strategy, num_tokens_from_string) + cks, result_images = [], [] + for g in groups: + text, image = _reconstruct_image_chunk(paragraphs, g) + cks.append(text) + result_images.append(image) logging.debug("naive_merge_with_images: %d texts -> %d chunks (delimiter=%r)", len(texts), len(cks), delimiter) - if cks and cks[0] == "": - cks = cks[1:] - result_images = result_images[1:] - tk_nums = tk_nums[1:] - return cks, result_images + return _apply_overlap_to_chunks(cks, overlapped_percent, chunk_token_num), result_images def docx_question_level(p, bull=-1): diff --git a/test/unit_test/deepdoc/parser/test_txt_parser.py b/test/unit_test/deepdoc/parser/test_txt_parser.py new file mode 100644 index 0000000000..0b571a1465 --- /dev/null +++ b/test/unit_test/deepdoc/parser/test_txt_parser.py @@ -0,0 +1,108 @@ +# +# Copyright 2025 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. +# + +"""Regression tests for ``RAGFlowTxtParser.parser_txt`` under the chunking contract. + +The contract (see ``rag.nlp.merge_paragraphs``, refs #17799): + +* delimiter = chunk boundary: delimiter text never enters a chunk; +* ``token_size`` = soft target + merge strategy (``OVER_CAP`` default); no + atom-split — a paragraph larger than ``chunk_token_num`` stands alone and the + model layer truncates it; +* ``UNDER_CAP`` is available as an explicit alternative strategy (never overflows + ``chunk_token_num``; ``OVER_CAP`` allows one boundary overflow). +""" + +from deepdoc.parser.txt_parser import RAGFlowTxtParser +import rag.nlp as nlp_mod + + +def _fake_word_tokens(s): + return len(s.split()) + + +def _nonempty(chunks): + return [c for c, _ in chunks if c.strip()] + + +def test_over_cap_accumulates_adjacent_paragraphs(monkeypatch): + monkeypatch.setattr(nlp_mod, "num_tokens_from_string", _fake_word_tokens) + text = "\n".join(["alpha beta gamma delta" for _ in range(8)]) # 4 tokens each + chunks = _nonempty(RAGFlowTxtParser.parser_txt(text, chunk_token_num=50, delimiter="\n")) + # OVER_CAP greedily accumulates adjacent paragraphs while under cap, instead + # of capping at fixed pairs: 8 * 4 = 32 tokens all fit under 50 -> 1 chunk. + assert len(chunks) == 1 + assert len(chunks[0].split()) == 32 + # Content is preserved (32 tokens total). + assert sum(len(c.split()) for c in chunks) == 32 + + +def test_oversize_unit_not_atom_split(monkeypatch): + monkeypatch.setattr(nlp_mod, "num_tokens_from_string", _fake_word_tokens) + text = "word " * 200 # ~200 tokens, no delimiter -> one paragraph + chunks = _nonempty(RAGFlowTxtParser.parser_txt(text, chunk_token_num=30, delimiter="\n!?;。；！？")) + # No atom-split: the whole unit is a single chunk. + assert len(chunks) == 1 + assert "".join(chunks).count("word") == 200 + + +def test_delimiter_text_not_in_chunk(monkeypatch): + monkeypatch.setattr(nlp_mod, "num_tokens_from_string", _fake_word_tokens) + text = "first##second##third" + chunks = _nonempty(RAGFlowTxtParser.parser_txt(text, chunk_token_num=1000, delimiter="##")) + assert all("##" not in c for c in chunks) + joined = "\n".join(chunks) + assert "first" in joined and "second" in joined and "third" in joined + + +def test_consecutive_delimiters_do_not_leak_delimiter_text(monkeypatch): + monkeypatch.setattr(nlp_mod, "num_tokens_from_string", _fake_word_tokens) + # pattern "##": consecutive delimiters must not glue the sides with "##". + text = "A####B" + chunks = _nonempty(RAGFlowTxtParser.parser_txt(text, chunk_token_num=1000, delimiter="##")) + joined = "\n".join(chunks) + assert "##" not in joined + assert "A" in joined and "B" in joined + + +def test_token_size_zero_keeps_each_paragraph_alone(monkeypatch): + monkeypatch.setattr(nlp_mod, "num_tokens_from_string", _fake_word_tokens) + text = "first second third" + chunks = _nonempty(RAGFlowTxtParser.parser_txt(text, chunk_token_num=0, delimiter=" ")) + assert chunks == ["first", "second", "third"] + + +def test_delimiter_boundary_when_segment_exceeds_cap(monkeypatch): + monkeypatch.setattr(nlp_mod, "num_tokens_from_string", _fake_word_tokens) + # Each paragraph is 2 tokens (> cap=1) -> its own chunk. + text = "aa aa\nbb bb\ncc cc" + chunks = _nonempty(RAGFlowTxtParser.parser_txt(text, chunk_token_num=1, delimiter="\n")) + assert chunks == ["aa aa", "bb bb", "cc cc"] + + +def test_keep_delimiters_preserves_delimiter(monkeypatch): + monkeypatch.setattr(nlp_mod, "num_tokens_from_string", _fake_word_tokens) + text = "first|second" + chunks = _nonempty(RAGFlowTxtParser.parser_txt(text, chunk_token_num=1000, delimiter="|", keep_delimiters=True)) + # When keep_delimiters=True the delimiter is retained in the chunk. + assert any("|" in c for c in chunks) + joined = "\n".join(chunks) + assert "first" in joined and "second" in joined + + +def test_empty_text_returns_empty(): + assert RAGFlowTxtParser.parser_txt("", chunk_token_num=128) == [] diff --git a/test/unit_test/rag/test_delimiter_case_sensitive.py b/test/unit_test/rag/test_delimiter_case_sensitive.py index 1b09f9eae5..b41f5a0b85 100644 --- a/test/unit_test/rag/test_delimiter_case_sensitive.py +++ b/test/unit_test/rag/test_delimiter_case_sensitive.py @@ -110,14 +110,25 @@ def force_every_section_above_budget(monkeypatch): def test_naive_merge_bare_char_a_splits_only_at_lowercase_a(): - """Bare-char ``a`` must split only at lowercase ``a``, not at ``A``.""" + """Bare-char ``a`` must split only at lowercase ``a``, not at ``A``. + + The delimiter produces two paragraphs ("B", "Ab") which the default + OVER_CAP merge pairs into one chunk (pairing may exceed cap). The + assertion therefore checks the *split point*: lowercase 'a' separates + "B" from "Ab" while the uppercase 'A' stays inline. + """ chunks = naive_merge(["BaAb"], chunk_token_num=8, delimiter="a") - assert [c.strip() for c in chunks if c.strip()] == ["B", "Ab"] + joined = "".join(chunks) + assert joined == "\nB\nAb" + # Case-insensitive matching would have split at 'A' too -> "Ba\\nb". + assert "Ba\nb" not in joined def test_naive_merge_bare_char_A_splits_only_at_uppercase_A(): chunks = naive_merge(["BaAb"], chunk_token_num=8, delimiter="A") - assert [c.strip() for c in chunks if c.strip()] == ["Ba", "b"] + joined = "".join(chunks) + assert joined == "\nBa\nb" + assert "B\nAb" not in joined def test_naive_merge_backtick_end_splits_only_at_lowercase_end(): diff --git a/test/unit_test/rag/test_merge_paragraphs.py b/test/unit_test/rag/test_merge_paragraphs.py new file mode 100644 index 0000000000..b9bce5e5da --- /dev/null +++ b/test/unit_test/rag/test_merge_paragraphs.py @@ -0,0 +1,206 @@ +# +# Copyright 2025 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. +# + +"""Unit tests for ``merge_paragraphs`` / ``MergeStrategy``. + +``merge_paragraphs`` is the single pure function that groups delimiter-split +paragraphs (no delimiter text) into chunks. It implements the two merge +strategies (see ``rag.nlp.merge_paragraphs`` for the full contract, refs #17799): + +* ``UNDER_CAP``: only merge the next paragraph when the projected total still + fits the soft ``token_size`` target. +* ``OVER_CAP`` (default): pair adjacent paragraphs even when the pair exceeds + ``token_size``; a paragraph larger than ``token_size`` stands alone. + +Neither strategy ever atom-splits a paragraph. +""" + +from rag.nlp import MergeStrategy, merge_paragraphs + + +def _paras_with_sizes(sizes): + """Build unique paragraph strings tagged with their token size.""" + paras = [f"p{i}" for i in range(len(sizes))] + + def size(p: str) -> int: + return sizes[int(p[1:])] # type: ignore[assignment] + + return paras, size + + +def _flatten(groups): + return [p for g in groups for p in g] + + +# --------------------------------------------------------------------------- # +# Contract acceptance examples (refs #17799) +# --------------------------------------------------------------------------- # + + +def test_under_cap_example(): + # cap=100, paragraph sizes 150/60/50/30 -> [[150], [60], [80]] + paras, size = _paras_with_sizes([150, 60, 50, 30]) + groups = merge_paragraphs(paras, 100, MergeStrategy.UNDER_CAP, size=size) + assert groups == [["p0"], ["p1"], ["p2", "p3"]] + + +def test_over_cap_example(): + # cap=100, paragraph sizes 150/60/50/30 -> [[150], [110], [30]] + paras, size = _paras_with_sizes([150, 60, 50, 30]) + groups = merge_paragraphs(paras, 100, MergeStrategy.OVER_CAP, size=size) + assert groups == [["p0"], ["p1", "p2"], ["p3"]] + + +def test_default_strategy_is_over_cap(): + paras, size = _paras_with_sizes([150, 60, 50, 30]) + groups = merge_paragraphs(paras, 100, size=size) + assert groups == [["p0"], ["p1", "p2"], ["p3"]] + + +# --------------------------------------------------------------------------- # +# Edge cases +# --------------------------------------------------------------------------- # + + +def test_empty_input(): + assert merge_paragraphs([], 100) == [] + + +def test_single_paragraph(): + assert merge_paragraphs(["only"], 100) == [["only"]] + + +def test_all_paragraphs_over_cap_stand_alone(): + paras, size = _paras_with_sizes([200, 300, 150]) + for strategy in (MergeStrategy.UNDER_CAP, MergeStrategy.OVER_CAP): + groups = merge_paragraphs(paras, 100, strategy, size=size) + assert groups == [["p0"], ["p1"], ["p2"]] + + +def test_all_under_cap_mergeable(): + paras, size = _paras_with_sizes([10, 20, 30]) + # UNDER_CAP: 10+20+30=60 <= 100 -> one chunk. + assert merge_paragraphs(paras, 100, MergeStrategy.UNDER_CAP, size=size) == [["p0", "p1", "p2"]] + # OVER_CAP: all 60 <= 100 -> one chunk (NOT pairwise [[10,20],[30]]). + assert merge_paragraphs(paras, 100, MergeStrategy.OVER_CAP, size=size) == [["p0", "p1", "p2"]] + + +def test_alternating_non_mergeable(): + # 60,60,60,60 with cap=100. + paras, size = _paras_with_sizes([60, 60, 60, 60]) + # UNDER_CAP: 60+60=120 > 100 -> every paragraph alone. + assert merge_paragraphs(paras, 100, MergeStrategy.UNDER_CAP, size=size) == [["p0"], ["p1"], ["p2"], ["p3"]] + # OVER_CAP: pairs. + assert merge_paragraphs(paras, 100, MergeStrategy.OVER_CAP, size=size) == [["p0", "p1"], ["p2", "p3"]] + + +def test_token_size_zero_every_paragraph_alone(): + paras, size = _paras_with_sizes([3, 2, 4]) + for strategy in (MergeStrategy.UNDER_CAP, MergeStrategy.OVER_CAP): + groups = merge_paragraphs(paras, 0, strategy, size=size) + assert len(groups) == 3 + assert all(len(g) == 1 for g in groups) + + +# --------------------------------------------------------------------------- # +# OVER_CAP greedy accumulation (contract: merge while projected total <= cap, +# allow one boundary overflow; oversized paragraph stands alone). +# See memory: feedback_over_cap_contract. +# --------------------------------------------------------------------------- # + + +def test_over_cap_accumulates_beyond_two(): + # 8 paragraphs of size 10 (total 80) under cap 128 must become ONE chunk, + # proving OVER_CAP accumulates past pairs instead of stopping at two. + paras, size = _paras_with_sizes([10] * 8) + groups = merge_paragraphs(paras, 128, MergeStrategy.OVER_CAP, size=size) + assert groups == [paras] + + +def test_over_cap_boundary_overflow(): + # 100+60 exceeds 128 -> OVER_CAP allows the boundary pair to overflow. + paras, size = _paras_with_sizes([100, 60, 100]) + groups = merge_paragraphs(paras, 128, MergeStrategy.OVER_CAP, size=size) + assert groups == [["p0", "p1"], ["p2"]] + + +def test_over_cap_vs_under_cap_boundary(): + # The ONLY semantic difference: OVER_CAP permits the boundary overflow. + paras, size = _paras_with_sizes([100, 60, 100]) + assert merge_paragraphs(paras, 128, MergeStrategy.UNDER_CAP, size=size) == [["p0"], ["p1"], ["p2"]] + assert merge_paragraphs(paras, 128, MergeStrategy.OVER_CAP, size=size) == [["p0", "p1"], ["p2"]] + + +def test_over_cap_oversized_stands_alone(): + # A paragraph larger than cap must never be paired (Bug A). + paras, size = _paras_with_sizes([60, 150, 60]) + groups = merge_paragraphs(paras, 128, MergeStrategy.OVER_CAP, size=size) + assert groups == [["p0"], ["p1"], ["p2"]] + + +def test_over_cap_oversized_then_accumulate(): + # Oversized boundary followed by normal accumulation in one input. + paras, size = _paras_with_sizes([10, 200, 10, 10, 10]) + groups = merge_paragraphs(paras, 128, MergeStrategy.OVER_CAP, size=size) + assert groups == [["p0"], ["p1"], ["p2", "p3", "p4"]] + + +def test_over_cap_single_oversized(): + paras, size = _paras_with_sizes([200]) + groups = merge_paragraphs(paras, 128, MergeStrategy.OVER_CAP, size=size) + assert groups == [["p0"]] + + +def test_token_size_one_delimiter_boundaries_not_one_token(): + # Token_size=1 on delimiter segments [3,2,4]: each segment is its own chunk + # (delimiter boundary), NEVER atom-split into 1-token pieces. + paras, size = _paras_with_sizes([3, 2, 4]) + for strategy in (MergeStrategy.UNDER_CAP, MergeStrategy.OVER_CAP): + groups = merge_paragraphs(paras, 1, strategy, size=size) + assert len(groups) == 3 + assert _flatten(groups) == paras + + +# --------------------------------------------------------------------------- # +# Invariants +# --------------------------------------------------------------------------- # + + +def test_whitespace_preserved_no_strip(): + paras = [" leading space", "internal space", "trailing space "] + groups = merge_paragraphs(paras, 100, MergeStrategy.OVER_CAP) + flat = _flatten(groups) + assert flat == paras + assert " leading space" in flat + + +def test_output_is_permutation_of_input_no_atom_split(): + # Every output paragraph must be exactly one input paragraph: no splitting, + # no duplication, no reordering. + paras, size = _paras_with_sizes([5, 17, 3, 42, 9]) + groups = merge_paragraphs(paras, 20, MergeStrategy.OVER_CAP, size=size) + flat = _flatten(groups) + assert sorted(flat) == sorted(paras) + # No paragraph was broken apart: each group entry is a whole input paragraph. + assert all(p in paras for g in groups for p in g) + + +def test_no_delimiter_text_introduced(): + paras = ["alpha", "beta", "gamma"] + groups = merge_paragraphs(paras, 100, MergeStrategy.OVER_CAP) + flat = _flatten(groups) + assert all("##" not in p for p in flat) diff --git a/test/unit_test/rag/test_naive_merge.py b/test/unit_test/rag/test_naive_merge.py index c61e78024d..ccd6ac6091 100644 --- a/test/unit_test/rag/test_naive_merge.py +++ b/test/unit_test/rag/test_naive_merge.py @@ -32,7 +32,7 @@ import re import pytest from rag import nlp -from rag.nlp import naive_merge, naive_merge_with_images +from rag.nlp import naive_merge, naive_merge_with_images, MergeStrategy DEFAULT_DELIMITER = "\n!?。；！？" @@ -77,21 +77,57 @@ 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 - # 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) + # OVER_CAP (default) allows at most one boundary paragraph to overflow the + # soft cap: a chunk may exceed ``chunk_token_num`` by one paragraph (10 + # tokens here) but never more. The old pairwise code packed to strictly + # ``<= cap``; greedy OVER_CAP instead closes the chunk right after the + # overflowing paragraph. + assert all(_tok(c) <= 50 + 10 for c in chunks) # Content is preserved. assert "".join(chunks).count("word") == 200 @pytest.mark.p2 -def test_small_sections_are_merged_not_oversplit(): +def test_small_section_is_split_at_delimiter_boundary(): + # A small section (well under chunk_token_num) that contains a delimiter + # must still be broken at the delimiter: the delimiter is a chunk boundary + # and its text must never leak into a chunk. The old code kept the whole + # section when it fit, so the delimiter text survived inside one chunk. + small_section = "first part。second part。third part" # 6 words, 2 delimiters + chunks = _nonempty(naive_merge([small_section], chunk_token_num=128, delimiter=DEFAULT_DELIMITER)) + # Delimiter text never appears inside any chunk. + assert all("。" not in c for c in chunks) + # Every delimiter-separated piece is present (content preserved). + joined = "".join(chunks) + assert "first part" in joined and "second part" in joined and "third part" in joined + + +@pytest.mark.p2 +def test_small_section_with_images_split_at_delimiter_boundary(): + # Same guarantee for the image path: a small text carrying an image is + # still split at the delimiter so the delimiter text does not leak. + small_section = "alpha。beta。gamma" # 3 words, 2 delimiters + texts = [(small_section, "")] + images = [object()] + chunks, imgs = naive_merge_with_images(texts, images, chunk_token_num=128, delimiter=DEFAULT_DELIMITER) + nonempty = _nonempty(chunks) + assert all("。" not in c for c in nonempty) + # The single image travels with its (split) text. + assert len(chunks) == len(imgs) + + +@pytest.mark.p2 +def test_small_sections_accumulate_under_over_cap(): + # Default strategy is OVER_CAP: adjacent small paragraphs are greedily + # accumulated while the projected total stays under chunk_token_num, not + # capped at fixed-size pairs. No atom-split is performed; the delimiter + # boundary (paragraph) is the unit. sentences = ["alpha beta gamma delta" for _ in range(8)] # 4 tokens each chunks = _nonempty(naive_merge(sentences, chunk_token_num=50, delimiter=DEFAULT_DELIMITER)) - # All 32 tokens comfortably fit one chunk. assert len(chunks) == 1 assert _tok(chunks[0]) == 32 + # Content is preserved (32 tokens total). + assert sum(_tok(c) for c in chunks) == 32 @pytest.mark.p2 @@ -121,11 +157,14 @@ def test_overlap_prefix_is_counted_in_token_budget(): # 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)) + # UNDER_CAP (strict): content chunks never overflow chunk_token_num, so the + # overlap-prefix budget check is the only thing under test here. + chunks = _nonempty(naive_merge(sentences, chunk_token_num=50, delimiter=DEFAULT_DELIMITER, overlapped_percent=20, strategy=MergeStrategy.UNDER_CAP)) assert len(chunks) > 1 - # Each 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. + # Each content chunk stays within the budget. Sentences are 10 tokens, the + # budget is 50, so a 5-sentence chunk is exactly 50; a 10-token overlap + # prefix (20% of 50) would push it to 60 and is therefore dropped at the + # boundary rather than letting the chunk overshoot. assert all(_tok(c) <= 50 for c in chunks) @@ -168,7 +207,8 @@ 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 for c in nonempty) + # OVER_CAP allows one boundary paragraph (10 tokens) to overflow the cap. + assert all(_tok(c) <= 50 + 10 for c in nonempty) @pytest.mark.p2 @@ -233,22 +273,29 @@ def test_images_distinct_lazyimages_are_concatenated(): @pytest.mark.p2 def test_strict_cap_no_overlap_packs_to_budget(): + # "strict cap" == UNDER_CAP: chunks never overflow chunk_token_num. sections = [" ".join(["w"] * 25) for _ in range(8)] - chunks = _nonempty(naive_merge(sections, chunk_token_num=50, delimiter=DEFAULT_DELIMITER)) + chunks = _nonempty(naive_merge(sections, chunk_token_num=50, delimiter=DEFAULT_DELIMITER, strategy=MergeStrategy.UNDER_CAP)) 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(): + # UNDER_CAP chunks are exactly 20 tokens (two 10-token sentences). A 20% + # overlap prefix is 4 tokens; 20 + 4 > 20, so the prefix is dropped at the + # boundary instead of letting the chunk overshoot the strict cap. 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) + chunks = _nonempty(naive_merge(sentences, chunk_token_num=20, delimiter=DEFAULT_DELIMITER, overlapped_percent=20, strategy=MergeStrategy.UNDER_CAP)) + assert len(chunks) > 1 + assert all(_tok(c) <= 20 for c in chunks) @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 test_no_atom_split_keeps_oversize_unit_whole(monkeypatch): + # The strict-cap atom sub-splitter is gone. A single unbroken unit that + # exceeds chunk_token_num is kept whole (the model layer truncates); this + # is the fix for the token_size=1 -> 1-token-per-chunk regression. def char_count_tokens(s): return len(s or "") @@ -256,15 +303,17 @@ def test_strict_cap_single_overlong_section_is_sub_split_on_whitespace(monkeypat 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 + assert len(chunks) == 1 + assert "".join(chunks).strip() == 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)) + # UNDER_CAP packs two 7-token sentences into a 14-token chunk, leaving + # headroom. A 20% overlap prefix is 4 tokens; 14 + 4 <= 20, so the prefix is + # kept (the overlap is chosen because it fits the strict cap). + sentences = [" ".join(["w"] * 7) for _ in range(20)] + chunks = _nonempty(naive_merge(sentences, chunk_token_num=20, delimiter=DEFAULT_DELIMITER, overlapped_percent=20, strategy=MergeStrategy.UNDER_CAP)) assert all(_tok(c) <= 20 for c in chunks) overlap_seen = False for a, b in zip(chunks, chunks[1:]): @@ -280,7 +329,7 @@ def test_strict_cap_overlap_chosen_when_it_fits(): 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) + chunks, imgs = naive_merge_with_images(sections, images, chunk_token_num=50, delimiter=DEFAULT_DELIMITER, strategy=MergeStrategy.UNDER_CAP) nonempty = _nonempty(chunks) assert all(_tok(c) <= 50 for c in nonempty) assert len(chunks) == len(imgs) @@ -295,30 +344,34 @@ def test_strict_cap_pos_text_does_not_overshoot_budget(monkeypatch): 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. + # section is ~15 chars, pos is 10 chars. chunk_token_num is 20. + # NOTE: ``pos`` is attached post-merge (in ``_reconstruct_text_chunk``), so it + # is NOT counted in the merge-paragraphs token budget. Greedy UNDER_CAP packs + # the text up to the cap; the reconstructed chunk then gains the pos tag on + # top. The honest bound is therefore ``cap + len(pos)`` — the pos tag is not + # budgeted away. (Accounting pos in the merge budget would be a separate fix.) 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) + chunks = _nonempty(naive_merge(sections, chunk_token_num=20, delimiter=DEFAULT_DELIMITER, strategy=MergeStrategy.UNDER_CAP)) + assert all(char_count_tokens(c) <= 20 + len(pos_tag) 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. +def test_empty_delimiter_keeps_unit_whole(): + # Empty delimiter -> no split; the whole section is one chunk (no atom-split). + # The model layer truncates oversize units. 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) + assert len(chunks) == 1 + assert "".join(chunks).count("word") == 100 @pytest.mark.p2 -def test_images_empty_delimiter_oversized_section_strictly_capped(): +def test_images_empty_delimiter_keeps_unit_whole(): 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(nonempty) == 1 + assert "".join(nonempty).count("word") == 100 assert len(chunks) == len(imgs)