unsloth/tests/python/test_rag_retrieval.py
Roland Tannous 92994e8b83 Studio: add RAG with hybrid search, reranker, chat integration
Backend (studio/backend/):
- core/rag/: parsers (PDF/TXT/MD/DOCX/HTML via pypdf/python-docx/bs4),
  recursive token-aware chunker, embeddings singleton via
  FastSentenceTransformer.from_pretrained(for_inference=True), Qdrant
  local vector store, bm25s lexical index, RRF hybrid retrieval,
  spawn-subprocess ingestion job with SSE progress, optional
  CrossEncoder reranker (off-by-default).
- routes/rag.py: KB CRUD, doc upload (KB + per-thread), doc list/delete,
  ingestion SSE, hybrid+rerank search, thread-index list/clear.
- routes/chat_history.py: purge thread RAG artifacts on thread delete
  and clear-all (rag_documents has no FK cascade to chat_threads so
  uploads work on un-persisted threads).
- studio.db gains 4 RAG tables; storage_roots gains rag_*() helpers.
- auth/authentication.py: get_current_subject_sse accepts ?token=... so
  EventSource can stream ingestion progress.

Frontend (studio/frontend/):
- features/rag/: api client, Zustand store, hooks, dropzone, KB list,
  doc rows, ingestion-progress, thread-index list components.
- Settings dialog gains a Knowledge Bases tab (master/detail + thread
  documents list); /knowledge-bases deep-links to it.
- features/chat/: per-thread ragSource/enableRerank/ragTopK state in
  chat-runtime-store; Retrieval section in chat-settings-sheet with KB
  DropdownMenu (active highlight + per-row trash), thread doc list with
  Clear-thread-index button, RAG Top K slider, reranker toggle;
  chat-adapter retrieves before /v1/chat/completions and injects hits
  as a system block; shared-composer + button routes documents into
  pendingDocs (auto-uploads, send blocked while indexing).
2026-05-23 18:46:15 +04:00

46 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]
# b appears at rank 2 in both -> highest fused score
assert ids[0] == "b"
assert set(ids) == {"a", "b", "c"} or set(ids) == {"b", "c", "a"}
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 order.
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