Studio: layout-aware RAG parsers + heading-aware chunking (Phase 3A)
Replace bare-pypdf/python-docx/BeautifulSoup extraction with Markdown-
preserving parsers so the chunker can split on real heading boundaries
instead of running paragraphs together.
Parsers
- pdf.py: pymupdf + pymupdf4llm.to_markdown() per page; pypdf kept as
fallback when pymupdf can't open the file.
- docx.py: mammoth.convert_to_html() + markdownify, with an explicit
style_map so Title/Heading 1..6 become h1..h6 in the output.
- html.py: BeautifulSoup pre-scrub (drop script/style) then markdownify
so <h*>, <table>, <ul> convert faithfully.
- text.py: signature update only; TXT/MD pass through unchanged.
- parsers/__init__.py: new ParsedImage + ParseResult dataclass; parse()
signature is now parse(path, *, want_images=False) -> ParseResult.
ParseResult is iterable over .pages for backward compat.
Chunker
- chunking.py: prepend Markdown heading separators ("\n# " .. "\n#### ")
to the priority list so heading-aware splits happen for free once the
parsers emit Markdown.
Ingestion
- ingestion.py: single call site updated to consume ParseResult.pages.
Deps (no-torch-runtime.txt)
+ pymupdf>=1.24, pymupdf4llm>=0.0.17, mammoth>=1.7, markdownify>=0.13
- pypdf kept as a fallback path.
Tests
- test_rag_parsers.py asserts Markdown headings survive PDF/DOCX/HTML
extraction; also exercises ParseResult iteration backward-compat.
- test_rag_chunking.py: new case verifying chunks start at Markdown
heading boundaries when the input is Markdown.
Foundation for Phase 3B-late (heading-aware spans for late chunking)
and Phase 3B-multimodal (want_images=True enables image extraction in
the same parser layer). No schema or opt-in flags in this commit.
This commit is contained in:
parent
92994e8b83
commit
c4b5889e53
10 changed files with 478 additions and 67 deletions
|
|
@ -82,3 +82,32 @@ def test_chunk_pages_overlap_produces_repeated_tokens():
|
|||
second_head_words = set(chunks[1].text.split()[:4])
|
||||
# At least one word should appear in both
|
||||
assert first_tail_words & second_head_words
|
||||
|
||||
|
||||
def test_chunk_pages_splits_on_markdown_headings():
|
||||
# Phase 3A: heading separators take priority over paragraph breaks
|
||||
# so chunks start at section boundaries when the parser emits
|
||||
# Markdown.
|
||||
md = (
|
||||
"# First Section\n\n"
|
||||
+ "alpha " * 30
|
||||
+ "\n\n## Subsection A\n\n"
|
||||
+ "beta " * 30
|
||||
+ "\n\n# Second Section\n\n"
|
||||
+ "gamma " * 30
|
||||
)
|
||||
chunks = chunk_pages(
|
||||
[ParsedPage(text = md)],
|
||||
max_tokens = 25,
|
||||
overlap_tokens = 0,
|
||||
token_counter = _wc_counter,
|
||||
)
|
||||
# We expect multiple chunks and at least one to begin at a heading.
|
||||
assert len(chunks) >= 2
|
||||
starts_at_heading = sum(
|
||||
1 for c in chunks if c.text.lstrip().startswith(("# ", "## "))
|
||||
)
|
||||
assert starts_at_heading >= 1, (
|
||||
f"expected at least one chunk to start at a Markdown heading; "
|
||||
f"got starts: {[c.text[:20] for c in chunks]}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,4 +1,8 @@
|
|||
"""Document parser tests — each format skipped if its lib is unavailable."""
|
||||
"""Document parser tests — each format skipped if its lib is unavailable.
|
||||
|
||||
Phase 3A: parsers now return ParseResult (iterable over .pages) and
|
||||
emit Markdown so the chunker can split on heading boundaries.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
|
@ -16,19 +20,22 @@ def test_text_parser_utf8(tmp_path):
|
|||
|
||||
file = tmp_path / "sample.txt"
|
||||
file.write_text("hello world\n\nsecond paragraph", encoding = "utf-8")
|
||||
pages = parse(file)
|
||||
assert len(pages) == 1
|
||||
assert "hello world" in pages[0].text
|
||||
assert "second paragraph" in pages[0].text
|
||||
result = parse(file)
|
||||
assert len(result) == 1
|
||||
assert "hello world" in result.pages[0].text
|
||||
assert "second paragraph" in result.pages[0].text
|
||||
assert result.images == []
|
||||
|
||||
|
||||
def test_markdown_parser_treated_as_text(tmp_path):
|
||||
def test_markdown_parser_preserves_headings(tmp_path):
|
||||
from core.rag.parsers import parse
|
||||
|
||||
file = tmp_path / "sample.md"
|
||||
file.write_text("# Title\n\nBody text with **emphasis**.", encoding = "utf-8")
|
||||
pages = parse(file)
|
||||
assert pages and "Title" in pages[0].text
|
||||
result = parse(file)
|
||||
assert result.pages
|
||||
# Markdown should pass through unchanged — heading marker preserved.
|
||||
assert "# Title" in result.pages[0].text
|
||||
|
||||
|
||||
def test_unsupported_format_raises(tmp_path):
|
||||
|
|
@ -40,23 +47,41 @@ def test_unsupported_format_raises(tmp_path):
|
|||
parse(file)
|
||||
|
||||
|
||||
def test_html_parser_strips_scripts(tmp_path):
|
||||
def test_html_parser_emits_markdown_headings(tmp_path):
|
||||
pytest.importorskip("bs4")
|
||||
pytest.importorskip("lxml")
|
||||
pytest.importorskip("markdownify")
|
||||
from core.rag.parsers import parse
|
||||
|
||||
file = tmp_path / "sample.html"
|
||||
file.write_text(
|
||||
"<html><body><script>alert(1)</script><p>visible text</p></body></html>",
|
||||
"<html><body>"
|
||||
"<script>alert(1)</script>"
|
||||
"<h1>Main Title</h1>"
|
||||
"<h2>Sub Section</h2>"
|
||||
"<p>visible text</p>"
|
||||
"<ul><li>one</li><li>two</li></ul>"
|
||||
"</body></html>",
|
||||
encoding = "utf-8",
|
||||
)
|
||||
pages = parse(file)
|
||||
assert pages
|
||||
assert "visible text" in pages[0].text
|
||||
assert "alert" not in pages[0].text
|
||||
result = parse(file)
|
||||
assert result.pages
|
||||
md = result.pages[0].text
|
||||
# markdownify converts <h1> → '# ', <h2> → '## '
|
||||
assert "# Main Title" in md
|
||||
assert "## Sub Section" in md
|
||||
assert "visible text" in md
|
||||
# script content scrubbed
|
||||
assert "alert" not in md
|
||||
# list items become Markdown bullets
|
||||
assert "one" in md and "two" in md
|
||||
|
||||
|
||||
def test_pdf_parser_extracts_pages(tmp_path):
|
||||
pytest.importorskip("pymupdf")
|
||||
pytest.importorskip("pymupdf4llm")
|
||||
from core.rag.parsers import parse
|
||||
|
||||
pypdf = pytest.importorskip("pypdf")
|
||||
from pypdf import PdfWriter
|
||||
|
||||
|
|
@ -65,25 +90,46 @@ def test_pdf_parser_extracts_pages(tmp_path):
|
|||
writer.add_blank_page(width = 72, height = 72)
|
||||
with open(file, "wb") as f:
|
||||
writer.write(f)
|
||||
from core.rag.parsers import parse
|
||||
|
||||
# blank page yields no extractable text — should return [] without error
|
||||
pages = parse(file)
|
||||
assert isinstance(pages, list)
|
||||
# Blank page yields no extractable text — should return empty pages
|
||||
# without error.
|
||||
result = parse(file)
|
||||
assert isinstance(result.pages, list)
|
||||
assert isinstance(result.images, list)
|
||||
|
||||
|
||||
def test_docx_parser_extracts_paragraphs(tmp_path):
|
||||
docx = pytest.importorskip("docx")
|
||||
def test_docx_parser_emits_markdown_headings(tmp_path):
|
||||
pytest.importorskip("docx")
|
||||
pytest.importorskip("mammoth")
|
||||
pytest.importorskip("markdownify")
|
||||
from docx import Document
|
||||
|
||||
file = tmp_path / "sample.docx"
|
||||
doc = Document()
|
||||
doc.add_heading("Top Level Heading", level = 1)
|
||||
doc.add_paragraph("First paragraph here.")
|
||||
doc.add_heading("Sub Heading", level = 2)
|
||||
doc.add_paragraph("Second paragraph here.")
|
||||
doc.save(str(file))
|
||||
from core.rag.parsers import parse
|
||||
|
||||
pages = parse(file)
|
||||
assert pages
|
||||
assert "First paragraph" in pages[0].text
|
||||
assert "Second paragraph" in pages[0].text
|
||||
result = parse(file)
|
||||
assert result.pages
|
||||
md = result.pages[0].text
|
||||
# mammoth via _STYLE_MAP maps Heading 1/2 → h1/h2 → '# '/'## '.
|
||||
assert "# Top Level Heading" in md
|
||||
assert "## Sub Heading" in md
|
||||
assert "First paragraph" in md
|
||||
assert "Second paragraph" in md
|
||||
|
||||
|
||||
def test_parse_result_is_iterable_for_backcompat(tmp_path):
|
||||
"""Code that does `for page in parse(path)` should keep working."""
|
||||
from core.rag.parsers import parse
|
||||
|
||||
file = tmp_path / "sample.txt"
|
||||
file.write_text("hello", encoding = "utf-8")
|
||||
result = parse(file)
|
||||
pages = list(result)
|
||||
assert len(pages) == 1
|
||||
assert pages[0].text == "hello"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue