198 lines
6.7 KiB
Python
198 lines
6.7 KiB
Python
"""Unit: retriever — ordering, dedup, and the context cap (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, max_chars=10_000)
|
|
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, max_chars=10_000)
|
|
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, max_chars=10_000)
|
|
assert [d.path for d in out] == ["d0.md", "d1.md"]
|
|
|
|
|
|
def test_combined_content_capped_with_truncation_marker() -> None:
|
|
big = _doc("big.md", "B" * 100)
|
|
small = _doc("small.md", "S" * 100)
|
|
chunks = [_chunk(big, 0.9), _chunk(small, 0.6)]
|
|
out = select_documents(chunks, n=2, max_chars=150)
|
|
# Best doc stays intact; the overflowing one is truncated in place.
|
|
assert out[0].content == "B" * 100
|
|
assert out[1].content.endswith(TRUNCATION_MARKER)
|
|
assert out[1].content.startswith("S")
|
|
assert len(out[0].content) + len(out[1].content) <= 150
|
|
|
|
|
|
def test_single_doc_over_budget_is_truncated_to_budget() -> None:
|
|
big = _doc("big.md", "Z" * 200)
|
|
out = select_documents([_chunk(big, 0.9)], n=2, max_chars=50)
|
|
assert len(out[0].content) == 50
|
|
assert out[0].content.endswith(TRUNCATION_MARKER)
|
|
|
|
|
|
def test_under_budget_no_truncation() -> None:
|
|
a = _doc("a.md", "A" * 80)
|
|
b = _doc("b.md", "B" * 60)
|
|
out = select_documents([_chunk(b, 0.5), _chunk(a, 0.9)], n=2, max_chars=200)
|
|
assert [d.path for d in out] == ["a.md", "b.md"]
|
|
assert a.content == "A" * 80 and b.content == "B" * 60
|
|
assert TRUNCATION_MARKER not in a.content + b.content
|
|
|
|
|
|
def test_empty_hits_yield_no_documents() -> None:
|
|
assert select_documents([], n=2, max_chars=24_000) == []
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# 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) == []
|