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

"""Unit: retriever — ordering and dedup, no context cap (A7 revised) (fake rows).

The SQL side of :func:`app.rag.retriever.retrieve` is exercised by the
chat integration tests against real Postgres; the pure mapping logic in
:func:`select_documents` is tested here with in-memory rows.
"""
from __future__ import annotations

import uuid
from types import SimpleNamespace

import pytest

from app.models import Document
from app.rag.retriever import TRUNCATION_MARKER, RetrievedChunk, select_documents


def _doc(path: str, content: str, source: str = "Homelab", title: str | None = None) -> Document:
    return Document(
        id=uuid.uuid4(),
        source=source,
        path=path,
        full_path=f"/tmp/{path}",
        title=title or path,
        content=content,
        content_hash="0" * 64,
    )


def _chunk(doc: Document, score: float, position: int = 0) -> RetrievedChunk:
    return RetrievedChunk(
        chunk_id=uuid.uuid4(),
        position=position,
        content=doc.content[:40],
        score=score,
        document=doc,
    )


def test_ranks_by_best_chunk_score_not_first_hit() -> None:
    """A doc whose *later* chunk scores highest must still rank first."""
    a = _doc("a.md", "A" * 50)
    b = _doc("b.md", "B" * 50)
    c = _doc("c.md", "C" * 50)
    chunks = [
        _chunk(a, 0.4, position=0),  # a's weak chunk comes first
        _chunk(b, 0.8),
        _chunk(a, 0.9, position=2),  # a's best chunk comes last
        _chunk(c, 0.5),
    ]
    docs = select_documents(chunks, n=3)
    assert [d.path for d in docs] == ["a.md", "b.md", "c.md"]


def test_dedups_to_one_document_per_hit_set() -> None:
    a = _doc("a.md", "A" * 50)
    chunks = [_chunk(a, 0.2), _chunk(a, 0.7), _chunk(a, 0.5)]
    docs = select_documents(chunks, n=2)
    assert len(docs) == 1
    assert docs[0] is a


def test_caps_at_n_documents() -> None:
    docs_in = [_doc(f"d{i}.md", "X" * 20) for i in range(4)]
    chunks = [_chunk(d, 0.5 - 0.1 * i) for i, d in enumerate(docs_in)]
    out = select_documents(chunks, n=2)
    assert [d.path for d in out] == ["d0.md", "d1.md"]


def test_content_never_truncated_even_past_old_budget() -> None:
    """Whole documents, never truncated (A7 revised, owner permission 2026-08-24).

    Two documents of 20 000 + 15 000 chars — 35 000 combined, well past
    the old 24 000 context budget — come back with content
    **byte-identical** to the originals, and the truncation marker is
    absent from both.
    """
    big = _doc("big.md", "B" * 20_000)
    small = _doc("small.md", "S" * 15_000)
    chunks = [_chunk(big, 0.9), _chunk(small, 0.6)]
    out = select_documents(chunks, n=2)
    assert [d.path for d in out] == ["big.md", "small.md"]
    assert out[0].content == "B" * 20_000
    assert out[1].content == "S" * 15_000
    assert TRUNCATION_MARKER not in out[0].content
    assert TRUNCATION_MARKER not in out[1].content


def test_empty_hits_yield_no_documents() -> None:
    assert select_documents([], n=2) == []


# ---------------------------------------------------------------------------
# Hybrid retrieval (A7): RRF fusion + lexical tsquery
# ---------------------------------------------------------------------------

from app.rag.retriever import fuse, lexical_tsquery  # noqa: E402


def _rc(
    doc_path: str, cosine: float = 0.0, fts_hit: bool = False, position: int = 0,
    is_summary: bool = False,
) -> RetrievedChunk:
    return RetrievedChunk(
        chunk_id=uuid.uuid4(),
        position=position,
        content="x" * 20,
        score=0.0,
        document=_doc(doc_path, "x" * 20),
        cosine=cosine,
        fts_hit=fts_hit,
        is_summary=is_summary,
    )


