File size: 1,449 Bytes
fd3cc1a | 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 | """High-level facade tying ingestion, retrieval, and answering together."""
from __future__ import annotations
from pathlib import Path
from .config import settings
from .embeddings import get_embedder
from .ingestion import ingest_documents
from .llm import Answer, answer
from .retrieval import Retriever, RetrievalMode
from .store import VectorStore
class RagEngine:
def __init__(self, store: VectorStore | None = None) -> None:
self.embedder = get_embedder()
self.store = store or VectorStore(dim=self.embedder.dim)
self.retriever = Retriever(self.store, self.embedder)
def add_documents(self, documents: list[dict], **kwargs) -> int:
ingest_documents(documents, store=self.store, embedder=self.embedder, **kwargs)
self.retriever._bm25 = None # invalidate cached BM25 index
return len(self.store)
def query(self, question: str, k: int = 5, mode: RetrievalMode = "hybrid") -> Answer:
chunks = self.retriever.retrieve(question, k=k, mode=mode)
return answer(question, chunks)
# --- persistence -------------------------------------------------------
def save(self, directory: str | Path | None = None) -> None:
self.store.save(directory or settings.index_dir)
@classmethod
def load(cls, directory: str | Path | None = None) -> "RagEngine":
store = VectorStore.load(directory or settings.index_dir)
return cls(store=store)
|