brain-of-reese/tests/unit/test_chunker.py

"""Unit tests: format-aware chunker (PLAN §5 policy, A9 formats).

The markdown policy tests are the original contract (md output stays
unchanged); the per-format tests cover the phase-09 dispatcher
(yaml/yml, json, py, txt) and the 1200-char hard cap for every format.
"""
from __future__ import annotations

from itertools import pairwise
from pathlib import Path

import pytest

from app.rag.chunker import (
    HARD_MAX_CHARS,
    chunk_document,
    chunk_json,
    chunk_markdown,
    chunk_python,
    chunk_text,
    chunk_yaml,
    extract_title,
)

ANCHOR = "## Big"
ANCHOR_PREFIX = f"{ANCHOR}\n\n"


def _paras(n: int, char: str = "l", width: int = 300) -> list[str]:
    return [f"paragraph {i} " + char * (width - 12) for i in range(n)]


def test_short_document_is_single_chunk() -> None:
    doc = "# Title\n\nJust some intro, no section headings at all."
    chunks = chunk_markdown(doc)
    assert chunks == [doc.strip()]


def test_empty_and_whitespace_only_content() -> None:
    assert chunk_markdown("") == []
    assert chunk_markdown("   \n\n  \n") == []


def test_invalid_params_raise() -> None:
    with pytest.raises(ValueError):
        chunk_markdown("# x", target_chars=0)
    with pytest.raises(ValueError):
        chunk_markdown("# x", overlap_chars=-1)


def test_splits_on_headings_and_keeps_nearest_heading() -> None:
    doc = (
        "# Title\n"
        "intro line\n"
        "## Alpha\n"
        "alpha body\n"
        "### Beta\n"
        "beta body\n"
        "## Gamma\n"
        "gamma body\n"
    )
    chunks = chunk_markdown(doc)
    assert chunks[0] == "# Title\nintro line"
    assert chunks[1] == "## Alpha\nalpha body"
    assert chunks[2] == "### Beta\nbeta body"
    assert chunks[3] == "## Gamma\ngamma body"


def test_document_without_h1_starts_at_first_section() -> None:
    chunks = chunk_markdown("## Only\n\nbody")
    assert chunks == ["## Only\n\nbody"]


def test_long_section_splits_with_overlap_and_anchor_on_every_chunk() -> None:
    body = "\n\n".join(_paras(10, width=138))
    doc = f"{ANCHOR}\n\n{body}"
    chunks = chunk_markdown(doc, target_chars=800, overlap_chars=100)

    assert len(chunks) == 3
    # Every chunk keeps its nearest preceding heading (the section anchor).
    assert all(c.startswith(ANCHOR) for c in chunks)
    # All chunks respect the target budget (anchor + packed body).
    assert all(len(c) <= 800 for c in chunks)
    # Overlap: the tail of each chunk is at the start of the next one.
    for prev, nxt in pairwise(chunks):
        assert nxt[len(ANCHOR_PREFIX) :].startswith(prev[-100:])


def test_overlap_zero_disables_tail_carryover() -> None:
    body = "\n\n".join(_paras(8, width=200))
    chunks = chunk_markdown(f"{ANCHOR}\n\n{body}", target_chars=800, overlap_chars=0)
    assert len(chunks) >= 2
    for prev, nxt in pairwise(chunks):
        assert not nxt[len(ANCHOR_PREFIX) :].startswith(prev[-50:])


def test_code_fences_stay_intact() -> None:
    doc = (
        "## Section\n"
        "before fence\n"
        "```\n"
        "## fake heading inside fence\n"
        "\n"
        "still in fence\n"
        "```\n"
        "after fence\n"
        "## Other\n"
        "other body\n"
    )
    chunks = chunk_markdown(doc)
    assert any(c.startswith("## Other") for c in chunks)
    # The fake heading inside the fence never opens a section…
    assert not any(c.startswith("## fake heading") for c in chunks)
    # …and the fence itself is whole in the chunk that contains it.
    fenced = [c for c in chunks if "still in fence" in c]
    assert len(fenced) == 1
    assert "## fake heading inside fence" in fenced[0]
    assert fenced[0].count("```") == 2
    # Blank lines inside the fence did not create extra paragraph chunks.
    assert not any(c.startswith("before fence\n\n") for c in chunks)


