src/retrieval/rag_pipeline.py

"""
RAG Pipeline — Retrieval-Augmented Generation
Uses FAISS vector store + HuggingFace embeddings + local LLM for generation.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import numpy as np

# Lazy imports — only loaded when needed to keep startup fast
_faiss = None
_SentenceTransformer = None
_pipeline_hf = None


def _get_faiss():
    global _faiss
    if _faiss is None:
        import faiss
        _faiss = faiss
    return _faiss


def _get_encoder(model_name: str):
    global _SentenceTransformer
    if _SentenceTransformer is None:
        from sentence_transformers import SentenceTransformer
        _SentenceTransformer = SentenceTransformer
    return _SentenceTransformer(model_name)


class DocumentChunk:
    """A single chunk of text from the knowledge base."""

    def __init__(self, text: str, metadata: dict):
        self.text = text
        self.metadata = metadata  # {source, title, authors, year, domain}

    def to_dict(self) -> dict:
        return {"text": self.text, "metadata": self.metadata}

    @classmethod
    def from_dict(cls, d: dict) -> "DocumentChunk":
        return cls(d["text"], d["metadata"])


class RAGPipeline:
    """
    Full RAG pipeline:
      1. Document ingestion + chunking
      2. Embedding with SentenceTransformers
      3. FAISS vector index for semantic search
      4. HuggingFace LLM for answer generation
    """

    INDEX_PATH = Path("data/index/faiss.index")
    CHUNKS_PATH = Path("data/index/chunks.json")
    DOCS_PATH = Path("data/papers")

    def __init__(
        self,
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
        llm_model: str = "google/flan-t5-base",
    ):
        self.embedding_model_name = embedding_model
        self.llm_model_name = llm_model
        self.encoder = None
        self.index = None
        self.chunks: list[DocumentChunk] = []
        self._llm = None

    # ── Encoder ──────────────────────────────────────────────────────────────

    def _load_encoder(self):
        if self.encoder is None:
            self.encoder = _get_encoder(self.embedding_model_name)

    def _embed(self, texts: list[str]) -> np.ndarray:
        self._load_encoder()
        return self.encoder.encode(texts, normalize_embeddings=True,
                                   show_progress_bar=False)

    # ── Index management ─────────────────────────────────────────────────────

    def load_or_build_index(self):
        """Load existing FAISS index or build from documents."""
        if self.INDEX_PATH.exists() and self.CHUNKS_PATH.exists():
            self._load_index()
        else:
            self._build_index()

    def _load_index(self):
        faiss = _get_faiss()
        self.index = faiss.read_index(str(self.INDEX_PATH))
        with open(self.CHUNKS_PATH) as f:
            self.chunks = [DocumentChunk.from_dict(d) for d in json.load(f)]

    def _build_index(self):
        """Ingest documents, embed, and build FAISS index."""
        from src.ingestion.document_loader import load_all_documents

        self.INDEX_PATH.parent.mkdir(parents=True, exist_ok=True)

        # Load and chunk documents
        self.chunks = load_all_documents(self.DOCS_PATH)

        if not self.chunks:
            # Fall back to built-in knowledge base
            from src.ingestion.builtin_knowledge import get_builtin_chunks
            self.chunks = get_builtin_chunks()

        # Embed all chunks
        texts = [c.text for c in self.chunks]
        embeddings = self._embed(texts).astype("float32")

        # Build FAISS flat L2 index
        faiss = _get_faiss()
        dim = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dim)  # Inner product (cosine on normalized)
        self.index.add(embeddings)

        # Persist
        faiss.write_index(self.index, str(self.INDEX_PATH))
        with open(self.CHUNKS_PATH, "w") as f:
            json.dump([c.to_dict() for c in self.chunks], f, indent=2)

    # ── Retrieval ─────────────────────────────────────────────────────────────

    def retrieve(self, query: str, top_k: int = 3) -> list[dict]:
        """Retrieve top-k most relevant document chunks."""
        q_embed = self._embed([query]).astype("float32")
        scores, indices = self.index.search(q_embed, top_k)

        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx < 0:
                continue
            chunk = self.chunks[idx]
            results.append({
                "text": chunk.text,
                "score": float(score),
                "source": chunk.metadata.get("source", "Unknown"),
                "title": chunk.metadata.get("title", ""),
                "authors": chunk.metadata.get("authors", ""),
                "year": chunk.metadata.get("year", ""),
                "domain": chunk.metadata.get("domain", ""),
            })
        return results

    # ── Generation ────────────────────────────────────────────────────────────

    def _load_llm(self):
        if self._llm is None:
            from transformers import pipeline as hf_pipeline
            self._llm = hf_pipeline(
                "text2text-generation",
                model=self.llm_model_name,
                max_new_tokens=512,
                temperature=0.3,
            )

    def _build_prompt(self, query: str, contexts: list[dict]) -> str:
        context_text = "\n\n".join([
            f"[Source {i+1}: {c['title']} ({c['year']})]\n{c['text']}"
            for i, c in enumerate(contexts)
        ])
        return f"""You are an expert AI assistant specializing in railroad condition monitoring,
distributed acoustic sensing (DAS), and AI-driven manufacturing systems.

Answer the question based on the provided research context. Be precise and technical.
If the answer is not in the context, say so clearly.

Context:
{context_text}

Question: {query}

Answer:"""

    def query(
        self,
        question: str,
        top_k: int = 3,
        temperature: float = 0.3,
    ) -> dict[str, Any]:
        """Full RAG pipeline: retrieve → prompt → generate."""
        sources = self.retrieve(question, top_k=top_k)

        if not sources:
            return {
                "answer": "No relevant documents found. Please add documents to the knowledge base.",
                "sources": [],
            }

        prompt = self._build_prompt(question, sources)

        try:
            self._load_llm()
            result = self._llm(prompt, temperature=temperature)
            answer = result[0]["generated_text"].strip()
        except Exception as e:
            # Graceful fallback — return context directly
            answer = self._fallback_answer(question, sources)

        return {"answer": answer, "sources": sources}

    def _fallback_answer(self, question: str, sources: list[dict]) -> str:
        """Context-only fallback when LLM is unavailable."""
        top = sources[0]
        return (
            f"**Based on retrieved context from '{top['title']}' ({top['year']}):**\n\n"
            f"{top['text']}\n\n"
            f"*Retrieved {len(sources)} relevant passages. "
            f"For full LLM-generated answers, ensure the model is loaded.*"
        )