wehub-resource-sync
c56bef871b
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
383 lines
17 KiB
Python
383 lines
17 KiB
Python
# 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}
|