def test_fence_block_is_atomic_across_forced_split() -> None:
    fence = "```\n" + "\n".join(f"code line {i}" for i in range(60)) + "\n```"
    doc = (
        f"{ANCHOR}\n\npara A "
        + "a" * 300
        + f"\n\n{fence}\n\npara B "
        + "b" * 300
        + "\n\npara C "
        + "c" * 300
    )
    chunks = chunk_markdown(doc, target_chars=1000, overlap_chars=100)
    assert len(chunks) >= 2
    # The whole fence (first and last code line) lives in one chunk — a
    # chunk boundary never falls inside a code block.
    assert any("code line 0" in c and "code line 59" in c for c in chunks)


def test_oversized_fence_block_is_split_to_stay_under_hard_cap() -> None:
    """aipi's embedding endpoint caps requests at ~1024 input tokens — a
    multi-KB fenced code block must not survive chunking as one piece."""
    code = "\n".join(f"int value_{i:03d} = {i}; // padding to grow the line" for i in range(160))
    doc = (
        "# Big Doc\n\n"
        "## Usage Example\n\n"
        f"```cpp\n{code}\n```\n\n"
        "## After\n\nDone.\n"
    )
    chunks = chunk_markdown(doc)
    assert len(chunks) >= 3
    # No chunk exceeds the hard cap (heading anchor adds a little).
    assert all(len(c) <= HARD_MAX_CHARS + 60 for c in chunks)
    # Content survives the split, and later sections are untouched.
    joined = "\n".join(chunks)
    assert "value_000" in joined
    assert "value_159" in joined
    assert any(c.startswith("## After") for c in chunks)


def test_unclosed_fence_does_not_break_sections() -> None:
    doc = "## A\n\n```\nunterminated fence\n\n## B\n\nbody\n"
    chunks = chunk_markdown(doc)
    # "## B" is inside the unterminated fence → not a real heading.
    assert len(chunks) == 1
    assert "## B" in chunks[0]


def test_extract_title_prefers_h1() -> None:
    assert extract_title("# My Title\n\nbody") == "My Title"
    assert extract_title("  # Indented H1\nbody") == ""  # ATX must be at col 0
    assert extract_title("## not a title\n\nbody") == ""
    assert extract_title("## sub only", fallback="stem") == "stem"
    assert extract_title("", fallback="fallback") == "fallback"


# ---------------------------------------------------------------------------
# Format dispatcher (chunk_document) — A9 multi-format ingestion
# ---------------------------------------------------------------------------


def test_dispatch_by_lowercased_suffix() -> None:
    md = "# T\n\n## A\n\nbody\n"
    assert chunk_document(md, "notes/Doc.MD") == chunk_markdown(md)
    assert chunk_document(md, "notes/doc.MARKDOWN") == chunk_markdown(md)
    assert chunk_document("p1\n\np2\n", "x.TXT") == chunk_text("p1\n\np2\n")
    assert chunk_document("a: 1\n", "x.YAML") == chunk_yaml("a: 1\n")
    assert chunk_document("a: 1\n", "x.Yml") == chunk_yaml("a: 1\n")
    assert chunk_document('{"a": 1}', "x.Json") == chunk_json('{"a": 1}')
    assert chunk_document("def f(): pass\n", "x.PY") == chunk_python("def f(): pass\n")


def test_dispatch_unknown_suffix_falls_back_to_paragraphs() -> None:
    assert chunk_document("hello\n\nworld", "data.csv") == ["hello\nworld"]


def test_dispatch_ignores_directory_part_of_path() -> None:
    assert chunk_document("def f(): pass\n", "a/b/c/script.py") == chunk_python("def f(): pass\n")


# ---------------------------------------------------------------------------
# yaml / yml
# ---------------------------------------------------------------------------


def test_yaml_splits_on_top_level_keys_and_keeps_key_anchors() -> None:
    doc = (
        "# leading comment\n"
        "services:\n"
        "  gitlab:\n"
        "    image: gitlab/gitlab-ce\n"
        "  prometheus:\n"
        "    image: prom/prometheus\n"
        "volumes:\n"
        "  gitlab-data:\n"
    )
    chunks = chunk_yaml(doc)
    joined = "\n".join(chunks)
    for key in ("services:", "volumes:"):
        assert key in joined
    # Indented keys are NOT block starts — they stay inside their parent block.
    assert not any(c.startswith("  gitlab:") for c in chunks)
    # The leading comment stays with the first block (preamble).
    assert chunks[0].startswith("# leading comment")
    assert "gitlab/gitlab-ce" in joined and "prom/prometheus" in joined


