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 <yashpn62@gmail.com>

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:

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