src/rag.py

import os
import csv
import glob
from pypdf import PdfReader
from llama_index.core import VectorStoreIndex, Settings
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import Document
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# -------------------------------------------------------------------------
# LlamaIndex configuration
# - Embeddings: sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2 (Supports 50+ languages)
# - LLM: disabled — we manage the LLM ourselves in src/llm.py
# -------------------------------------------------------------------------
Settings.embed_model = HuggingFaceEmbedding(
    model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
)
Settings.llm = None

_index: VectorStoreIndex | None = None


# ── File loaders ─────────────────────────────────────────────────────────

def _load_csv(filepath: str) -> list[Document]:
    """Loads a CSV with 'Question' (and optionally 'Answer') columns."""
    docs = []
    try:
        with open(filepath, newline="", encoding="utf-8") as fh:
            reader = csv.DictReader(fh)
            if reader.fieldnames is None or "Question" not in reader.fieldnames:
                print(f"[RAG] Skipping {filepath}: no 'Question' column.")
                return docs
            has_answer = "Answer" in (reader.fieldnames or [])
            for row in reader:
                question = row.get("Question", "").strip()
                answer = row.get("Answer", "").strip() if has_answer else ""
                if question:
                    text = f"Q: {question}\nA: {answer}" if answer else f"Q: {question}"
                    docs.append(Document(text=text, metadata={"source": filepath}))
    except Exception as exc:
        print(f"[RAG] Error reading CSV {filepath}: {exc}")
    return docs


def _load_pdf(filepath: str) -> list[Document]:
    """Extracts text from each page of a PDF as separate documents."""
    docs = []
    try:
        reader = PdfReader(filepath)
        basename = os.path.basename(filepath)
        for i, page in enumerate(reader.pages):
            text = page.extract_text()
            if text and text.strip():
                docs.append(Document(
                    text=text.strip(),
                    metadata={"source": basename, "page": i + 1},
                ))
        print(f"[RAG] Loaded {len(docs)} pages from {basename}")
    except Exception as exc:
        print(f"[RAG] Error reading PDF {filepath}: {exc}")
    return docs


def _load_text(filepath: str) -> list[Document]:
    """Loads a plain text file as a single document."""
    docs = []
    try:
        with open(filepath, "r", encoding="utf-8") as fh:
            text = fh.read().strip()
            if text:
                docs.append(Document(text=text, metadata={"source": filepath}))
    except Exception as exc:
        print(f"[RAG] Error reading text file {filepath}: {exc}")
    return docs


# ── Vector store init ────────────────────────────────────────────────────

def init_vector_store(data_path: str = "data") -> None:
    """
    Scans the data directory and loads ALL supported files:
      - .csv  → FAQ-style (Question/Answer columns)
      - .pdf  → RBI guidelines, circulars, etc.
      - .txt  → Plain text documents

    Builds an in-memory vector index from all loaded content.
    """
    global _index

    # Support both file path and directory path for backward compatibility
    if os.path.isfile(data_path):
        data_dir = os.path.dirname(data_path)
    else:
        data_dir = data_path

    if not os.path.isdir(data_dir):
        print(f"[RAG] Warning: data directory not found at '{data_dir}'. Vector store empty.")
        return

    documents: list[Document] = []

    # Auto-discover all supported files
    for pattern, loader in [("*.csv", _load_csv), ("*.pdf", _load_pdf), ("*.txt", _load_text)]:
        for filepath in sorted(glob.glob(os.path.join(data_dir, pattern))):
            loaded = loader(filepath)
            documents.extend(loaded)

    if not documents:
        print("[RAG] Warning: No documents loaded from any file.")
        return

    try:
        splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50)
        nodes = splitter.get_nodes_from_documents(documents)
        _index = VectorStoreIndex(nodes)
        print(f"[RAG] Loaded {len(documents)} documents → {len(nodes)} chunks into vector store.")
    except Exception as exc:
        print(f"[RAG] Error building vector index: {exc}")


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

def retrieve_context(query: str, k: int = 5) -> str:
    """
    Returns the top-k most relevant passages for a query string.
    Falls back to an informative message if the index is not ready.
    """
    if not _index:
        return "Knowledge base not available."

    if not query or not query.strip():
        return "Knowledge base not available."

    try:
        retriever = _index.as_retriever(similarity_top_k=k)
        results = retriever.retrieve(query)
        if not results:
            return "No relevant information found in the knowledge base."
        return "\n\n---\n\n".join(node.get_content() for node in results)
    except Exception as exc:
        print(f"[RAG] Error during retrieval: {exc}")
        return "Knowledge base lookup failed."


def retrieve_context_multi(english_query: str, hindi_text: str, k: int = 15) -> str:
    """
    Searches using BOTH the English query and the original Hindi text,
    deduplicates, and returns the top-k most relevant passages.
    """
    if not _index:
        return "Knowledge base not available."

    seen_content: set[str] = set()
    all_results = []

    try:
        retriever = _index.as_retriever(similarity_top_k=k)

        # Search with English query
        if english_query and english_query.strip():
            for node in retriever.retrieve(english_query):
                content = node.get_content()
                if content not in seen_content:
                    seen_content.add(content)
                    all_results.append((node.get_score(), content))

        # Search with Hindi text (backup)
        if hindi_text and hindi_text.strip():
            for node in retriever.retrieve(hindi_text):
                content = node.get_content()
                if content not in seen_content:
                    seen_content.add(content)
                    all_results.append((node.get_score(), content))

        if not all_results:
            return "No relevant information found in the knowledge base."

        all_results.sort(key=lambda x: x[0], reverse=True)
        top_results = [content for _, content in all_results[:k]]
        print(f"[RAG] Retrieved {len(top_results)} unique passages from dual query")
        return "\n\n---\n\n".join(top_results)

    except Exception as exc:
        print(f"[RAG] Error during retrieval: {exc}")
        return "Knowledge base lookup failed."