def test_yaml_document_separators_start_new_blocks() -> None:
    a = "site_a: " + "a" * 500 + "\n"
    b = "site_b: " + "b" * 500 + "\n"
    chunks = chunk_yaml(a + "---\n" + b, target_chars=600, overlap_chars=0)
    # Each site is long enough to force its own chunk; the separator must not
    # glue them into one over-budget chunk.
    assert len(chunks) >= 2
    assert all(len(c) <= 600 for c in chunks)
    assert not any("site_a" in c and "site_b" in c for c in chunks)


def test_yaml_oversized_key_block_is_split_under_hard_cap() -> None:
    doc = "big_list:\n" + ("  - " + "x" * 60 + "\n") * 60  # one ~3800-char block
    chunks = chunk_yaml(doc)
    assert len(chunks) >= 2
    assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
    # Overlap re-prints (≤50 chars per split), so only a little content is
    # re-stated — the bulk of the block must survive.
    assert sum(len(c) for c in chunks) >= len(doc) - 300


def test_yaml_empty_content() -> None:
    assert chunk_yaml("") == []
    assert chunk_yaml("\n\n  \n") == []


# ---------------------------------------------------------------------------
# json
# ---------------------------------------------------------------------------


def test_json_splits_on_top_level_keys_pretty_printed() -> None:
    doc = '{"hosts": {"kafkabridge": "10.0.3.7"}, "count": 3}'
    chunks = chunk_json(doc, target_chars=45, overlap_chars=0)  # force 1 chunk/block
    assert len(chunks) == 2
    first, second = chunks
    assert '"hosts"' in first and "kafkabridge" in first
    assert '"count"' in second
    # Pretty-printed (indent=2), not the compact input form.
    assert '"kafkabridge": "10.0.3.7"' in first
    assert not any('{"hosts"' in c for c in chunks)


def test_json_each_key_block_is_self_contained() -> None:
    doc = '{"a": "x", "b": "y"}'
    chunks = chunk_json(doc, target_chars=13, overlap_chars=0)  # force 1 chunk/block
    assert [c for c in chunks if '"a"' in c] and [c for c in chunks if '"b"' in c]
    assert not any('"a"' in c and '"b"' in c for c in chunks)


def test_json_oversized_value_falls_under_hard_cap() -> None:
    doc = '{"blob": "' + "z" * 4000 + '"}'
    chunks = chunk_json(doc)
    assert len(chunks) >= 2
    assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
    assert "".join(chunks).count("z") >= 4000


def test_json_top_level_list_is_one_pretty_block() -> None:
    chunks = chunk_json("[1, 2, 3]")
    assert chunks == ["[\n  1,\n  2,\n  3\n]"]


def test_json_unparseable_falls_back_to_paragraph_packing() -> None:
    doc = "{broken json\n\nsecond paragraph here\n"
    assert chunk_json(doc) == chunk_text(doc)
    assert chunk_json("not json at all") == chunk_text("not json at all")


# ---------------------------------------------------------------------------
# python
# ---------------------------------------------------------------------------


def test_python_splits_on_top_level_defs_and_classes() -> None:
    doc = (
        '"""Module doc."""\n'
        "import asyncio\n"
        "\n"
        "CONST = 1\n"
        "\n"
        "def alpha():\n"
        "    return 1\n"
        "\n"
        "class Beta:\n"
        "    def run(self):\n"
        "        return 2\n"
    )
    chunks = chunk_python(doc, target_chars=60, overlap_chars=0)  # force 1 chunk/block
    assert len(chunks) == 3
    assert chunks[0].startswith('"""Module doc."""')
    assert "CONST = 1" in chunks[0]  # preamble ends at the first def/class
    assert chunks[1].startswith("def alpha")
    assert chunks[2].startswith("class Beta")
    assert "def run" in chunks[2]  # nested def stays inside the class block


def test_python_decorators_stay_with_their_definition() -> None:
    doc = "@app.get('/x')\ndef handler():\n    return 'x'\n"
    chunks = chunk_python(doc)
    assert chunks[0].startswith("@app.get")


def test_python_oversized_function_falls_back_to_line_packing() -> None:
    doc = "def big():\n" + "\n".join(f"    val_{i:03d} = {i}  # padding" for i in range(80))
    chunks = chunk_python(doc)
    assert len(chunks) >= 2
    assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
    assert "val_000" in chunks[0]
    assert "val_079" in chunks[-1]
    assert sum(len(c) for c in chunks) >= len(doc) - 100


def test_python_unparseable_source_falls_back_to_paragraphs() -> None:
    src = "def broken(:\n\nstill text\n"
    assert chunk_python(src) == chunk_text(src)