def test_lexical_tsquery_tokens_lowercased_deduped_in_order() -> None:
    assert lexical_tsquery("How did I Install GITLAB gitlab?") == "how | did | i | install | gitlab"


def test_lexical_tsquery_punctuation_and_umlauts_ignored() -> None:
    assert lexical_tsquery("c3-r00t? -- what's up!") == "c3 | r00t | what | s | up"


def test_lexical_tsquery_pure_symbols_return_none() -> None:
    assert lexical_tsquery("??? ???") is None
    assert lexical_tsquery("") is None


def test_lexical_tsquery_stopwords_left_to_postgres() -> None:
    # lexical_tsquery passes raw tokens through; Postgres's to_tsquery
    # lexing drops the stopwords (verified against real PG in
    # test_retrieve_empty_kb / integration tests).
    assert lexical_tsquery("how do i") == "how | do | i"


def test_lexical_tsquery_dotted_tokens_kept_whole() -> None:
    """The 2026-09-05 incident: the default parser lexes dotted words
    as ONE lexeme ("llama.cpp" → 'llama.cpp', "Qwen 3.8" → '3.8'), so
    the query carries them whole — split tokens (llama | cpp) can never
    match the document side."""
    assert lexical_tsquery(
        "What are the correct llama.cpp arguments for Qwen 3.8?"
    ) == "what | are | the | correct | llama.cpp | arguments | for | qwen | 3.8"
    # The dash still splits (only dots group): ai | internal.network.
    assert lexical_tsquery("how did I set up ai-internal.network?") == (
        "how | did | i | set | up | ai | internal.network"
    )


def test_fuse_combines_both_lists_for_double_hits() -> None:
    v1 = _rc("a.md", cosine=0.9)
    v2 = _rc("b.md", cosine=0.5)
    l1 = _rc("a.md", cosine=0.1)  # same chunk id -> matched in place
    a_id = v1.chunk_id
    l1.chunk_id = a_id
    out = fuse([v1, v2], [l1], k=60)
    by_id = {rc.chunk_id: rc for rc in out}
    # a: 1/61 (vector rank 1) + 1/61 (lexical rank 1); b: 1/62 only.
    assert by_id[a_id].score == pytest.approx(2 / 61)
    assert by_id[a_id].fts_hit is True
    assert by_id[v2.chunk_id].score == pytest.approx(1 / 62)
    assert by_id[v2.chunk_id].fts_hit is False
    assert [rc.chunk_id for rc in out] == [a_id, v2.chunk_id]


def test_fuse_lexical_only_chunks_enter_with_zero_cosine() -> None:
    vector = [_rc("a.md", cosine=0.8)]
    lexical = [_rc("b.md", cosine=0.0, fts_hit=True)]
    out = fuse(vector, lexical, k=60)
    assert len(out) == 2
    b = next(rc for rc in out if rc.document.path == "b.md")
    assert b.cosine == 0.0
    assert b.fts_hit is True
    # Still ranked by its (only) RRF term.
    assert b.score == pytest.approx(1 / 61)


def test_fuse_orders_by_score_then_cosine_then_path() -> None:
    # Two chunks share an RRF score (both rank 1 in different lists):
    # the higher-cosine one must sort first.
    hi = _rc("z.md", cosine=0.9)
    lo = _rc("a.md", cosine=0.2)
    out = fuse([hi], [lo], k=60)
    assert [rc.document.path for rc in out] == ["z.md", "a.md"]
    # Equal score AND cosine -> path order.
    p1 = _rc("b.md", cosine=0.5)
    p2 = _rc("a.md", cosine=0.5)
    out = fuse([p1], [p2], k=60)
    assert [rc.document.path for rc in out] == ["a.md", "b.md"]
    # Equal score, cosine, path -> position order.
    s1 = _rc("a.md", cosine=0.5, position=1)
    s2 = _rc("a.md", cosine=0.5, position=0)
    out = fuse([s1], [s2], k=60)
    assert [rc.position for rc in out] == [0, 1]


