"""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 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 ) -> 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, ) 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_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) == []