FabledGrimoire/tests/test_chunking.py

"""chunk_document — the document shape a record is embedded as.

The model window is 512 tokens and fastembed truncates silently, so without
chunking everything past ~400 words of a record would be invisible to
semantic search. These tests pin the two halves of the contract: short
records keep the single-document 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.
"""
import re

from scribe.services.embeddings import (
    _CHUNK_CHAR_BUDGET,
    _MIN_CHUNK_CONTENT,
    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
    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 chunking: every part of a 2,000-word dev-log can
    influence retrieval. 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


# --- no chunk is a heading alone ---------------------------------------------
#
# A heading or a one-line preamble as a whole chunk ranks at the top of vague
# queries — text about nothing embeds close to everything. Each case below is
# one of those shapes.

_HEADING = re.compile(r"^#{1,6}\s")


def _said(chunk: str, title: str) -> str:
    """A chunk's text once its title prefix and heading lines are set aside."""
    body = chunk[len(title) + 1:] if chunk.startswith(title + "\n") else chunk
    return "\n".join(ln for ln in body.splitlines() if not _HEADING.match(ln)).strip()


def _assert_nothing_thin_and_nothing_lost(title: str, body: str) -> list[str]:
    chunks = chunk_document(title, body)
    assert len(chunks) > 1, "the case must be long enough to split"
    for chunk in chunks:
        assert len(_said(chunk, title)) >= _MIN_CHUNK_CONTENT, (
            f"a chunk says almost nothing: {chunk[:120]!r}")
        assert len(chunk) <= _CHUNK_CHAR_BUDGET + len(title) + 1 + 80
    joined = "\n".join(chunks)
    for line in (ln.strip() for ln in body.splitlines()):
        # A line longer than a chunk is cut across two, so only its ends can
        # be looked for whole.
        for part in ([line] if len(line) < 400 else [line[:60], line[-60:]]):
            assert part in joined, f"content dropped: {part[:60]!r}"
    return chunks


def test_a_heading_over_one_long_paragraph_is_not_a_chunk_of_its_own():
    """A work log whose entry is one long paragraph: the split used to flush
    the heading as its own piece before cutting the paragraph."""
    entry = " ".join(["The deploy was verified against the live instance."] * 40)
    _assert_nothing_thin_and_nothing_lost(
        "Step 6 — webhook drift-flagging",
        "The step body.\n\n## Work log — 2026-08-16\n\n" + entry,
    )


def test_a_heading_with_nothing_under_it_joins_what_follows():
    """'## Findings' directly above its '### …' parts is a section with no
    body, and stood alone whenever the section before it was full."""
    body = (
        # Five paragraphs fill the section before it, so the heading cannot
        # merge backwards and has to be carried forward.
        "## Context\n\n" + _long_section("context", paragraphs=5)
        + "\n\n## Findings worth carrying forward\n\n"
        + "### One\n\n" + _long_section("one", paragraphs=5)
        + "\n\n### Two\n\n" + _long_section("two", paragraphs=5)
    )
    _assert_nothing_thin_and_nothing_lost("A milestone dev-log", body)


def test_a_lone_lead_in_line_joins_the_section_after_it():
    _assert_nothing_thin_and_nothing_lost(
        "2026-06-10 — UI batch",
        "Two dev→main PRs this session.\n\n## PR 89\n\n"
        + _long_section("pr89", paragraphs=8),
    )


def test_a_thin_last_section_joins_the_one_before_it():
    _assert_nothing_thin_and_nothing_lost(
        "T",
        "## Body\n\n" + _long_section("body", paragraphs=8) + "\n\n## Next\n\nTBD.",
    )


def test_a_continuation_piece_repeats_the_heading_it_sits_under():
    """A thin heading merged forward puts two headings in one section; the
    pieces of the second must carry the second, not the first."""
    body = "## Parent\n\n## Child\n\n" + _long_section("child", paragraphs=40)
    chunks = chunk_document("T", body)
    assert len(chunks) > 2
    for chunk in chunks[1:]:
        assert "## Child" in chunk


# --- the read path: best chunk wins ------------------------------------------


async def test_search_collapses_chunk_rows_to_best_chunk_per_note():
    """Rows arrive at CHUNK grain ordered by distance; a note appearing via
    several chunks must come back ONCE, scored by its best chunk — otherwise a
    long record fills the top-k with copies of itself."""
    from unittest.mock import AsyncMock, MagicMock, patch

    from scribe.services import embeddings as emb

    note_a, note_b = MagicMock(id=1), MagicMock(id=2)
    # Rows are (Note, distance, chunk_index, chunk_text) — the chunk columns
    # ride along so the collapse can report WHICH passage won.
    rows = [
        (note_a, 0.10, 3, "the passage that actually matched"),
        (note_b, 0.20, 0, "b's best"),
        (note_a, 0.25, 7, "a worse chunk of a"),
        (note_a, 0.30, 1, "a worse chunk of a"),
    ]
    result = MagicMock()
    result.all.return_value = rows
    session, ctx = _session_ctx()
    session.execute = AsyncMock(return_value=result)

    report: dict = {}
    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(emb, "get_embedding", AsyncMock(return_value=[0.0] * 384)),
    ):
        out = await emb.semantic_search_notes(
            1, "a query", limit=8, demote_superseded=False, report=report
        )

    assert [note.id for _s, note in out] == [1, 2]
    assert out[0][0] == 1.0 - 0.10  # the BEST chunk's score, not a later one

    # And the winning chunk is reported, not merely used for scoring. Without
    # this a caller can only preview the head of the body — a span this query
    # has already ranked lower than the one that won.
    assert report["best_chunk"][1] == {
        "index": 3, "text": "the passage that actually matched",
    }
    assert report["best_chunk"][2]["index"] == 0