# ---------------------------------------------------------------------------
# txt
# ---------------------------------------------------------------------------


def test_txt_paragraph_packing() -> None:
    doc = "alpha\n\nbeta\n\ngamma\n"
    chunks = chunk_text(doc)
    assert chunks == ["alpha\nbeta\ngamma"]  # all three fit the target


def test_txt_long_doc_packs_with_overlap() -> None:
    doc = "\n\n".join(f"para {i} " + "l" * 300 for i in range(6))
    chunks = chunk_text(doc, target_chars=800, overlap_chars=100)
    assert len(chunks) >= 2
    assert all(len(c) <= 800 for c in chunks)
    assert all(f"para {i}" in "\n".join(chunks) for i in range(6))


# ---------------------------------------------------------------------------
# quadlet family + j2 (A9 revised 2026-08-27) — plain-text dispatch
# ---------------------------------------------------------------------------

#: The ten suffixes added 2026-08-27 (owner decision R1): the full Podman
#: quadlet family + Jinja templates, all chunked as plain text.
NEW_FORMAT_SUFFIXES = [
    ".container", ".network", ".volume", ".image", ".pod",
    ".kube", ".swap", ".os", ".endpoint", ".j2",
]

REPO_ROOT = Path(__file__).resolve().parents[2]
FIXTURE_DOCS = REPO_ROOT / "tests" / "fixtures" / "docs" / "homelab"


@pytest.mark.parametrize("suffix", NEW_FORMAT_SUFFIXES)
def test_dispatch_new_suffixes_match_plain_text(suffix: str) -> None:
    """Every new suffix dispatches to plain-text paragraph packing — the
    chunks are identical to a direct ``chunk_text`` call."""
    content = "[Unit]\nDescription=one\n\n[Service]\nRestart=always\n\n[Container]\nImage=busybox\n"
    name = suffix.lstrip(".")
    assert chunk_document(content, f"x/{name}{suffix}") == chunk_text(content)


def test_container_fixture_chunks_past_single_paragraph_pack() -> None:
    """The realistic quadlet TOML fixture (≥1 500 chars) splits into ≥2
    chunks, every chunk stays under the hard cap, and the sentinel token
    survives the split."""
    rel = "quadlet/compose.container"
    content = (FIXTURE_DOCS / rel).read_text(encoding="utf-8")
    assert len(content) >= 1500  # the fixture must exercise sub-splitting
    chunks = chunk_document(content, rel)
    assert len(chunks) >= 2
    assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
    assert any("RESE-QUADLET-SENTINEL-77aa" in c for c in chunks)
    joined = "\n".join(chunks)
    assert "[Container]" in joined and "Restart=always" in joined


def test_j2_fixture_chunks_braces_as_plain_text() -> None:
    """Jinja braces are just text — no special handling; the ``{{ … }}``
    line and the ``{% set %}`` sentinel appear verbatim in the chunks."""
    rel = "templates/deploy.j2"
    content = (FIXTURE_DOCS / rel).read_text(encoding="utf-8")
    chunks = chunk_document(content, rel)
    assert chunks
    assert all(len(c) <= HARD_MAX_CHARS for c in chunks)
    joined = "\n".join(chunks)
    assert "RESE-JINJA-SENTINEL-33dd" in joined
    assert any("[{{ svc.name }}]" in c for c in chunks)


def test_dispatch_still_falls_back_for_unknown_suffix() -> None:
    content = "alpha\n\nbeta\n"
    assert chunk_document(content, "x/notes.whatever") == chunk_text(content)


# ---------------------------------------------------------------------------
# pdf (phase 126, task 03, LOCKED A2) — the page-aware chunker
# ---------------------------------------------------------------------------


def test_pdf_single_page_chunks_with_page_anchor() -> None:
    content = "---PAGE-1---\nHello single page.\n"
    chunks = chunk_document(content, "report.pdf")
    assert chunks == ["Page 1:\nHello single page."]
    assert "---PAGE-" not in "\n".join(chunks)


def test_pdf_multi_page_keeps_the_marker_page_number() -> None:
    """Blank pages carry no marker — the anchor follows the MARKER's
    number, not the chunk count (page 3 stays page 3 when page 2 was
    blank)."""
    content = (
        "---PAGE-1---\nAlpha text one.\n\n"
        "---PAGE-3---\nGamma text three.\n"
    )
    chunks = chunk_document(content, "report.pdf")
    assert chunks == ["Page 1:\nAlpha text one.", "Page 3:\nGamma text three."]
    assert "Page 2:" not in "\n".join(chunks)


