326 lines
11 KiB
Python
326 lines
11 KiB
Python
"""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_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 _lexical_candidates, _vector_candidates # noqa: E402
|
||
|
||
|
||
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."""
|
||
|
||
def __init__(self, rows: list) -> None:
|
||
self._rows = rows
|
||
self.statements: list = []
|
||
|
||
def execute(self, stmt, params: dict | None = None) -> _FakeResult:
|
||
self.statements.append((stmt, params))
|
||
return _FakeResult(self._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."""
|
||
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)
|