chore: import upstream snapshot with attribution

haystack/components/preprocessors/markdown_header_splitter.py

@@ -0,0 +1,382 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from typing import Literal
+
+from haystack import Document, component, logging
+from haystack.components.preprocessors import DocumentSplitter
+
+logger = logging.getLogger(__name__)
+
+
+@component
+class MarkdownHeaderSplitter:
+    """
+    Split documents at ATX-style Markdown headers (#), with optional secondary splitting.
+
+    This component processes text documents by:
+    - Splitting them into chunks at Markdown headers (e.g., '#', '##', etc.), preserving header hierarchy as metadata.
+    - Optionally applying a secondary split (by word, passage, period, or line) to each chunk
+      (using haystack's DocumentSplitter).
+    - Preserving and propagating metadata such as parent headers, page numbers, and split IDs.
+    """
+
+    def __init__(
+        self,
+        *,
+        page_break_character: str = "\f",
+        keep_headers: bool = True,
+        header_split_levels: list[int] | None = None,
+        secondary_split: Literal["word", "passage", "period", "line"] | None = None,
+        split_length: int = 200,
+        split_overlap: int = 0,
+        split_threshold: int = 0,
+        skip_empty_documents: bool = True,
+    ) -> None:
+        """
+        Initialize the MarkdownHeaderSplitter.
+
+        :param page_break_character: Character used to identify page breaks. Defaults to form feed ("\f").
+        :param keep_headers: If True, headers are kept in the content. If False, headers are moved to metadata.
+            Defaults to True.
+        :param header_split_levels: List of header levels (1–6) to split on. For example, `[1, 2]` splits only
+            on `#` and `##` headers, merging content under deeper headers into the preceding chunk. Defaults to
+            all levels `[1, 2, 3, 4, 5, 6]`.
+        :param secondary_split: Optional secondary split condition after header splitting.
+            Options are None, "word", "passage", "period", "line". Defaults to None.
+        :param split_length: The maximum number of units in each split when using secondary splitting. Defaults to 200.
+        :param split_overlap: The number of overlapping units for each split when using secondary splitting.
+            Defaults to 0.
+        :param split_threshold: The minimum number of units per split when using secondary splitting. Defaults to 0.
+        :param skip_empty_documents: Choose whether to skip documents with empty content. Default is True.
+            Set to False when downstream components in the Pipeline (like LLMDocumentContentExtractor) can extract text
+            from non-textual documents.
+        """
+        if header_split_levels is None:
+            header_split_levels = [1, 2, 3, 4, 5, 6]
+
+        if not isinstance(header_split_levels, list) or len(header_split_levels) == 0:
+            raise ValueError("header_split_levels must be a non-empty list.")
+        invalid = [lvl for lvl in header_split_levels if not isinstance(lvl, int) or lvl < 1 or lvl > 6]
+        if invalid:
+            raise ValueError(
+                f"header_split_levels contains invalid values: {invalid}. All levels must be integers between 1 and 6."
+            )
+        if len(header_split_levels) != len(set(header_split_levels)):
+            raise ValueError("header_split_levels must not contain duplicate values.")
+
+        self.page_break_character = page_break_character
+        self.secondary_split = secondary_split
+        self.split_length = split_length
+        self.split_overlap = split_overlap
+        self.split_threshold = split_threshold
+        self.skip_empty_documents = skip_empty_documents
+        self.keep_headers = keep_headers
+        self.header_split_levels = header_split_levels
+        self._header_split_levels_set = set(header_split_levels)
+        self._header_pattern = re.compile(r"(?m)^(#{1,6}) (.+)$")  # ATX-style .md-headers
+
+        # Matches fenced code blocks delimited by triple backticks (```) or triple tildes (~~~).
+        # Broken down:
+        #   ^                 - fence must start at the beginning of a line (MULTILINE)
+        #   (?P<fence>`{3,}|~{3,})
+        #                     - named capture group "fence": three or more backticks OR three or
+        #                       more tildes. Capturing it allows the closing fence to be matched
+        #                       with a backreference, so ```-opened blocks must close with ```
+        #                       and ~~~-opened blocks must close with ~~~.
+        #   [^\n]*            - optional language identifier (e.g. "python") and any other text
+        #                       on the opening fence line, up to the newline
+        #   \n                - newline ending the opening fence line
+        #   .*?               - the code block body, matched lazily (DOTALL so . matches newlines)
+        #   ^(?P=fence)       - closing fence: must be identical to the opening fence (backreference),
+        #                       and must start at the beginning of a line
+        #   \s*$              - optional trailing whitespace after the closing fence
+        self._code_block_pattern = re.compile(
+            r"^(?P<fence>`{3,}|~{3,})[^\n]*\n.*?^(?P=fence)\s*$", re.MULTILINE | re.DOTALL
+        )
+
+        self._is_warmed_up = False
+
+        # initialize secondary_splitter only if needed
+        if self.secondary_split:
+            self.secondary_splitter = DocumentSplitter(
+                split_by=self.secondary_split,
+                split_length=self.split_length,
+                split_overlap=self.split_overlap,
+                split_threshold=self.split_threshold,
+            )
+
+    def warm_up(self) -> None:
+        """
+        Warm up the MarkdownHeaderSplitter.
+        """
+        if self.secondary_split and not self._is_warmed_up:
+            self.secondary_splitter.warm_up()
+            self._is_warmed_up = True
+
+    def _code_block_spans(self, text: str) -> list[tuple[int, int]]:
+        """Return the (start, end) character spans of all fenced code blocks in text."""
+        return [(m.start(), m.end()) for m in self._code_block_pattern.finditer(text)]
+
+    def _split_text_by_markdown_headers(self, text: str, doc_id: str) -> list[dict]:
+        """Split text by ATX-style headers (#) and create chunks with appropriate metadata."""
+        logger.debug("Splitting text by markdown headers")
+
+        # Pre-compute fenced code block spans so that # lines inside code blocks (e.g. Python comments) are not
+        # mistaken for Markdown headers.
+        code_spans = self._code_block_spans(text)
+
+        # find headers at the configured levels only, excluding any that fall inside a code block. Content between
+        # skipped headers is absorbed into the preceding chunk's span since end = next_match.start().
+        matches = [
+            m
+            for m in re.finditer(self._header_pattern, text)
+            if len(m.group(1)) in self._header_split_levels_set
+            and not any(start <= m.start() < end for start, end in code_spans)
+        ]
+
+        # return unsplit if no headers found
+        if not matches:
+            logger.info(
+                "No headers found in document {doc_id}; returning full document as single chunk.", doc_id=doc_id
+            )
+            return [{"content": text, "meta": {}}]
+
+        # process headers and build chunks
+        chunks: list[dict] = []
+        header_stack: list[str | None] = [None] * 6
+        pending_headers: list[str] = []  # store empty headers to prepend to next content
+        has_content = False  # flag to track if any header has content
+
+        for i, match in enumerate(matches):
+            # extract header info
+            header_prefix = match.group(1)
+            header_text = match.group(2)
+            level = len(header_prefix)
+
+            # get content
+            start = match.end()
+            end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
+            content = text[start:end]
+
+            # update header stack to track nesting
+            header_stack[level - 1] = header_text
+            for j in range(level, 6):
+                header_stack[j] = None
+
+            # skip splits w/o content
+            if not content.strip():  # this strip is needed to avoid counting whitespace as content
+                if self.keep_headers:
+                    header_line = f"{header_prefix} {header_text}"
+                    pending_headers.append(header_line)
+                continue
+
+            has_content = True  # at least one header has content
+            # Build parent metadata from the current header stack so the first child of a
+            # contentful section still inherits its full ancestor chain.
+            parent_headers = [h for h in header_stack[: level - 1] if h is not None]
+
+            logger.debug(
+                "Creating chunk for header '{header_text}' at level {level}", header_text=header_text, level=level
+            )
+
+            if self.keep_headers:
+                header_line = f"{header_prefix} {header_text}"
+                # add pending & current header to content
+                chunk_content = ""
+                if pending_headers:
+                    chunk_content += "\n".join(pending_headers) + "\n"
+                chunk_content += f"{header_line}{content}"
+                chunks.append(
+                    {"content": chunk_content, "meta": {"header": header_text, "parent_headers": parent_headers}}
+                )
+                pending_headers = []  # reset pending headers
+            else:
+                chunks.append({"content": content, "meta": {"header": header_text, "parent_headers": parent_headers}})
+
+        # return doc unchunked if no headers have content
+        if not has_content:
+            logger.info(
+                "Document {doc_id} contains only headers with no content; returning original document.", doc_id=doc_id
+            )
+            return [{"content": text, "meta": {}}]
+
+        return chunks
+
+    def _apply_secondary_splitting(self, documents: list[Document]) -> list[Document]:
+        """
+        Apply secondary splitting while preserving header metadata and structure.
+
+        Ensures page counting is maintained across splits.
+        """
+        result_docs = []
+        current_split_id = 0  # track split_id across all secondary splits from the same parent
+
+        for doc in documents:
+            if doc.content is None:
+                result_docs.append(doc)
+                continue
+
+            content_for_splitting: str = doc.content
+
+            if not self.keep_headers:  # skip header extraction if keep_headers
+                # extract header information
+                header_match = re.match(self._header_pattern, doc.content)
+                if header_match:
+                    content_for_splitting = doc.content[header_match.end() :]
+
+            # track page from meta
+            current_page = doc.meta.get("page_number", 1)
+
+            # create a clean meta dict without split_id for secondary splitting
+            clean_meta = {k: v for k, v in doc.meta.items() if k != "split_id"}
+
+            secondary_splits = self.secondary_splitter.run(
+                documents=[Document(content=content_for_splitting, meta=clean_meta)]
+            )["documents"]
+
+            # split processing
+            for i, split in enumerate(secondary_splits):
+                # calculate page number for this split
+                if i > 0 and secondary_splits[i - 1].content:
+                    current_page = self._update_page_number_with_breaks(secondary_splits[i - 1].content, current_page)
+
+                # set page number and split_id to meta
+                split.meta["page_number"] = current_page
+                split.meta["split_id"] = current_split_id
+                # ensure source_id is preserved from the original document
+                if "source_id" in doc.meta:
+                    split.meta["source_id"] = doc.meta["source_id"]
+                current_split_id += 1
+
+                # preserve header metadata if we're not keeping headers in content
+                if not self.keep_headers:
+                    for key in ["header", "parent_headers"]:
+                        if key in doc.meta:
+                            split.meta[key] = doc.meta[key]
+
+                result_docs.append(split)
+
+        logger.debug(
+            "Secondary splitting complete. Final count: {final_count} documents.", final_count=len(result_docs)
+        )
+        return result_docs
+
+    def _update_page_number_with_breaks(self, content: str | None, current_page: int) -> int:
+        """
+        Update page number based on page breaks in content.
+
+        :param content: Content to check for page breaks
+        :param current_page: Current page number
+        :return: New current page number
+        """
+        if not isinstance(content, str):
+            return current_page
+
+        page_breaks = content.count(self.page_break_character)
+        new_page_number = current_page + page_breaks
+
+        if page_breaks > 0:
+            logger.debug(
+                "Found {page_breaks} page breaks, page number updated: {old} → {new}",
+                page_breaks=page_breaks,
+                old=current_page,
+                new=new_page_number,
+            )
+
+        return new_page_number
+
+    def _split_documents_by_markdown_headers(self, documents: list[Document]) -> list[Document]:
+        """Split a list of documents by markdown headers, preserving metadata."""
+
+        result_docs = []
+        for doc in documents:
+            logger.debug("Splitting document with id={doc_id}", doc_id=doc.id)
+            # mypy: doc.content is Optional[str], so we must check for None before passing to splitting method
+            if doc.content is None:
+                continue
+            splits = self._split_text_by_markdown_headers(doc.content, doc.id)
+            docs = []
+
+            current_page = doc.meta.get("page_number", 1) if doc.meta else 1
+            total_page_breaks = doc.content.count(self.page_break_character)
+            logger.debug(
+                "Processing document with id={doc_id}: starting at page {start_page}, "
+                "contains {page_breaks} page breaks in total",
+                doc_id=doc.id,
+                start_page=current_page,
+                page_breaks=total_page_breaks,
+            )
+            for split_idx, split in enumerate(splits):
+                meta = doc.meta.copy() if doc.meta else {}
+                meta.update({"source_id": doc.id, "page_number": current_page, "split_id": split_idx})
+                if split.get("meta"):
+                    meta.update(split["meta"])
+                current_page = self._update_page_number_with_breaks(split["content"], current_page)
+                docs.append(Document(content=split["content"], meta=meta))
+            logger.debug(
+                "Split into {num_docs} documents for id={doc_id}, final page: {current_page}",
+                num_docs=len(docs),
+                doc_id=doc.id,
+                current_page=current_page,
+            )
+            result_docs.extend(docs)
+        return result_docs
+
+    @component.output_types(documents=list[Document])
+    def run(self, documents: list[Document]) -> dict[str, list[Document]]:
+        """
+        Run the markdown header splitter with optional secondary splitting.
+
+        :param documents: List of documents to split
+
+        :returns: A dictionary with the following key:
+            - `documents`: List of documents with the split texts. Each document includes:
+                - A metadata field `source_id` to track the original document.
+                - A metadata field `page_number` to track the original page number.
+                - A metadata field `split_id` to identify the split chunk index within its parent document.
+                - All other metadata copied from the original document.
+        :raises ValueError: If a document has `None` content.
+        :raises TypeError: If a document's content is not a string.
+        """
+        if self.secondary_split and not self._is_warmed_up:
+            self.warm_up()
+        # validate input documents
+        for doc in documents:
+            if doc.content is None:
+                raise ValueError(
+                    "MarkdownHeaderSplitter only works with text documents but content for document ID"
+                    f" {doc.id} is None."
+                )
+            if not isinstance(doc.content, str):
+                raise TypeError("MarkdownHeaderSplitter only works with text documents (str content).")
+
+        final_docs = []
+        for doc in documents:
+            # handle empty documents
+            if not doc.content or not doc.content.strip():  # avoid counting whitespace as content
+                if self.skip_empty_documents:
+                    logger.warning("Document ID {doc_id} has an empty content. Skipping this document.", doc_id=doc.id)
+                    continue
+                # keep empty documents
+                final_docs.append(doc)
+                logger.warning(
+                    "Document ID {doc_id} has an empty content. Keeping this document as per configuration.",
+                    doc_id=doc.id,
+                )
+                continue
+
+            # split this document by headers
+            header_split_docs = self._split_documents_by_markdown_headers([doc])
+
+            # apply secondary splitting if configured
+            if self.secondary_split:
+                doc_splits = self._apply_secondary_splitting(header_split_docs)
+            else:
+                doc_splits = header_split_docs
+
+            final_docs.extend(doc_splits)
+
+        return {"documents": final_docs}