zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-07 16:11:16 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix: delimiter is chunk boundary, drop token_size atom-split (OVER_CAP default) (#17808)

Browse Source 
## 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 <noreply@tencent.com>
This commit is contained in:
Jack
2026-08-05 11:50:07 +08:00
committed by GitHub
parent 302e611a43
commit 9b05e5c67e
9 changed files with 879 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
deepdoc/parser/txt_parser.py
View File

@@ -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]

77
rag/flow/chunker/token_chunker.py
View File

@@ -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))

178
rag/flow/tests/test_token_chunker.py
View File

@@ -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}"

15
rag/flow/tests/test_token_chunker_delimiter.py
View File

@@ -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"

408
rag/nlp/__init__.py
View File

@@ -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):

108
test/unit_test/deepdoc/parser/test_txt_parser.py Normal file
View File

@@ -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) == []

17
test/unit_test/rag/test_delimiter_case_sensitive.py
View File

@@ -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():

206
test/unit_test/rag/test_merge_paragraphs.py Normal file
View File

@@ -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)

123
test/unit_test/rag/test_naive_merge.py
View File

@@ -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)