def test_fuse_rejects_nonpositive_k() -> None:
    with pytest.raises(ValueError):
        fuse([], [], k=0)
    with pytest.raises(ValueError):
        fuse([], [], k=-1)


def test_fuse_empty_lists() -> None:
    assert fuse([], [], k=60) == []


# ---------------------------------------------------------------------------
# Phase 30: is_summary survives both candidate lists and the fusion
# ---------------------------------------------------------------------------

from app.models import Chunk  # noqa: E402
from app.rag.retriever import (  # noqa: E402
    NAME_HIT_LIMIT,
    _lexical_candidates,
    _name_hit_chunks,
    _normalize_name,
    _vector_candidates,
    name_hit_tokens,
)


class _FakeResult:
    """Stands in for SQLAlchemy's RowMapping result (``.all()`` only)."""

    def __init__(self, rows: list) -> None:
        self._rows = rows

    def all(self) -> list:
        return self._rows


class _FakeSession:
    """Returns canned rows from ``execute`` without touching Postgres.

    One list of rows (legacy form) is returned for EVERY call; several
    lists (one per successive ``execute``) model a query sequence — the
    name-hit lexical path (2026-09-05) issues the document-projection
    query and, when hits exist, the LATERAL chunk query, BEFORE the FTS
    query.
    """

    def __init__(self, *rowsets: list) -> None:
        if len(rowsets) == 1 and not (
            rowsets[0] and isinstance(rowsets[0][0], list)
        ):
            rowsets = (rowsets[0],)  # the single-rowset legacy form
        self._rowsets = rowsets
        self._call = 0
        self.statements: list = []

    def execute(self, stmt, params: dict | None = None) -> _FakeResult:
        self.statements.append((stmt, params))
        rows = self._rowsets[min(self._call, len(self._rowsets) - 1)]
        self._call += 1
        return _FakeResult(rows)


def _chunk_row(is_summary: bool) -> Chunk:
    doc = _doc("summary-src.yaml", "RAW_YAML_CONTENT")
    return Chunk(
        id=uuid.uuid4(),
        document_id=doc.id,
        position=-1,  # the summary chunk's position (phase 30)
        content="Summary text",
        is_summary=is_summary,
    )


def test_vector_candidates_carry_is_summary_flag() -> None:
    """The vector list copies ``Chunk.is_summary`` onto each candidate."""
    doc = _doc("summary-src.yaml", "RAW_YAML_CONTENT")
    summary = _chunk_row(is_summary=True)
    ordinary = _chunk_row(is_summary=False)
    ordinary.position = 0
    ordinary.content = "ordinary content"
    rows = [
        (summary, 0.123456, doc),
        (ordinary, 0.2, doc),
    ]
    out = _vector_candidates(_FakeSession(rows), [0.0] * 768, limit=5)  # pyright: ignore[reportArgumentType]
    assert len(out) == 2
    by_pos = {rc.position: rc for rc in out}
    assert by_pos[-1].is_summary is True  # the summary chunk (position −1)
    assert by_pos[0].is_summary is False  # ordinary content chunk
    assert by_pos[-1].cosine == pytest.approx(0.876544)  # 1 − distance, still rounded


def test_vector_candidates_default_is_summary_false_for_legacy_chunks() -> None:
    """Pre-phase-30 rows have ``is_summary=false`` — candidates stay False."""
    doc = _doc("legacy.md", "LEGACY")
    legacy = Chunk(
        id=uuid.uuid4(),
        document_id=doc.id,
        position=0,
        content="legacy content",
        is_summary=False,
    )
    out = _vector_candidates(_FakeSession([(legacy, 0.5, doc)]), [0.0] * 768, limit=5)  # pyright: ignore[reportArgumentType]
    assert out[0].is_summary is False


