mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-05 15:20:30 +08:00
1f8e9af644
## Summary
`TokenChunker._invoke` discarded the user's configured `delimiters`
whenever no backtick-wrapped delimiter was present: it passed a
hardcoded `""` to `naive_merge`, so the configured delimiters (including
the default `["\n"]`) were never forwarded. `naive_merge` then ignored
the newline sentence boundary and cut chunks **mid-sentence** once the
token budget was exceeded.
## Root cause
`token_chunker.py:326` called `naive_merge(payload, chunk_token_size,
"", overlapped_percent)`. `_compile_delimiter_pattern` intentionally
returns `""` for bare (non-backtick) delimiters — that return value is a
*path selector* (empty → token-budget merge; non-empty → hard
`_split_text_by_pattern` split). The bug was not in that selector but in
the `else` branch, which threw away `self._param.delimiters` instead of
forwarding it.
## Fix
Forward the configured delimiters as a soft boundary:
```python
else naive_merge(
payload,
self._param.chunk_token_size,
"".join(self._param.delimiters),
overlapped_percent,
)
```
`naive_merge` already parses the string via the canonical
`parse_delimiter_field`, so bare and backtick-wrapped delimiters are
honored as soft boundaries while the token budget is still respected.
The path-selection role of `_compile_delimiter_pattern` is untouched:
backtick-wrapped delimiters still select the hard
`_split_text_by_pattern` path; bare delimiters still take the
token-budget merge path. No regression for any previously-working
(wrapped-delimiter) configuration.
## Test
Adds `rag/flow/tests/test_token_chunker_delimiter.py`:
- `test_token_chunker_token_size_mode_does_not_split_sentences` — fails
on the old code (3 sentences cut mid-stream), passes after the fix.
- `test_naive_merge_empty_delimiter_ignores_newline_break` — root-cause
companion asserting `naive_merge("")` cuts while `naive_merge("\n")`
preserves boundaries.
Both tests skip when the tokenizer is unavailable (dead-tokenizer
guard).
## Note
This branch contains only this one-line fix on top of `main`; it is
intentionally independent of the unrelated tokenizer/offline-BPE work.
Co-authored-by: CodeBuddy <noreply@codebuddy.ai>
118 lines
5.3 KiB
Python
118 lines
5.3 KiB
Python
"""Regression test for token_size-mode delimiter handling in TokenChunker.
|
||
|
||
The default ``delimiters`` a user supplies (e.g. ``["\\n"]``) are only honoured
|
||
when they are backtick-wrapped; a bare delimiter makes
|
||
``_compile_delimiter_pattern`` return ``""``, so ``TokenChunker._invoke`` falls
|
||
through to ``naive_merge(payload, chunk_token_size, "")`` (token_chunker.py:323).
|
||
|
||
Passing ``""`` discards the sentence-boundary preference that ``naive_merge``
|
||
carries by default (``"\\n。；！？"``). Without a forced ``\\n`` break, the token
|
||
budget flush can land *inside* a sentence, so a line like
|
||
``"Sentence number 7. alpha beta gamma delta epsilon zeta eta theta iota kappa"``
|
||
gets split across two chunks. The Go port (``sentenceDelimiter = (\\n|[!?。；！？])``
|
||
in token.go:340) always breaks on ``\\n``, so it never cuts a sentence mid-stream.
|
||
|
||
This test exposes that divergence: with ``delimiters=["\\n"]`` in token_size mode,
|
||
every original sentence must survive intact in a single chunk. It currently FAILS
|
||
because of the ``""`` passed at token_chunker.py:326; forwarding
|
||
``"".join(self._param.delimiters)`` to ``naive_merge`` (or at least
|
||
``"\\n。；！？"``) makes it pass and aligns Python with the Go side.
|
||
|
||
Note: the test needs a live tiktoken encoder. ``common.token_utils`` returns 0 on
|
||
any encoder failure, which would make every token budget "never exceeded" and
|
||
hide the bug behind a single-chunk baseline. We skip rather than pass under a
|
||
dead tokenizer, mirroring ``capture_golden.py::_assert_tokenizer_alive``.
|
||
"""
|
||
|
||
import asyncio
|
||
import types
|
||
|
||
from rag.flow.chunker.token_chunker import TokenChunker, TokenChunkerParam
|
||
from rag.nlp import naive_merge
|
||
from common.token_utils import num_tokens_from_string
|
||
|
||
|
||
def _build_token_chunker(param: dict) -> TokenChunker:
|
||
"""Construct a TokenChunker without the real Graph the normal __init__ wants.
|
||
|
||
``ComponentBase.__init__`` asserts ``canvas`` is a real ``Graph``; a unit test
|
||
has none. Bypass it and attach the three attributes ``_invoke`` actually reads
|
||
(``_canvas``, ``_param``, ``callback``) so the genuine chunking path runs with
|
||
the real ``naive_merge`` and the real ``num_tokens_from_string``. Mirrors
|
||
``capture_golden.py::_build_component``.
|
||
"""
|
||
p = TokenChunkerParam()
|
||
for key, value in param.items():
|
||
setattr(p, key, value)
|
||
p.check() # __new__ bypassed the constructor's validation; reproduce it.
|
||
|
||
comp = TokenChunker.__new__(TokenChunker)
|
||
comp._canvas = types.SimpleNamespace(_doc_id=None, _tenant_id="t")
|
||
comp._param = p
|
||
comp.callback = lambda *_a, **_kw: None
|
||
return comp
|
||
|
||
|
||
def _invoke_text(comp: TokenChunker, text: str) -> list[dict]:
|
||
asyncio.run(comp._invoke(name="t", output_format="text", text=text))
|
||
return comp._param.outputs["chunks"]["value"]
|
||
|
||
|
||
def _make_sentences(n: int) -> list[str]:
|
||
# Each sentence is ~16 tokens of plain English; well under the 128 budget,
|
||
# so several fit per chunk and the token-budget flush lands mid-sentence.
|
||
return [f"Sentence number {i}. alpha beta gamma delta epsilon zeta eta theta iota kappa" for i in range(n)]
|
||
|
||
|
||
def test_token_chunker_token_size_mode_does_not_split_sentences():
|
||
# Guard: a dead tokenizer would collapse everything to one chunk and hide
|
||
# the bug. Skip instead of recording a poisoned pass.
|
||
if num_tokens_from_string("alive tokenizer probe sentence") <= 0:
|
||
import pytest
|
||
|
||
pytest.skip("tiktoken encoder unavailable; num_tokens_from_string returned 0")
|
||
|
||
sentences = _make_sentences(30)
|
||
payload = "\n".join(sentences)
|
||
|
||
comp = _build_token_chunker({"chunk_token_size": 128, "delimiters": ["\n"]})
|
||
chunks = _invoke_text(comp, payload)
|
||
texts = [c["text"] for c in chunks]
|
||
|
||
# An original sentence is "intact" when its full text appears in exactly one
|
||
# chunk. If the budget flush fell inside it, the sentence is split across two
|
||
# chunks and this assertion fails.
|
||
split_sentences = [s for s in sentences if not any(s in t for t in texts)]
|
||
assert not split_sentences, (
|
||
f"{len(split_sentences)} sentence(s) were cut mid-stream by "
|
||
f"token_size mode with delimiters=['\\n'] (got {len(chunks)} chunks). "
|
||
"naive_merge was called with an empty delimiter (token_chunker.py:326), "
|
||
"dropping the '\\n' sentence-boundary preference. Example split: "
|
||
f"{split_sentences[0]!r}"
|
||
)
|
||
|
||
|
||
def test_naive_merge_empty_delimiter_ignores_newline_break():
|
||
"""Root-cause check: naive_merge('') can cut mid-sentence; naive_merge('\\n') cannot.
|
||
|
||
Documents why forwarding the configured delimiters fixes the TokenChunker bug.
|
||
"""
|
||
if num_tokens_from_string("alive tokenizer probe sentence") <= 0:
|
||
import pytest
|
||
|
||
pytest.skip("tiktoken encoder unavailable; num_tokens_from_string returned 0")
|
||
|
||
sentences = _make_sentences(30)
|
||
payload = "\n".join(sentences)
|
||
|
||
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)]
|
||
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"
|