"""Mermaid-aware chunking helpers shared by hierarchical_chunk_* DAG tasks. The line-based custom chunker (`chunk_with_atomic_mermaid`) is used **only** by the Excel branch: Docling-xlsx markdown has almost no `\n\n` paragraph anchors, so the stock `SentenceSplitter(paragraph_separator="\n\n")` falls through to char-level splits that bisect table rows mid-cell. Word / PDF / PPTX and `.md` inputs have normal paragraph structure and use a stock 2-level `MarkdownNodeParser -> SentenceSplitter(512)` hierarchy via `build_text_parser_map`. Public API: chunk_with_atomic_mermaid(text, target, max_size, tokenize) Excel-only line-based chunker. Treats each ```mermaid fence as an atomic line-group; pure-prose chunks stay <= target, chunks that contain a fence may grow up to max_size. Backward-merges tail chunks < MIN_CHUNK_TOKENS within the same section. split_oversize_mermaid(block_lines, tokenize, limit) Split a single fence whose total tokens exceed `limit` along its body lines, rewrapping each part as its own ```mermaid fence. prune_header_only_nodes(nodes) Drop any node whose stripped text matches HEADER_ONLY_RE; remove ancestors whose entire subtree was pruned. build_excel_parser_map(tokenize_fn, embed_max_tokens, target_tokens=512) Parser map for the Excel branch ({markdown, chunk_size_512}). build_text_parser_map(tokenizer) Stock 2-level parser map for non-md / md branches ({markdown, chunk_size_512}). """ from __future__ import annotations import re from typing import Any, Callable, Dict, List, Sequence, Tuple MERMAID_FENCE_START_RE = re.compile(r"^```mermaid\b") FENCE_END_RE = re.compile(r"^```\s*$") HEADER_ONLY_RE = re.compile(r"^\s*#{1,6} [^\n]+\s*$") MIN_CHUNK_TOKENS = 64 EXCEL_PARSER_IDS: List[str] = ["markdown", "chunk_size_512"] TEXT_PARSER_IDS: List[str] = ["markdown", "chunk_size_512"] def split_oversize_mermaid( block_lines: List[str], tokenize: Callable[[str], int], limit: int ) -> List[str]: """Split a single mermaid fence whose total tokens exceed `limit`. Splits along the block's inner body lines and rewraps each part as its own ```mermaid...``` fence. The opener (e.g. ```` ```mermaid ````) and closer (```` ``` ````) are preserved verbatim per part. Raises ValueError if a single inner body line by itself exceeds the limit (i.e. cannot be made to fit even alone inside the wrapper). """ opener, closer = block_lines[0], block_lines[-1] inner = block_lines[1:-1] wrapper_tok = tokenize(opener + "\n\n" + closer) parts: List[str] = [] cur: List[str] = [] cur_tok = wrapper_tok for ln in inner: ln_tok = tokenize(ln + "\n") if wrapper_tok + ln_tok > limit: raise ValueError( f"single mermaid line exceeds embedding context ({ln_tok} tokens)" ) if cur_tok + ln_tok > limit and cur: parts.append(opener + "\n" + "\n".join(cur) + "\n" + closer) cur = [ln] cur_tok = wrapper_tok + ln_tok else: cur.append(ln) cur_tok += ln_tok if cur: parts.append(opener + "\n" + "\n".join(cur) + "\n" + closer) return parts def chunk_with_atomic_mermaid( text: str, target: int, max_size: int, tokenize: Callable[[str], int], ) -> List[Tuple[str, bool]]: """Line-based chunker that treats ```mermaid fences as atomic line-groups. Returns a list of `(chunk_text, contains_mermaid)` tuples. Pure-prose chunks stay <= `target`. A chunk containing at least one mermaid fence may grow up to `max_size` so the diagram stays adjacent to its explanatory prose. A single fence > `max_size` is split internally via `split_oversize_mermaid`. After chunking, a backward-merge pass folds chunks < MIN_CHUNK_TOKENS into the previous chunk if the combined size still fits within `max_size`. This eliminates tail-fragment leaves (e.g. lone table rows that fell past a flush boundary). Cross-section merging is the caller's responsibility -- this function never sees section boundaries. """ lines = text.split("\n") chunks: List[Tuple[str, bool]] = [] cur: List[str] = [] cur_tok = 0 cur_has_mermaid = False def flush() -> None: nonlocal cur, cur_tok, cur_has_mermaid if cur: chunks.append(("\n".join(cur), cur_has_mermaid)) cur, cur_tok, cur_has_mermaid = [], 0, False i = 0 n = len(lines) while i < n: ln = lines[i] if MERMAID_FENCE_START_RE.match(ln): # Locate the closing ``` line; tolerate unterminated fences by # consuming to end-of-text. j = i + 1 while j < n and not FENCE_END_RE.match(lines[j]): j += 1 end = j if j < n else n - 1 block_lines = lines[i:end + 1] block_text = "\n".join(block_lines) block_tok = tokenize(block_text) if block_tok > max_size: flush() for part in split_oversize_mermaid(block_lines, tokenize, max_size): chunks.append((part, True)) i = end + 1 continue if cur_tok + block_tok > max_size and cur: flush() cur.extend(block_lines) cur_tok += block_tok cur_has_mermaid = True i = end + 1 else: tok_n = tokenize(ln + "\n") limit = max_size if cur_has_mermaid else target if cur_tok + tok_n > limit and cur: flush() cur.append(ln) cur_tok += tok_n i += 1 flush() # Backward tail-merge within the supplied section. merged: List[Tuple[str, bool]] = [] for chunk_text, has_m in chunks: if merged and tokenize(chunk_text) < MIN_CHUNK_TOKENS: prev_text, prev_has = merged[-1] combined = prev_text + "\n" + chunk_text if tokenize(combined) <= max_size: merged[-1] = (combined, prev_has or has_m) continue merged.append((chunk_text, has_m)) return merged def prune_header_only_nodes(nodes: Sequence[Any]) -> List[Any]: """Drop nodes whose stripped text is just an ATX header line. Iteratively removes nodes matching HEADER_ONLY_RE plus any ancestor whose entire surviving subtree was pruned. Preserves topological order. Operates on llama_index BaseNode-like objects (uses `node_id`, `text`, and `relationships[CHILD]`). """ from llama_index.core.schema import NodeRelationship drop: set = set() for n in nodes: if HEADER_ONLY_RE.match(n.text): drop.add(n.node_id) changed = True while changed: changed = False for n in nodes: if n.node_id in drop: continue child_rel = n.relationships.get(NodeRelationship.CHILD) if not child_rel: continue children = child_rel if isinstance(child_rel, list) else [child_rel] if children and all(r.node_id in drop for r in children): drop.add(n.node_id) changed = True return [n for n in nodes if n.node_id not in drop] def build_text_parser_map(tokenizer: Callable[[str], List[int]]) -> Dict[str, Any]: """Stock 2-level hierarchy for Word / PDF / PPTX / .md branches. Returns a parser map suitable for `HierarchicalNodeParser` with `node_parser_ids = TEXT_PARSER_IDS`. The `SentenceSplitter` uses 20% overlap and `paragraph_separator="\\n\\n"` (markdown anchors). """ from llama_index.core.node_parser import MarkdownNodeParser, SentenceSplitter return { "markdown": MarkdownNodeParser(), "chunk_size_512": SentenceSplitter( chunk_size=512, chunk_overlap=int(512 * 0.2), tokenizer=tokenizer, paragraph_separator="\n\n", ), } def build_excel_parser_map( tokenize_fn: Callable[[str], int], embed_max_tokens: int, ) -> Dict[str, Any]: """Parser map for the Excel branch (keys: `EXCEL_PARSER_IDS`). Returns a dict whose `chunk_size_512` entry is a `MermaidAwareLineSplitter` -- a `TextSplitter` subclass that delegates to `chunk_with_atomic_mermaid` and stamps `has_mermaid` metadata on every leaf. The 512-token target is fixed by spec; the splitter still grows mermaid-bearing chunks up to `embed_max_tokens`. Callers should pair this with `HierarchicalNodeParser(node_parser_ids=EXCEL_PARSER_IDS, ...)` and apply `prune_header_only_nodes` to the result. """ from llama_index.core.node_parser import MarkdownNodeParser return { "markdown": MarkdownNodeParser(), "chunk_size_512": MermaidAwareLineSplitter( tokenize_fn=tokenize_fn, target_tokens=512, max_tokens=embed_max_tokens, ), } def _make_mermaid_aware_splitter_class(): """Defer class definition to first use so importing this module does not require llama_index at module-load time (helps `@task.external_python` environments that load the helper as a string before resolving deps). """ from llama_index.core.bridge.pydantic import PrivateAttr from llama_index.core.node_parser import TextSplitter class _MermaidAwareLineSplitter(TextSplitter): """TextSplitter wrapping `chunk_with_atomic_mermaid` for Excel chunking. Field semantics: target_tokens: pure-prose chunk budget. max_tokens: hard ceiling -- mermaid-bearing chunks may grow to here. The `tokenize_fn` callable is held in a PrivateAttr so it bypasses pydantic schema validation; it is not serialized. """ target_tokens: int = 512 max_tokens: int = 8191 _tokenize_fn: Callable[[str], int] = PrivateAttr() def __init__( self, *, tokenize_fn: Callable[[str], int], target_tokens: int = 512, max_tokens: int = 8191, **kwargs: Any, ) -> None: super().__init__( target_tokens=target_tokens, max_tokens=max_tokens, **kwargs ) self._tokenize_fn = tokenize_fn @classmethod def class_name(cls) -> str: return "MermaidAwareLineSplitter" def split_text(self, text: str) -> List[str]: return [ t for t, _ in chunk_with_atomic_mermaid( text, self.target_tokens, self.max_tokens, self._tokenize_fn ) ] def _parse_nodes( self, nodes: Sequence[Any], show_progress: bool = False, **kwargs: Any, ) -> List[Any]: out = super()._parse_nodes(nodes, show_progress=show_progress, **kwargs) for leaf in out: leaf.metadata["has_mermaid"] = "```mermaid" in leaf.text return out return _MermaidAwareLineSplitter class _LazySplitterProxy: """Lazy-resolves to the real `_MermaidAwareLineSplitter` on first call. Avoids importing llama_index at module-load time when the helper is inlined into `@task.external_python` source. """ _cls = None def __call__(self, *args: Any, **kwargs: Any) -> Any: if _LazySplitterProxy._cls is None: _LazySplitterProxy._cls = _make_mermaid_aware_splitter_class() return _LazySplitterProxy._cls(*args, **kwargs) MermaidAwareLineSplitter = _LazySplitterProxy()