def _lexical_row(is_summary: bool, doc_path: str) -> object:
    """One row of ``_LEXICAL_SQL`` (attribute access, as SQLAlchemy returns)."""
    doc = _doc(doc_path, "DOC_BODY")
    return SimpleNamespace(
        chunk_id=uuid.uuid4(),
        position=-1 if is_summary else 0,
        content="summary chunk text" if is_summary else "content chunk text",
        doc_id=doc.id,
        source=doc.source,
        path=doc.path,
        full_path=doc.full_path,
        title=doc.title,
        doc_content=doc.content,
        content_hash=doc.content_hash,
        indexed_at=None,
        is_summary=is_summary,
        rank=0.33,
    )


def test_lexical_candidates_carry_is_summary_flag() -> None:
    """The lexical list reads ``c.is_summary`` from the raw row.

    The question carries no digit-bearing name token (no bare, no
    numeric-join), so the name-hit path issues NO queries at all — the
    single FTS rowset answers the only (FTS) call, and the list is the
    plain FTS rows: the pre-name-hit behavior, unchanged.
    """
    rows = [_lexical_row(True, "summary-src.yaml"), _lexical_row(False, "other.md")]
    out = _lexical_candidates(
        _FakeSession(rows), "how do i configure the thing", limit=10  # pyright: ignore[reportArgumentType]
    )
    assert len(out) == 2
    by_path = {rc.document.path: rc for rc in out}
    assert by_path["summary-src.yaml"].is_summary is True
    assert by_path["summary-src.yaml"].position == -1
    assert by_path["other.md"].is_summary is False
    assert all(rc.fts_hit is True for rc in out)


def test_fuse_keeps_is_summary_on_double_hit() -> None:
    """A summary chunk in both lists keeps the flag after fusion."""
    v1 = _rc("s.yaml", cosine=0.9, is_summary=True)
    l1 = _rc("s.yaml", cosine=0.9, is_summary=True)  # lexical copy of the same chunk
    l1.chunk_id = v1.chunk_id
    out = fuse([v1], [l1], k=60)
    assert len(out) == 1
    assert out[0].is_summary is True
    assert out[0].fts_hit is True
    assert out[0].score == pytest.approx(2 / 61)


def test_fuse_keeps_is_summary_on_lexical_only_hit() -> None:
    """A summary-only lexical hit (no vector rank) keeps the flag."""
    out = fuse([], [_rc("s.yaml", is_summary=True)], k=60)
    assert len(out) == 1
    assert out[0].is_summary is True
    assert out[0].fts_hit is True
    assert out[0].cosine == 0.0


def test_fuse_default_is_summary_stays_false_for_legacy_chunks() -> None:
    """Neither list flagged ⇒ fusion never invents a summary flag."""
    out = fuse([_rc("a.md", cosine=0.8)], [_rc("b.md")], k=60)
    assert len(out) == 2
    assert all(rc.is_summary is False for rc in out)


# ---------------------------------------------------------------------------
# Name-hit lexical signal (the 2026-09-05 incident — the versioned-name
# case the default parser lexes incompatibly: "Qwen 3.8" → qwen/3/8 can
# never match a document's qwen3/8/27b tokens)
# ---------------------------------------------------------------------------

INCIDENT_QUESTION = "What are the correct llama.cpp arguments for Qwen 3.8?"


def test_normalize_name() -> None:
    assert _normalize_name("Qwen 3.8") == "qwen38"
    assert _normalize_name("qwen3.8-27b-juggernaut-vulkan") == "qwen3827bjuggernautvulkan"
    assert _normalize_name("Mixed CASE-99") == "mixedcase99"
    assert _normalize_name("!!!") == ""


def test_name_hit_tokens_incident_question() -> None:
    """The incident question yields EXACTLY the versioned join
    ``qwen38`` — the token the document names actually carry. Plain
    prose words (``what``, ``llamacpp``, ``arguments``, ``server`` —
    no digit) never name-match (the precision guard); the single
    digits ("3", "8") and the bare "38" are < 4 chars; the
    digit-leading ``38show`` boundary artifact is dropped."""
    tokens = name_hit_tokens(INCIDENT_QUESTION)
    assert tokens == ["qwen38"]
    for absent in ("what", "qwen", "llamacpp", "arguments", "3", "8", "38", "38show", "server"):
        assert absent not in tokens


