210 lines
8.3 KiB
Python
210 lines
8.3 KiB
Python
"""Chunker policy: heading sections, breadcrumbs, sibling merge, block-aligned split."""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
SCRIPTS = Path(__file__).parent.parent / "scripts"
|
|
sys.path.insert(0, str(SCRIPTS))
|
|
|
|
from wiki_chunker import MERGE_BELOW, OVERLAP, SPLIT_ABOVE, parse_markdown, token_estimate
|
|
|
|
FILLER = "Tohle je odstavec s dostatkem textu na to, aby se do velikosti chunku počítal. "
|
|
|
|
|
|
def _paragraph(tokens: int) -> str:
|
|
return (FILLER * (1 + tokens * 4 // len(FILLER)))[: tokens * 4]
|
|
|
|
|
|
def _crumbs(parsed):
|
|
return [chunk.breadcrumb for chunk in parsed.chunks]
|
|
|
|
|
|
def test_heading_hierarchy_builds_breadcrumbs():
|
|
doc = "\n".join(
|
|
[
|
|
"# Kořen",
|
|
"",
|
|
_paragraph(MERGE_BELOW),
|
|
"",
|
|
"## Sekce A",
|
|
"",
|
|
_paragraph(MERGE_BELOW),
|
|
"",
|
|
"### Podsekce",
|
|
"",
|
|
_paragraph(MERGE_BELOW),
|
|
"",
|
|
"## Sekce B",
|
|
"",
|
|
_paragraph(MERGE_BELOW),
|
|
]
|
|
)
|
|
parsed = parse_markdown(doc, "notes/doc.md")
|
|
assert _crumbs(parsed) == [
|
|
"notes/doc.md > Kořen",
|
|
"notes/doc.md > Kořen > Sekce A",
|
|
"notes/doc.md > Kořen > Sekce A > Podsekce",
|
|
"notes/doc.md > Kořen > Sekce B",
|
|
]
|
|
assert parsed.headings == ["Kořen", "Sekce A", "Podsekce", "Sekce B"]
|
|
|
|
|
|
def test_frontmatter_title_and_tags():
|
|
doc = "---\ntitle: Tokyo metro tips\ntags: [transit, jr-pass]\n---\n\n" + _paragraph(MERGE_BELOW)
|
|
parsed = parse_markdown(doc, "japan/metro.md")
|
|
assert parsed.title == "Tokyo metro tips"
|
|
assert parsed.tags == ["transit", "jr-pass"]
|
|
# Only the title reaches the embedded text; tags stay metadata for filtering.
|
|
assert _crumbs(parsed) == ["japan/metro.md > Tokyo metro tips"]
|
|
assert "transit" not in parsed.chunks[0].text
|
|
|
|
|
|
def test_frontmatter_title_matching_h1_is_not_duplicated():
|
|
doc = "---\ntitle: Zálohy\n---\n\n# Zálohy\n\n" + _paragraph(MERGE_BELOW)
|
|
parsed = parse_markdown(doc, "a.md")
|
|
assert _crumbs(parsed) == ["a.md > Zálohy"]
|
|
|
|
|
|
def test_comma_separated_tags():
|
|
doc = "---\ntags: gear, safety\n---\n\ntext"
|
|
assert parse_markdown(doc, "a.md").tags == ["gear", "safety"]
|
|
|
|
|
|
def test_malformed_frontmatter_stays_body():
|
|
"""Unparseable frontmatter is not consumed — it stays body text, tags included."""
|
|
doc = "---\ntitle: [unclosed\ntags: [a\n---\n\ntext"
|
|
parsed = parse_markdown(doc, "a.md")
|
|
assert parsed.tags == []
|
|
assert "title: [unclosed" in parsed.chunks[-1].text # the raw block survived as prose
|
|
|
|
|
|
def test_small_adjacent_siblings_merge():
|
|
doc = "\n".join(["# Root", "", "## A", "", "krátké A", "", "## B", "", "krátké B"])
|
|
parsed = parse_markdown(doc, "a.md")
|
|
# A and B share the parent `a.md > Root` and are both under MERGE_BELOW.
|
|
# `# Root` carries no prose of its own, so it contributes no chunk.
|
|
assert _crumbs(parsed) == ["a.md > Root > A"]
|
|
assert "krátké A" in parsed.chunks[0].text
|
|
assert "krátké B" in parsed.chunks[0].text
|
|
|
|
|
|
def test_heading_without_own_prose_yields_no_chunk():
|
|
"""`# Titul` immediately followed by `## Sekce` must not embed a bare title."""
|
|
doc = "\n".join(["# Titul", "", "## Sekce", "", _paragraph(MERGE_BELOW)])
|
|
parsed = parse_markdown(doc, "a.md")
|
|
assert _crumbs(parsed) == ["a.md > Titul > Sekce"]
|
|
# The dropped heading still reaches the index through the breadcrumb and headings list.
|
|
assert parsed.headings == ["Titul", "Sekce"]
|
|
|
|
|
|
def test_merge_does_not_cross_parents():
|
|
"""A small H3 must not be glued onto the next H2 — different parents."""
|
|
doc = "\n".join(["# Root", "", "## A", "", "### A1", "", "malé", "", "## B", "", "malé B"])
|
|
parsed = parse_markdown(doc, "a.md")
|
|
assert "a.md > Root > A > A1" in _crumbs(parsed)
|
|
a1 = next(c for c in parsed.chunks if c.breadcrumb.endswith("A1"))
|
|
assert "malé B" not in a1.text
|
|
|
|
|
|
def test_merge_never_exceeds_split_threshold():
|
|
big = _paragraph(SPLIT_ABOVE - 100)
|
|
doc = "\n".join(["# Root", "", "## A", "", "malé", "", "## B", "", big])
|
|
parsed = parse_markdown(doc, "a.md")
|
|
for chunk in parsed.chunks:
|
|
assert token_estimate(chunk.text) <= SPLIT_ABOVE + token_estimate(chunk.breadcrumb) + 1
|
|
|
|
|
|
def test_h4_stays_inside_its_h3_section():
|
|
doc = "\n".join(["# Root", "", "### Trojka", "", "text", "", "#### Čtyřka", "", "hluboký text"])
|
|
parsed = parse_markdown(doc, "a.md")
|
|
assert "a.md > Root > Trojka" in _crumbs(parsed)
|
|
assert not any("Čtyřka" in crumb for crumb in _crumbs(parsed))
|
|
section = next(c for c in parsed.chunks if c.breadcrumb.endswith("Trojka"))
|
|
assert "hluboký text" in section.text
|
|
|
|
|
|
def test_preamble_before_first_heading_is_kept():
|
|
doc = "úvodní odstavec bez nadpisu\n\n" + "# Root\n\n" + _paragraph(MERGE_BELOW)
|
|
parsed = parse_markdown(doc, "a.md")
|
|
assert "úvodní odstavec bez nadpisu" in parsed.chunks[0].text
|
|
|
|
|
|
def test_document_without_headings_is_one_chunk():
|
|
parsed = parse_markdown(_paragraph(MERGE_BELOW), "a.md")
|
|
assert _crumbs(parsed) == ["a.md"]
|
|
|
|
|
|
def test_empty_document_yields_no_chunks():
|
|
assert parse_markdown("", "a.md").chunks == []
|
|
assert parse_markdown(" \n\n \n", "a.md").chunks == []
|
|
|
|
|
|
def test_setext_heading_is_a_section():
|
|
doc = "Nadpis setext\n=============\n\n" + _paragraph(MERGE_BELOW)
|
|
parsed = parse_markdown(doc, "a.md")
|
|
assert _crumbs(parsed) == ["a.md > Nadpis setext"]
|
|
|
|
|
|
def test_oversized_section_splits_and_keeps_code_fence_whole():
|
|
fence = "```python\n" + "\n".join(f"value_{i} = {i} # a comment long enough to count" for i in range(60)) + "\n```"
|
|
doc = "# Velká sekce\n\n" + "\n\n".join([_paragraph(120)] * 6) + "\n\n" + fence + "\n\n" + _paragraph(120)
|
|
parsed = parse_markdown(doc, "big.md")
|
|
|
|
assert len(parsed.chunks) > 1
|
|
assert {c.breadcrumb for c in parsed.chunks} == {"big.md > Velká sekce"}
|
|
# The fence is larger than SPLIT_ABOVE on its own, so it must sit alone and unbroken.
|
|
assert sum(fence in chunk.text for chunk in parsed.chunks) == 1
|
|
|
|
|
|
def test_split_carries_block_aligned_overlap():
|
|
blocks = [_paragraph(40) + f" marker{i:03d}" for i in range(40)]
|
|
doc = "# Root\n\n" + "\n\n".join(blocks)
|
|
parsed = parse_markdown(doc, "a.md")
|
|
|
|
assert len(parsed.chunks) > 1
|
|
# Each boundary repeats at least one whole block, capped at OVERLAP tokens.
|
|
for earlier, later in zip(parsed.chunks, parsed.chunks[1:], strict=False):
|
|
shared = [b for b in blocks if b in earlier.text and b in later.text]
|
|
assert shared, "expected overlap blocks between consecutive chunks"
|
|
assert token_estimate("".join(shared)) <= OVERLAP
|
|
|
|
|
|
def test_table_is_not_broken():
|
|
table = "\n".join(["| a | b |", "|---|---|"] + [f"| {i} | {i * 2} |" for i in range(80)])
|
|
doc = "# Root\n\n" + "\n\n".join([_paragraph(150)] * 5) + "\n\n" + table
|
|
parsed = parse_markdown(doc, "a.md")
|
|
assert sum(table in chunk.text for chunk in parsed.chunks) == 1
|
|
|
|
|
|
def test_breadcrumb_is_prefixed_to_chunk_text():
|
|
parsed = parse_markdown("# Root\n\n" + _paragraph(MERGE_BELOW), "notes/a.md")
|
|
chunk = parsed.chunks[0]
|
|
assert chunk.text.startswith(chunk.breadcrumb + "\n\n")
|
|
|
|
|
|
def test_catalog_title_falls_back_to_first_heading():
|
|
"""Notes carry their title as `# H1` far more often than as frontmatter."""
|
|
doc = "# Zálohování dat\n\n" + _paragraph(MERGE_BELOW) + "\n\n## Retence\n\n" + _paragraph(MERGE_BELOW)
|
|
parsed = parse_markdown(doc, "notes/zalohy.md")
|
|
assert parsed.title == "Zálohování dat"
|
|
|
|
|
|
def test_frontmatter_title_wins_over_first_heading():
|
|
doc = "---\ntitle: Z frontmatteru\n---\n\n# Z nadpisu\n\n" + _paragraph(MERGE_BELOW)
|
|
assert parse_markdown(doc, "a.md").title == "Z frontmatteru"
|
|
|
|
|
|
def test_document_without_headings_has_no_title():
|
|
assert parse_markdown(_paragraph(MERGE_BELOW), "a.md").title is None
|
|
assert parse_markdown("", "a.md").title is None
|
|
|
|
|
|
def test_heading_title_fallback_does_not_change_any_chunk():
|
|
"""The fallback feeds only the catalog — touching the breadcrumb root would force --full."""
|
|
doc = "úvod bez nadpisu\n\n# Root\n\n" + _paragraph(MERGE_BELOW) + "\n\n## Sekce\n\n" + _paragraph(MERGE_BELOW)
|
|
parsed = parse_markdown(doc, "notes/a.md")
|
|
|
|
assert parsed.title == "Root"
|
|
# The root stays the bare path, so no chunk text is rewritten by the fallback.
|
|
assert _crumbs(parsed) == ["notes/a.md", "notes/a.md > Root", "notes/a.md > Root > Sekce"]
|