rag_engine.py

"""
RAG Engine for Sakhi.
Fetches top chunks from FAISS, deduplicates, and merges contiguous text blocks.
"""
import os
import pickle
import logging

from config import EMBED_MODEL, TOP_K_FETCH, TOP_K_MERGED, FAISS_INDEX_PATH, CHUNKS_PATH

logger = logging.getLogger(__name__)

class RAGEngine:
    def __init__(self):
        self.model = None
        self.index = None
        self.chunks = []
        self._initialized = False

    
    def initialize(self):

        if self._initialized:
            return

        from build_index import ensure_index_exists

        ensure_index_exists()

        try:

            import faiss
            import torch

            from sentence_transformers import SentenceTransformer

            device = "cuda" if torch.cuda.is_available() else "cpu"

            self.model = SentenceTransformer(
                EMBED_MODEL,
                device=device,
            )

            self.index = faiss.read_index(FAISS_INDEX_PATH)

            with open(CHUNKS_PATH, "rb") as f:
                self.chunks = pickle.load(f)

            logger.info(
                f"Loaded {len(self.chunks)} chunks."
            )

            self._initialized = True

        except Exception as e:

            logger.exception(e)

    def retrieve(self, query: str) -> list[dict]:
        self.initialize()
        if not self.chunks or self.index is None or self.model is None:
            return []

        try:
            import faiss
            # Fetch wider net of chunks initially
            query_embedding = self.model.encode([query], convert_to_numpy=True, normalize_embeddings=True)
            scores, indices = self.index.search(query_embedding, min(TOP_K_FETCH, len(self.chunks)))

            retrieved = []
            for score, idx in zip(scores[0], indices[0]):
                if 0 <= idx < len(self.chunks):
                    chunk = self.chunks[idx].copy()
                    chunk["score"] = float(score)
                    retrieved.append(chunk)

            # Intelligent Merging: Sort by source, then by chunk_idx
            # This ensures adjacent paragraphs in the PDF are merged into one readable block
            unique_chunks = {c['chunk_idx']: c for c in retrieved}.values()
            sorted_chunks = sorted(unique_chunks, key=lambda x: (x.get('source', ''), x.get('chunk_idx', 0)))

            merged_results = []
            curr = None

            for c in sorted_chunks:
                if not curr:
                    curr = c.copy()
                else:
                    # If chunks are from the same document and are adjacent/overlapping
                    if curr['source'] == c['source'] and (c['chunk_idx'] - curr['chunk_idx']) <= 2:
                        curr['text'] += "\n" + c['text']
                        curr['chunk_idx'] = c['chunk_idx'] # Move the end pointer
                    else:
                        merged_results.append(curr)
                        curr = c.copy()
            if curr:
                merged_results.append(curr)

            # Return top merged blocks
            return sorted(merged_results, key=lambda x: x.get('score', 0), reverse=True)[:TOP_K_MERGED]

        except Exception as e:
            logger.error(f"FAISS retrieval error: {e}")
            return []

    def format_context(self, chunks: list[dict]) -> str:
        if not chunks:
            return "No specific PDF context found. Please use general knowledge to explain."
        context_parts = []
        for i, chunk in enumerate(chunks, 1):
            source = chunk.get("source", "Unknown")
            context_parts.append(f"--- Document: {source} ---\n{chunk.get('text', '')}\n")
        return "\n".join(context_parts)