File size: 1,363 Bytes
978f645 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 | """Query → top-k chunks. Encapsulates the embedder + store pair so callers
don't have to assemble both. Loads from disk lazily.
"""
from __future__ import annotations
from pathlib import Path
from src.core.logger import get_logger
from src.rag.embed import EMBEDDING_DIM, Embedder
from src.rag.store import FAISSStore
logger = get_logger(__name__)
class RAGRetriever:
"""Bundle (embedder, store). Use `RAGRetriever.load(dir)` to construct."""
def __init__(self, store: FAISSStore, embedder: Embedder) -> None:
self._store = store
self._embedder = embedder
@classmethod
def load(cls, index_dir: Path) -> "RAGRetriever":
store = FAISSStore.load(Path(index_dir), dim=EMBEDDING_DIM)
return cls(store=store, embedder=Embedder())
def __len__(self) -> int:
return len(self._store)
def search(self, query: str, k: int = 5) -> list[dict]:
"""Return up to `k` chunks most relevant to `query`, sorted by score desc.
Each chunk dict carries `text`, `source`, `chunk_index`, `score`.
Returns [] for empty query or empty store.
"""
if not query.strip() or len(self._store) == 0:
return []
vec = self._embedder.encode([query])
hits = self._store.search(vec[0], k=k)
return [{**chunk, "score": score} for chunk, score in hits]
|