Create rag_engine.py

rag_engine.py

@@ -0,0 +1,222 @@
+"""
+RAG Engine — document ingestion, chunking, embedding, FAISS indexing, retrieval.
+"""
+
+import os
+import pickle
+from pathlib import Path
+from typing import Optional
+
+from langchain.docstore.document import Document
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import (
+    PyPDFLoader,
+    Docx2txtLoader,
+    TextLoader,
+)
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
+
+from config import cfg
+from logging_config import get_logger, setup_logging
+from utils import validate_file, file_checksum, ensure_dir, Timer
+
+setup_logging(log_dir=cfg.app.log_dir)
+logger = get_logger(__name__)
+
+
+class RAGEngine:
+    """
+    Handles the full RAG lifecycle:
+    load → chunk → embed → store → retrieve.
+    """
+
+    def __init__(self) -> None:
+        self.embeddings: Optional[HuggingFaceEmbeddings] = None
+        self.vector_store: Optional[FAISS] = None
+        self.ingested_checksums: set[str] = set()
+        self.all_chunks: list[Document] = []
+        self._splitter = RecursiveCharacterTextSplitter(
+            chunk_size=cfg.chunking.chunk_size,
+            chunk_overlap=cfg.chunking.chunk_overlap,
+            separators=cfg.chunking.separators,
+        )
+        logger.info("RAGEngine initialised.")
+
+    # ── Embedding model ───────────────────────────────────────────────────────
+
+    def load_embeddings(self) -> None:
+        if self.embeddings is not None:
+            return
+        logger.info("Loading embedding model: %s", cfg.embedding.model_name)
+        with Timer() as t:
+            self.embeddings = HuggingFaceEmbeddings(
+                model_name=cfg.embedding.model_name,
+                encode_kwargs={"normalize_embeddings": cfg.embedding.normalize_embeddings},
+            )
+        logger.info("Embedding model loaded in %s.", t)
+
+    # ── Document loaders ──────────────────────────────────────────────────────
+
+    def _load_single(self, filepath: str) -> list[Document]:
+        ext = Path(filepath).suffix.lower()
+        loaders = {
+            ".pdf":  lambda: PyPDFLoader(filepath),
+            ".docx": lambda: Docx2txtLoader(filepath),
+            ".txt":  lambda: TextLoader(filepath, encoding="utf-8"),
+        }
+        if ext not in loaders:
+            logger.warning("Unsupported file type: %s", ext)
+            return []
+        try:
+            loader = loaders[ext]()
+            docs = loader.load()
+            # Normalise metadata
+            for doc in docs:
+                doc.metadata["source"] = filepath
+            logger.info("Loaded %d page(s) from '%s'.", len(docs), filepath)
+            return docs
+        except Exception as exc:
+            logger.error("Failed to load '%s': %s", filepath, exc)
+            return []
+
+    def load_documents(self, paths: list[str]) -> list[Document]:
+        all_docs: list[Document] = []
+        for path in paths:
+            ok, reason = validate_file(
+                path,
+                allowed_extensions=cfg.app.allowed_extensions,
+                max_size_mb=cfg.app.max_file_size_mb,
+            )
+            if not ok:
+                logger.warning("Skipping '%s': %s", path, reason)
+                continue
+            chk = file_checksum(path)
+            if chk in self.ingested_checksums:
+                logger.info("Skipping duplicate file: %s", path)
+                continue
+            docs = self._load_single(path)
+            if docs:
+                self.ingested_checksums.add(chk)
+                all_docs.extend(docs)
+        return all_docs
+
+    # ── Chunking ──────────────────────────────────────────────────────────────
+
+    def chunk_documents(self, docs: list[Document]) -> list[Document]:
+        if not docs:
+            return []
+        with Timer() as t:
+            chunks = self._splitter.split_documents(docs)
+        logger.info(
+            "Produced %d chunks from %d documents in %s.", len(chunks), len(docs), t
+        )
+        return chunks
+
+    def update_splitter(self, chunk_size: int, chunk_overlap: int) -> None:
+        self._splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+            separators=cfg.chunking.separators,
+        )
+
+    # ── Vector store ──────────────────────────────────────────────────────────
+
+    def build_index(self, chunks: list[Document]) -> None:
+        if not chunks:
+            raise ValueError("Cannot build index: no chunks provided.")
+        self.load_embeddings()
+        logger.info("Building FAISS index from %d chunks…", len(chunks))
+        with Timer() as t:
+            self.vector_store = FAISS.from_documents(chunks, self.embeddings)
+        self.all_chunks = chunks
+        logger.info("FAISS index built in %s.", t)
+        self._save_index()
+
+    def add_documents_to_index(self, chunks: list[Document]) -> None:
+        """Incremental update — appends to an existing index."""
+        if not chunks:
+            return
+        self.load_embeddings()
+        if self.vector_store is None:
+            self.build_index(chunks)
+            return
+        logger.info("Incrementally adding %d chunks to existing index.", len(chunks))
+        self.vector_store.add_documents(chunks)
+        self.all_chunks.extend(chunks)
+        self._save_index()
+
+    def _save_index(self) -> None:
+        ensure_dir(cfg.retrieval.index_path)
+        self.vector_store.save_local(cfg.retrieval.index_path)
+        meta_path = cfg.retrieval.metadata_file
+        with open(meta_path, "wb") as f:
+            pickle.dump(
+                {
+                    "checksums": self.ingested_checksums,
+                    "chunks": self.all_chunks,
+                },
+                f,
+            )
+        logger.info("Index saved to '%s'.", cfg.retrieval.index_path)
+
+    def load_index(self) -> bool:
+        index_file = cfg.retrieval.index_file
+        meta_file = cfg.retrieval.metadata_file
+        if not (os.path.exists(index_file) and os.path.exists(meta_file)):
+            logger.info("No persisted index found at '%s'.", cfg.retrieval.index_path)
+            return False
+        self.load_embeddings()
+        try:
+            self.vector_store = FAISS.load_local(
+                cfg.retrieval.index_path,
+                self.embeddings,
+                allow_dangerous_deserialization=True,
+            )
+            with open(meta_file, "rb") as f:
+                meta = pickle.load(f)
+            self.ingested_checksums = meta.get("checksums", set())
+            self.all_chunks = meta.get("chunks", [])
+            logger.info(
+                "Index loaded: %d chunks, %d source files.",
+                len(self.all_chunks),
+                len(self.ingested_checksums),
+            )
+            return True
+        except Exception as exc:
+            logger.error("Failed to load index: %s", exc)
+            return False
+
+    # ── Retrieval ─────────────────────────────────────────────────────────────
+
+    def retrieve(
+        self,
+        query: str,
+        top_k: Optional[int] = None,
+    ) -> list[Document]:
+        if self.vector_store is None:
+            raise RuntimeError("Vector store is not initialised. Build or load an index first.")
+        k = top_k or cfg.retrieval.top_k
+        with Timer() as t:
+            results = self.vector_store.similarity_search(query, k=k)
+        logger.info(
+            "Retrieved %d chunks for query '%s…' in %s.",
+            len(results),
+            query[:60],
+            t,
+        )
+        return results
+
+    # ── Status ────────────────────────────────────────────────────────────────
+
+    @property
+    def is_ready(self) -> bool:
+        return self.vector_store is not None
+
+    @property
+    def doc_count(self) -> int:
+        return len(self.ingested_checksums)
+
+    @property
+    def chunk_count(self) -> int:
+        return len(self.all_chunks)