From e54e7ec7ef6c42b89e966473e7a2d96296997629 Mon Sep 17 00:00:00 2001 From: Yash Raj Pandey <55940078+devYRPauli@users.noreply.github.com> Date: Mon, 13 Jul 2026 10:04:56 -0400 Subject: [PATCH] Fix: TokenChunker discards TitleChunker chunks when output_format is 'chunks' (#16825) Fixes #16812 ### Problem In the `rag/flow` ingestion pipeline, when `TitleChunker` feeds `TokenChunker`, the chapter-aware chunks are silently discarded and the parser's raw flat json is re-chunked instead. `TitleChunker` emits `output_format="chunks"` and writes its chapter-aware output to the `chunks` field (`rag/flow/chunker/title_chunker/common.py`, `set_output("output_format", "chunks")`). But `TokenChunker._invoke` only handles `output_format` in `["markdown", "text", "html"]`, then falls through to the `# json` path which reads `from_upstream.json_result`. There is no branch for `"chunks"`, so `from_upstream.chunks` is never read. Downstream effects reported in #16812: PageIndex/TOC extraction receives flat line-level text instead of structured chapter blocks (incorrect/duplicate/missing chapters), and retrieval quality degrades because chunks are no longer aligned to document structure. ### Fix Select the source list based on `output_format`, mirroring the exact pattern already used in `title_chunker/common.py`: ```python json_result = (from_upstream.chunks if from_upstream.output_format == "chunks" else from_upstream.json_result) or [] ``` `chunks` items share the same dict shape as `json_result` items (both consumed via `.get("text")`, `.get("doc_type_kwd")`, etc.), so they flow through the existing token-sizing path unchanged. One-line change, no behavior change for the `json`/`markdown`/`text`/`html` paths. ### Test Adds `rag/flow/tests/test_token_chunker.py`, an isolated unit test that runs the real `TokenChunker._invoke` (heavy deps stubbed; real pydantic schema used when available) and asserts that with `output_format="chunks"` the upstream `chunks` are consumed rather than the raw parser `json`. Verified RED -> GREEN: the test fails against the current code (reads the raw json) and passes with the fix. Signed-off-by: Yash Raj Pandey --- rag/flow/chunker/token_chunker.py | 2 +- rag/flow/tests/test_token_chunker.py | 187 +++++++++++++++++++++++++++ 2 files changed, 188 insertions(+), 1 deletion(-) create mode 100644 rag/flow/tests/test_token_chunker.py diff --git a/rag/flow/chunker/token_chunker.py b/rag/flow/chunker/token_chunker.py index d39a6583e4..80e029a1d5 100644 --- a/rag/flow/chunker/token_chunker.py +++ b/rag/flow/chunker/token_chunker.py @@ -344,7 +344,7 @@ class TokenChunker(ProcessBase): return # json - json_result = from_upstream.json_result or [] + json_result = (from_upstream.chunks if from_upstream.output_format == "chunks" else from_upstream.json_result) or [] if self._param.delimiter_mode == "one": sections = [] for item in json_result: diff --git a/rag/flow/tests/test_token_chunker.py b/rag/flow/tests/test_token_chunker.py new file mode 100644 index 0000000000..e50dba6fc6 --- /dev/null +++ b/rag/flow/tests/test_token_chunker.py @@ -0,0 +1,187 @@ +import importlib.util +import asyncio +import sys +import types +from contextlib import contextmanager +from pathlib import Path + +@contextmanager +def _load_token_chunker_with_stubs(): + root = Path(__file__).resolve().parents[3] + original_modules = {} + + def _install(name: str, module: types.ModuleType): + original_modules.setdefault(name, sys.modules.get(name)) + sys.modules[name] = module + + try: + rag_pkg = types.ModuleType("rag") + rag_pkg.__path__ = [str(root / "rag")] + _install("rag", rag_pkg) + + rag_flow_pkg = types.ModuleType("rag.flow") + rag_flow_pkg.__package__ = "rag" + rag_flow_pkg.__path__ = [str(root / "rag" / "flow")] + _install("rag.flow", rag_flow_pkg) + + rag_flow_chunker_pkg = types.ModuleType("rag.flow.chunker") + rag_flow_chunker_pkg.__package__ = "rag.flow" + rag_flow_chunker_pkg.__path__ = [str(root / "rag" / "flow" / "chunker")] + _install("rag.flow.chunker", rag_flow_chunker_pkg) + + rag_flow_parser_pkg = types.ModuleType("rag.flow.parser") + rag_flow_parser_pkg.__package__ = "rag.flow" + rag_flow_parser_pkg.__path__ = [str(root / "rag" / "flow" / "parser")] + _install("rag.flow.parser", rag_flow_parser_pkg) + + common_pkg = types.ModuleType("common") + common_pkg.__path__ = [str(root / "common")] + _install("common", common_pkg) + + common_float_utils = types.ModuleType("common.float_utils") + common_float_utils.normalize_overlapped_percent = lambda value: value + _install("common.float_utils", common_float_utils) + + common_token_utils = types.ModuleType("common.token_utils") + common_token_utils.num_tokens_from_string = lambda text: 1 + _install("common.token_utils", common_token_utils) + + rag_nlp = types.ModuleType("rag.nlp") + rag_nlp.naive_merge = lambda *_args, **_kwargs: [] + _install("rag.nlp", rag_nlp) + + class ProcessParamBase: + def __init__(self): + pass + + class ProcessBase: + def __init__(self, _pipeline, _id, param): + self._pipeline = _pipeline + self._id = _id + self._param = param + self._outputs = {} + self.callback = lambda *_args, **_kwargs: None + + def set_output(self, key, value): + self._outputs[key] = value + + rag_flow_base = types.ModuleType("rag.flow.base") + rag_flow_base.ProcessBase = ProcessBase + rag_flow_base.ProcessParamBase = ProcessParamBase + _install("rag.flow.base", rag_flow_base) + + rag_flow_parser_pdf_metadata = types.ModuleType("rag.flow.parser.pdf_chunk_metadata") + rag_flow_parser_pdf_metadata.PDF_POSITIONS_KEY = "pdf_positions" + rag_flow_parser_pdf_metadata.extract_pdf_positions = lambda _item: [] + rag_flow_parser_pdf_metadata.finalize_pdf_chunk = lambda chunk: chunk + + async def restore_pdf_text_previews(*_args, **_kwargs): + return None + + rag_flow_parser_pdf_metadata.restore_pdf_text_previews = restore_pdf_text_previews + _install("rag.flow.parser.pdf_chunk_metadata", rag_flow_parser_pdf_metadata) + + try: + import pydantic # noqa: F401 + + schema_spec = importlib.util.spec_from_file_location( + "rag.flow.chunker.schema", + root / "rag" / "flow" / "chunker" / "schema.py", + ) + if schema_spec is None or schema_spec.loader is None: + raise RuntimeError("Failed to locate rag.flow.chunker.schema stub loader.") + schema_module = importlib.util.module_from_spec(schema_spec) + _install("rag.flow.chunker.schema", schema_module) + schema_spec.loader.exec_module(schema_module) + except Exception: + schema_module = types.ModuleType("rag.flow.chunker.schema") + + class TokenChunkerFromUpstream: + def __init__( + self, + name, + file=None, + chunks=None, + output_format=None, + json_result=None, + markdown_result=None, + text_result=None, + html_result=None, + _created_time=None, + _elapsed_time=None, + ): + self.name = name + self.file = file + self.chunks = chunks + self.output_format = output_format + self.json_result = json_result + self.json = json_result + self.markdown_result = markdown_result + self.markdown = markdown_result + self.text_result = text_result + self.text = text_result + self.html_result = html_result + self.html = html_result + self._created_time = _created_time + self._elapsed_time = _elapsed_time + + @classmethod + def model_validate(cls, data): + if isinstance(data, dict): + return cls( + name=data.get("name", ""), + file=data.get("file"), + chunks=data.get("chunks"), + output_format=data.get("output_format"), + json_result=data.get("json_result", data.get("json")), + markdown_result=data.get("markdown_result", data.get("markdown")), + text_result=data.get("text_result", data.get("text")), + html_result=data.get("html_result", data.get("html")), + _created_time=data.get("_created_time"), + _elapsed_time=data.get("_elapsed_time"), + ) + raise TypeError("TokenChunkerFromUpstream expects a dict payload.") + + schema_module.TokenChunkerFromUpstream = TokenChunkerFromUpstream + _install("rag.flow.chunker.schema", schema_module) + + token_chunker_spec = importlib.util.spec_from_file_location( + "rag.flow.chunker.token_chunker", + root / "rag" / "flow" / "chunker" / "token_chunker.py", + ) + if token_chunker_spec is None or token_chunker_spec.loader is None: + raise RuntimeError("Failed to locate rag.flow.chunker.token_chunker stub loader.") + token_chunker_module = importlib.util.module_from_spec(token_chunker_spec) + _install("rag.flow.chunker.token_chunker", token_chunker_module) + token_chunker_spec.loader.exec_module(token_chunker_module) + yield token_chunker_module + finally: + for module_name, original in original_modules.items(): + if original is None: + sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = original + + +def test_token_chunker_prefers_upstream_chunks_for_json_output_format_chunks(): + # Regression for #16812: when the upstream (e.g. TitleChunker) emits + # output_format="chunks", TokenChunker must consume from_upstream.chunks and + # not fall through to the raw parser json_result. Heavy deps are stubbed so + # the real TokenChunker._invoke runs against the real schema when pydantic is + # available (see title_chunker/common.py for the same chunks-vs-json branch). + with _load_token_chunker_with_stubs() as token_chunker_module: + token_chunker = token_chunker_module.TokenChunker + param = token_chunker_module.TokenChunkerParam() + param.delimiter_mode = "one" + chunker = token_chunker(None, "token_chunker", param) + + kwargs = { + "name": "token_chunker", + "output_format": "chunks", + "chunks": [{"text": "CHAPTER-AWARE"}], + "json": [{"text": "RAW-PARSER-JSON"}], + } + + asyncio.run(chunker._invoke(**kwargs)) + + assert chunker._outputs["chunks"] == [{"text": "CHAPTER-AWARE"}]