def test_name_hit_tokens_no_digit_question_returns_empty() -> None:
    """A question with no digit-bearing token (bare or joined) yields
    no name candidates — prose joins like ``correctllama`` never count."""
    assert name_hit_tokens("what is the correct caddy config") == []
    assert name_hit_tokens("a e i o u 3 8") == []


def test_name_hit_tokens_bare_digit_bearing_token() -> None:
    """A single written token that carries a digit (``1panel``) is a
    name candidate on its own — no join needed."""
    tokens = name_hit_tokens("what is my 1panel dashboard setup")
    assert tokens == ["1panel"]


def _name_row(doc: Document) -> tuple:
    """One row of the name-hit document projection (catalog order)."""
    return (doc.id, doc.source, doc.path, doc.title)


def _name_hit_lateral_row(doc: Document, is_summary: bool = False) -> SimpleNamespace:
    """One row of the name-hit LATERAL chunk query."""
    return SimpleNamespace(
        doc_id=doc.id,
        source=doc.source,
        path=doc.path,
        full_path=doc.full_path,
        title=doc.title,
        doc_content=doc.content,
        content_hash=doc.content_hash,
        indexed_at=None,
        chunk_id=uuid.uuid4(),
        position=-1 if is_summary else 0,
        content="summary chunk" if is_summary else "content chunk",
        is_summary=is_summary,
    )


def test_name_hit_chunks_no_tokens_skips_all_queries() -> None:
    """A question with no name tokens issues no queries at all."""
    session = _FakeSession([])  # any call would surface a statement
    assert _name_hit_chunks(session, "a e i o u 3 8") == []  # pyright: ignore[reportArgumentType]
    assert session.statements == []


def test_name_hit_chunks_no_matching_doc_returns_empty() -> None:
    """Name tokens exist but no document name carries one: the
    projection runs, the LATERAL fetch does not."""
    doc = _doc("quadlets/other.container", "body")
    name_rows = [_name_row(doc)]
    session = _FakeSession(name_rows, [])
    assert _name_hit_chunks(session, INCIDENT_QUESTION) == []  # pyright: ignore[reportArgumentType]
    assert len(session.statements) == 1  # projection only — no LATERAL fetch


def test_name_hit_chunks_ranked_by_count_length_catalog() -> None:
    """A two-candidate question (``qwen38`` + ``1panel``): the document
    whose name carries BOTH (2 matches, 12 total chars) leads; the two
    single-match documents tie on (1, 6) and fall to catalog order
    (``dashboards/1panel-notes.md`` before ``quadlets/qwen3.8…``).
    Hits carry ``fts_hit=True`` (the A8 gate answers), ``cosine=0.0``,
    and the summary flag of their representative chunk."""
    both = _doc("dashboards/1panel-qwen3.8.md", "body", title="1Panel Qwen 3.8")
    panel = _doc("dashboards/1panel-notes.md", "body")
    q38 = _doc("quadlets/qwen3.8-27b-juggernaut-vulkan.container", "body")
    name_rows = [_name_row(d) for d in (panel, both, q38)]  # catalog order
    question = "what are the correct llama.cpp arguments for qwen 3.8 and the 1panel dashboard?"
    lateral_rows = [
        _name_hit_lateral_row(q38, is_summary=True),  # LATERAL may return any order
        _name_hit_lateral_row(both),
        _name_hit_lateral_row(panel),
    ]
    session = _FakeSession(name_rows, lateral_rows)
    out = _name_hit_chunks(session, question)  # pyright: ignore[reportArgumentType]
    assert [rc.document.path for rc in out] == [
        "dashboards/1panel-qwen3.8.md",  # 2 matched tokens — leads
        "dashboards/1panel-notes.md",  # (1, 6) — catalog order
        "quadlets/qwen3.8-27b-juggernaut-vulkan.container",  # (1, 6) — after
    ]
    assert all(rc.fts_hit is True for rc in out)  # the lexical signal
    assert all(rc.cosine == 0.0 for rc in out)  # no vector rank
    assert all(rc.score == 0.0 for rc in out)  # fuse fills the score
    by_path = {rc.document.path: rc for rc in out}
    assert by_path["quadlets/qwen3.8-27b-juggernaut-vulkan.container"].is_summary is True
    assert by_path["quadlets/qwen3.8-27b-juggernaut-vulkan.container"].position == -1
    assert by_path["dashboards/1panel-notes.md"].is_summary is False