# --- the write path: one row per chunk ---------------------------------------


def _session_ctx(log_rows=()):
    """A session stand-in. `log_rows` answers the work-log read a task's
    document needs — answered explicitly rather than left to
    autovivify, because `list(result.all())` on a bare MagicMock raises and is
    swallowed, which would quietly make every one of these a no-log test."""
    from unittest.mock import AsyncMock, MagicMock

    session = MagicMock()
    result = MagicMock()
    result.all.return_value = list(log_rows)
    session.execute = AsyncMock(return_value=result)
    session.commit = AsyncMock()
    ctx = MagicMock()
    ctx.__aenter__ = AsyncMock(return_value=session)
    ctx.__aexit__ = AsyncMock(return_value=False)
    return session, ctx


async def test_upsert_stores_one_versioned_row_per_chunk():
    from unittest.mock import AsyncMock, patch

    from scribe.services import embeddings as emb

    body = "\n\n".join(
        f"## Section {i}\n\n{_long_section(f'sec-{i}')}" for i in range(6)
    )
    chunks = chunk_document("T", body)
    assert len(chunks) > 1

    session, ctx = _session_ctx()
    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(
            emb, "get_embeddings",
            AsyncMock(return_value=[[0.0] * 384 for _ in chunks]),
        ),
    ):
        await emb.upsert_note_embedding(7, 42, "T", body)

    rows = [call.args[0] for call in session.add.call_args_list]
    assert [r.chunk_index for r in rows] == list(range(len(chunks)))
    assert [r.chunk_text for r in rows] == chunks
    assert {r.chunker_version for r in rows} == {emb.CHUNKER_VERSION}
    assert {r.embedding_model for r in rows} == {emb.EMBEDDING_MODEL}
    assert {r.user_id for r in rows} == {42}
    session.execute.assert_awaited()  # the delete that makes replacement atomic


async def test_a_tasks_work_logs_reach_the_rows_it_is_stored_as():
    """The wiring half: the shaper is pure and tested next door, so what is
    pinned here is that the WRITER actually asks for the logs and embeds what
    comes back. A log that is written, readable and absent
    from the index is the same half-surface one layer down."""
    import datetime
    from unittest.mock import AsyncMock, patch

    from scribe.services import embeddings as emb

    session, ctx = _session_ctx(
        log_rows=[(datetime.datetime(2026, 9, 20), "ruled out the cache theory")]
    )
    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(emb, "get_embeddings", AsyncMock(return_value=[[0.0] * 384])),
    ):
        await emb.upsert_note_embedding(7, 42, "A task", "Its own prose.")

    rows = [call.args[0] for call in session.add.call_args_list]
    stored = "\n".join(r.chunk_text for r in rows)
    assert "ruled out the cache theory" in stored
    assert "Its own prose." in stored
    assert emb.WORK_LOG_HEADING in stored


async def test_upsert_of_an_emptied_record_clears_rows_instead_of_embedding():
    """An empty embedding is worse than none, and a STALE one is worse than
    that — a record emptied of content must stop being findable by what it no
    longer says."""
    from unittest.mock import patch

    from scribe.services import embeddings as emb

    session, ctx = _session_ctx()
    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(emb, "get_embeddings") as embedder,
    ):
        await emb.upsert_note_embedding(7, 42, "", "")

    embedder.assert_not_called()
    session.execute.assert_awaited()   # the delete
    session.add.assert_not_called()
    session.commit.assert_awaited()


async def test_backfill_reembeds_notes_with_a_stale_chunker_version():
    """The reason chunker_version exists: a shape change becomes a version
    bump that re-embeds exactly the stale notes, instead of another 0077-style
    table wipe. Only rows AT the current version count as done."""
    from unittest.mock import AsyncMock, MagicMock, patch

    from scribe.services import embeddings as emb

    def _ids(*values):
        r = MagicMock()
        r.fetchall.return_value = [(v,) for v in values]
        return r

    session, ctx = _session_ctx()
    # The opening scan, in order: ids already at the current version, every
    # note id, vectors older than their note, tasks logged since their vectors.
    # IDS ONLY — the text is re-read per note at embed time.
    session.execute = AsyncMock(side_effect=[
        _ids(1),            # note 1 is current
        _ids(1, 2),         # the corpus
        _ids(),             # nothing stale by timestamp
        _ids(),             # no task logged since its vectors
    ])

    with (
        patch.object(emb, "async_session", return_value=ctx),
        patch.object(emb, "_current_row",
                     AsyncMock(return_value=(42, "stale-version", "body", "note", None))),
        patch.object(emb, "upsert_note_embedding", AsyncMock()) as upsert,
        patch.object(emb.asyncio, "sleep", AsyncMock()),
    ):
        await emb.backfill_note_embeddings()

    embedded = [call.args[0] for call in upsert.call_args_list]
    assert embedded == [2], "only the stale note is re-embedded"


def test_a_row_is_current_only_in_the_live_models_space():
    """The column is a width, not an identity, so a same-width model
    swap writes a second geometry with no error. The backfill's "current"
    predicate has to name BOTH halves of the stamp, on every embedding table."""
    from scribe.models.embedding import (
        MilestoneEmbedding, NoteEmbedding, RuleEmbedding, SystemEmbedding,
    )
    from scribe.services import embeddings as emb

    for table in (NoteEmbedding, RuleEmbedding, MilestoneEmbedding, SystemEmbedding):
        clause = str(emb.is_current_stamp(table))
        assert "chunker_version" in clause and "embedding_model" in clause, table