unsloth/tests/python/test_rag_retrieval.py
Daniel Han ab0828b976 Studio: fix RAG correctness bugs
Backend:
- Deterministic SQLite connection cleanup. The RAG code used bare
  `with get_connection() as conn:`, which commits but never closes, leaning
  on GC to release handles (the rest of studio_db closes explicitly). Add a
  closing_connection() context manager that commits/rolls back like sqlite3's
  own manager and always closes, and route all 30 RAG call sites through it.
- filter_by_min_score no longer drops BM25-only and figure-ref hits. min_score
  is a cosine floor, so it now gates only hits that carry a dense_score;
  lexical and figure-ref hits (dense_score is None) pass through instead of
  being silently discarded when the floor is raised.
- Fix two tests that could not pass against the production code: the RRF
  fusion test asserted the wrong winner (c edges out b: 0.032266 vs 0.032258),
  and two tool-handler scope tests stubbed retrieve_hybrid without accepting
  the embedder_model kwarg the handler now passes (TypeError was swallowed,
  leaving captured["scope"] unset).

Frontend:
- Removing an in-flight upload chip now routes through the teardown thunk
  already registered for the aggregate-progress toast (abort, unsubscribe,
  release the index slot, delete the backend doc with the correct kb/thread
  scope key it closed over) and clears the toast entry. Deleting directly
  leaked the concurrency slot and hardcoded the thread scope, mis-targeting
  KB-scoped docs. Applied in both the composer hook and the compare-view
  composer; drop the now-vestigial chip-scope-key tracking and unused
  activeThreadId selectors. Add index-progress-store.remove(id).
2026-05-31 09:56:23 +00:00

47 lines
1.6 KiB
Python

"""Unit tests for RAG RRF fusion — no external deps."""
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[2]
STUDIO_BACKEND = REPO_ROOT / "studio" / "backend"
if str(STUDIO_BACKEND) not in sys.path:
sys.path.insert(0, str(STUDIO_BACKEND))
from core.rag.retrieval import Hit, _rrf_fuse
def test_rrf_fuses_two_rankings():
bm25 = [Hit("a", 10.0), Hit("b", 8.0), Hit("c", 5.0)]
dense = [Hit("c", 0.9), Hit("b", 0.8), Hit("d", 0.5)]
fused = _rrf_fuse([bm25, dense], rrf_k = 60, top_k = 3)
ids = [h.chunk_id for h in fused]
# c (bm25 rank2 1/63 + dense rank0 1/61 = 0.032266) narrowly beats
# b (rank1 in both = 2/62 = 0.032258); d falls outside top_k.
assert ids[0] == "c"
assert set(ids) == {"a", "b", "c"}
def test_rrf_top_k_limits_output():
rankings = [
[Hit(f"r1_{i}", 0.0) for i in range(20)],
[Hit(f"r2_{i}", 0.0) for i in range(20)],
]
fused = _rrf_fuse(rankings, rrf_k = 60, top_k = 5)
assert len(fused) == 5
def test_rrf_unique_ranking():
# Single ranking: fused order matches input.
ranking = [Hit("x", 0.0), Hit("y", 0.0), Hit("z", 0.0)]
fused = _rrf_fuse([ranking], rrf_k = 60, top_k = 3)
assert [h.chunk_id for h in fused] == ["x", "y", "z"]
def test_rrf_preserves_payload_from_first_ranking():
a = Hit("a", 1.0, document_id = "doc1", chunk_index = 5)
b = Hit("a", 2.0, document_id = "doc2", chunk_index = 7)
fused = _rrf_fuse([[a], [b]], rrf_k = 60, top_k = 1)
# First sighting wins for payload (deterministic).
assert fused[0].document_id == "doc1"
assert fused[0].chunk_index == 5