def test_pdf_markers_never_appear_in_chunks() -> None:
    content = (
        "---PAGE-1---\nAlpha text one.\n\n"
        "---PAGE-2---\nBeta text two.\n\n"
        "---PAGE-3---\nGamma text three.\n"
    )
    chunks = chunk_document(content, "deep/dir/report.pdf")
    assert len(chunks) == 3
    for chunk in chunks:
        assert "---PAGE-" not in chunk
    assert chunks[0].startswith("Page 1:")
    assert chunks[1].startswith("Page 2:")
    assert chunks[2].startswith("Page 3:")


def test_pdf_empty_and_markerless_content_yield_no_chunks() -> None:
    assert chunk_document("", "report.pdf") == []
    # No marker at all (the parser always emits one for text — belt and
    # braces): nothing to anchor, no chunks (never plain-text fallback —
    # a PDF chunk must carry its page).
    assert chunk_document("just words, no marker", "report.pdf") == []
    assert chunk_document("---PAGE-1---\n\n---PAGE-2---\n   ", "report.pdf") == []


def test_pdf_long_page_sub_splits_and_every_chunk_keeps_the_anchor() -> None:
    """A page larger than the target sub-splits through the paragraph
    packer — EVERY chunk of the page keeps the ``Page N:`` anchor and
    the anchored chunk never exceeds the target."""
    page = "\n\n".join(f"paragraph {i} " + "l" * 300 for i in range(6))
    content = f"---PAGE-1---\n{page}\n\n---PAGE-2---\nshort second page\n"
    chunks = chunk_document(content, "report.pdf", target_chars=800, overlap_chars=100)
    assert len(chunks) >= 3  # the long page sub-splits + the short page
    page1 = [c for c in chunks if c.startswith("Page 1:")]
    page2 = [c for c in chunks if c.startswith("Page 2:")]
    assert len(page1) >= 2 and page2 == ["Page 2:\nshort second page"]
    for chunk in chunks:
        assert len(chunk) <= 800, f"target exceeded: {len(chunk)}"
    # All six paragraphs survive the split.
    joined = "\n".join(chunks)
    for i in range(6):
        assert f"paragraph {i}" in joined


def test_pdf_overlap_carryover_stays_within_the_page() -> None:
    """The tail carry-over packs paragraphs INSIDE one page; a new page
    never inherits the previous page's tail (the page is the boundary)."""
    a = "\n\n".join(f"para-a-{i} " + "l" * 250 for i in range(3))
    b = "\n\n".join(f"para-b-{i} " + "l" * 250 for i in range(3))
    content = f"---PAGE-1---\n{a}\n\n---PAGE-2---\n{b}\n"
    chunks = chunk_document(content, "r.pdf", target_chars=700, overlap_chars=120)
    first_of_page2 = next(c for c in chunks if c.startswith("Page 2:"))
    assert "para-a-" not in first_of_page2, "page 2 must not inherit page 1's tail"
    assert all(len(c) <= 700 for c in chunks)


def test_pdf_dispatch_matches_direct_call() -> None:
    content = "---PAGE-1---\nalpha\n\n---PAGE-2---\nbeta\n"
    from app.rag.chunker import _chunk_pdf

    assert chunk_document(content, "x/report.PDF") == _chunk_pdf(content)


# ---------------------------------------------------------------------------
# Hard cap across every format (aipi ~1024-token request cap)
# ---------------------------------------------------------------------------


@pytest.mark.parametrize(
    ("content", "path"),
    [
        ("# T\n\n" + "word " * 1200, "big.md"),
        ("key: " + "v" * 5000 + "\n", "big.yaml"),
        ('{"blob": "' + "z" * 5000 + '"}', "big.json"),
        ("def f():\n" + "    x = 1\n" * 1000, "big.py"),
        ("line of text\n\n" * 800, "big.txt"),
        ("[Unit]\n" + "key=" + "v" * 5000 + "\n", "big.container"),  # phase 47
        ("{% for x in y %}" + "{{ x }} " * 400 + "{% endfor %}", "big.j2"),  # phase 47
        # phase 126: one long page (a single 12 500-char line) + one short.
        ("---PAGE-1---\n" + "word " * 2500 + "\n\n---PAGE-2---\nshort\n", "big.pdf"),
    ],
)
def test_hard_cap_holds_for_every_format(content: str, path: str) -> None:
    chunks = chunk_document(content, path)
    assert chunks, "expected at least one chunk"
    for c in chunks:
        assert len(c) <= HARD_MAX_CHARS, f"{path}: {len(c)} chars"