FabledScribe/tests/test_chunking.py

"""chunk_document — the document shape a record is embedded as (#280).

The model window is 512 tokens and fastembed truncates silently, so before
chunking, everything past ~400 words of a record was PERMANENTLY invisible to
semantic search. These tests pin the two halves of the fix's contract: short
records keep the exact historical shape (the corpus's sharpest vectors are
untouched), and long records lose NOTHING — every line of the body lands in
some chunk, each chunk inside the window budget, each carrying the title as
its topical anchor.
"""
from scribe.services.embeddings import (
    _CHUNK_CHAR_BUDGET,
    chunk_document,
    embedding_text,
)


def _long_section(tag: str, paragraphs: int = 6, sentence: str = None) -> str:
    sentence = sentence or f"This paragraph discusses {tag} in useful detail."
    para = " ".join([sentence] * 6)
    return "\n\n".join(f"{para} (p{i})" for i in range(paragraphs))


# --- the identity half: short records are byte-for-byte unaffected -----------


def test_a_short_record_yields_exactly_the_historical_shape():
    """Snippets and reference notes are the sharpest records in the corpus
    (#2485) precisely because of this shape — chunking must not touch them."""
    assert chunk_document("A title", "A short body") == [
        embedding_text("A title", "A short body")
    ]


def test_a_bodyless_record_is_one_chunk_of_its_title():
    assert chunk_document("Just a title", "") == ["Just a title"]


def test_an_empty_record_yields_no_chunks():
    """Callers gate on falsiness to skip embedding entirely."""
    assert chunk_document("", "") == []
    assert chunk_document(None, None) == []


def test_a_record_exactly_at_budget_stays_whole():
    body = "x" * (_CHUNK_CHAR_BUDGET - len("T\n"))
    assert chunk_document("T", body) == [embedding_text("T", body)]


# --- the no-lost-data half: this is what the build is FOR --------------------


def test_every_line_of_a_long_body_lands_in_some_chunk():
    """The point of #280. Before chunking, a 2,000-word dev-log's last three
    quarters could not influence retrieval at all. Nothing may be dropped."""
    sections = [
        f"## Topic {i}\n\n{_long_section(f'topic-{i}')}" for i in range(8)
    ]
    body = "Intro paragraph before any heading.\n\n" + "\n\n".join(sections)
    chunks = chunk_document("A very long dev-log", body)

    assert len(chunks) > 1
    joined = "\n".join(chunks)
    for line in body.splitlines():
        if line.strip():
            assert line.strip() in joined, f"content dropped: {line[:60]!r}"


def test_every_chunk_fits_the_window_budget():
    body = "\n\n".join(_long_section(f"t{i}") for i in range(10))
    for chunk in chunk_document("T", body):
        assert len(chunk) <= _CHUNK_CHAR_BUDGET + len("T") + 1


def test_every_chunk_is_anchored_by_the_title():
    """Each vector must carry its own topical anchor — the property that makes
    snippets discriminative. An unanchored mid-document chunk would embed as
    free-floating prose about nothing in particular."""
    body = "\n\n".join(
        f"## Section {i}\n\n{_long_section(f'sec-{i}')}" for i in range(6)
    )
    chunks = chunk_document("Retrieval reference", body)
    assert len(chunks) > 1
    for chunk in chunks:
        assert chunk.startswith("Retrieval reference\n")


# --- boundary behaviour ------------------------------------------------------


def test_sections_split_at_markdown_headings_and_stay_whole_when_they_fit():
    a = "## Alpha\n\nShort alpha content."
    b = "## Beta\n\n" + _long_section("beta")
    c = "## Gamma\n\n" + _long_section("gamma")
    chunks = chunk_document("T", f"{a}\n\n{b}\n\n{c}")
    # Beta's content never shares a chunk with Gamma's heading-onward content:
    # heading boundaries are chunk boundaries unless merging small sections.
    for chunk in chunks:
        assert not ("(p5)" in chunk and "## Gamma" in chunk and "beta" in chunk)


def test_small_adjacent_sections_merge_instead_of_each_spending_a_vector():
    body = (
        "\n\n".join(f"## S{i}\n\nTiny." for i in range(4))
        + "\n\n## Big\n\n"
        + _long_section("big", paragraphs=10)
    )
    chunks = chunk_document("T", body)
    tiny_chunks = [c for c in chunks if "Tiny." in c]
    assert len(tiny_chunks) == 1, "four tiny sections should share one chunk"


def test_a_heading_inside_a_code_fence_does_not_split():
    """A commented `# step` in a recorded shell snippet is content, not
    structure."""
    body = "Intro.\n\n```bash\n# not a heading\necho hi\n```\n\nOutro."
    chunks = chunk_document("T", body)
    assert chunks == [embedding_text("T", body)]


def test_pieces_subsplit_from_one_section_repeat_its_heading():
    """'Which part of which topic' must survive the split — a continuation
    piece without its heading embeds as context-free prose."""
    body = "## The Only Topic\n\n" + _long_section("only", paragraphs=40)
    chunks = chunk_document("T", body)
    assert len(chunks) > 1
    for chunk in chunks:
        assert "## The Only Topic" in chunk


def test_a_monster_single_paragraph_is_hard_split_not_dropped():
    body = "word " * 2000  # one paragraph, no newlines to split at
    chunks = chunk_document("T", body)
    assert len(chunks) > 1
    total_words = sum(chunk.count("word") for chunk in chunks)
    assert total_words == 2000