def test_name_hit_chunks_capped_at_limit() -> None:
    """Twelve tied name hits (one matched token each) yield exactly
    ``NAME_HIT_LIMIT`` of them — catalog order (the deterministic
    tie-break)."""
    docs = [_doc(f"quadlets/m{i:02d}.container", "body") for i in range(12)]
    for d in docs:  # give every document a name that carries the token
        d.title = "qwen38 model i"
    name_rows = [_name_row(d) for d in docs]
    # Only the ten winners (catalog order — the deterministic tie-break
    # of the twelve identical scores) reach the LATERAL fetch; the fake
    # answers with exactly those rows.
    lateral_rows = [_name_hit_lateral_row(d) for d in docs[:NAME_HIT_LIMIT]]
    session = _FakeSession(name_rows, lateral_rows)
    out = _name_hit_chunks(
        session, "tell me about the qwen 3.8 models"  # pyright: ignore[reportArgumentType]
    )
    assert len(out) == NAME_HIT_LIMIT
    assert [rc.document.path for rc in out] == [f"quadlets/m{i:02d}.container" for i in range(10)]


def test_lexical_candidates_name_hits_lead_and_dedupe_with_fts() -> None:
    """The full lexical list: name hits LEAD (their representative
    chunks), the FTS rows follow, and an FTS row sharing the name hit's
    chunk id appears exactly once (deduped)."""
    q38 = _doc("quadlets/qwen3.8-27b-juggernaut-vulkan.container", "body")
    other = _doc("quadlets/qwen38-other.container", "body")  # 1 matched token
    name_rows = [_name_row(q38), _name_row(other)]
    q38_chunk = uuid.uuid4()

    def _lateral(doc: Document) -> SimpleNamespace:
        row = _name_hit_lateral_row(doc)
        if doc is q38:
            row.chunk_id = q38_chunk
        return row

    lateral_rows = [_lateral(q38), _lateral(other)]
    fts_rows = [
        # an FTS hit on the SAME chunk as the q38 name hit (deduped away)
        SimpleNamespace(
            chunk_id=q38_chunk, position=1, content="c", doc_id=q38.id,
            source=q38.source, path=q38.path, full_path=q38.full_path,
            title=q38.title, doc_content=q38.content, content_hash=q38.content_hash,
            indexed_at=None, is_summary=False, rank=0.1,
        ),
        # an FTS hit on a different chunk of the OTHER doc (kept)
        _lexical_row(False, "quadlets/qwen38-other.container"),
    ]
    session = _FakeSession(name_rows, lateral_rows, fts_rows)
    out = _lexical_candidates(session, INCIDENT_QUESTION, limit=10)  # pyright: ignore[reportArgumentType]
    assert len(out) == 3  # q38 (once), other (name hit), other (FTS chunk)
    # Both name hits tie on (1, 6) — catalog order: "qwen3." (ASCII 46)
    # sorts before "qwen38" (ASCII 56).
    assert out[0].document.path == "quadlets/qwen3.8-27b-juggernaut-vulkan.container"
    assert out[1].document.path == "quadlets/qwen38-other.container"
    assert out[0].chunk_id == q38_chunk  # the name-hit representative row
    assert {
        rc.chunk_id for rc in out
    } == {q38_chunk, fts_rows[1].chunk_id, lateral_rows[1].chunk_id}
    assert all(rc.fts_hit is True for rc in out)