rag_engine.py

"""
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)