rag.py

"""
rag.py — Plexi RAG Engine
=========================
Handles everything related to the LlamaIndex vector index:
  - Downloading the pre-built index from GitHub
  - Loading HuggingFace sentence-transformer embeddings
  - Embedding queries and retrieving top-k chunks scoped by semester + subject
  - Extracting text from PDFs for full-context fallback
  - Formatting retrieved chunks for the LLM system prompt
"""

import io
import os
import shutil          # API-BUG-7: temp-dir cleanup after index is in memory
import tempfile
from pathlib import Path
from typing import TypedDict  # API-BUG-4: explicit typed return from retrieve_chunks

import requests

# ---------------------------------------------------------------------------
# Optional LlamaIndex — graceful degradation if not installed
# ---------------------------------------------------------------------------
try:
    from llama_index.core import Settings, StorageContext, load_index_from_storage
    from llama_index.embeddings.huggingface import HuggingFaceEmbedding

    LLAMA_INDEX_AVAILABLE = True
except ImportError:
    LLAMA_INDEX_AVAILABLE = False

# API-BUG-5: MetadataFilters let the vector store do scope-filtering internally,
# avoiding the over-fetch window that could miss relevant chunks.
try:
    from llama_index.core.vector_stores.types import MetadataFilter, MetadataFilters

    METADATA_FILTERS_AVAILABLE = True
except ImportError:
    try:
        from llama_index.core.vector_stores import MetadataFilter, MetadataFilters

        METADATA_FILTERS_AVAILABLE = True
    except ImportError:
        METADATA_FILTERS_AVAILABLE = False

try:
    import PyPDF2

    PYPDF2_AVAILABLE = True
except ImportError:
    PYPDF2_AVAILABLE = False

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
MATERIALS_REPO = os.getenv("MATERIALS_REPO", "KunalGupta25/plexi-materials")
MANIFEST_BRANCH = os.getenv("MANIFEST_BRANCH", "main")
EMBED_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"

INDEX_FILES = [
    "default__vector_store.json",
    "docstore.json",
    "graph_store.json",
    "image__vector_store.json",
    "index_store.json",
]

DEFAULT_TOP_K = 5


# ---------------------------------------------------------------------------
# Typed chunk return (API-BUG-4)
# ---------------------------------------------------------------------------
class ChunkDict(TypedDict):
    text: str
    score: float | None
    filename: str | None
    subject: str | None


# ---------------------------------------------------------------------------
# Index loading (called once at FastAPI startup via asyncio.to_thread)
# ---------------------------------------------------------------------------

def load_index():
    """
    Download the pre-built LlamaIndex from the materials repo and return a
    VectorStoreIndex ready for querying.

    Returns (index, error_msg). index is None if loading failed.

    API-BUG-7: The temp directory is always removed in the finally block.
    After load_index_from_storage() returns, all data is in memory — the
    files on disk are no longer needed.
    """
    if not LLAMA_INDEX_AVAILABLE:
        return None, "llama-index-core is not installed."

    index_base_url = (
        f"https://raw.githubusercontent.com/{MATERIALS_REPO}/{MANIFEST_BRANCH}/index"
    )
    index_dir = tempfile.mkdtemp(prefix="plexi_index_")

    try:
        # Download each index shard from GitHub
        for filename in INDEX_FILES:
            url = f"{index_base_url}/{filename}"
            try:
                resp = requests.get(url, timeout=30)
                resp.raise_for_status()
                with open(os.path.join(index_dir, filename), "wb") as fh:
                    fh.write(resp.content)
            except Exception as err:
                return None, f"Failed to download index file '{filename}': {err}"

        # Build the in-memory index from the downloaded shards
        try:
            embed_model = HuggingFaceEmbedding(model_name=EMBED_MODEL_ID)
            Settings.embed_model = embed_model
            Settings.llm = None

            storage_ctx = StorageContext.from_defaults(persist_dir=index_dir)
            index = load_index_from_storage(storage_ctx)
            return index, None
        except Exception as err:
            return None, f"Failed to load index from storage: {err}"

    finally:
        # API-BUG-7: always wipe the temp dir — data is now in RAM.
        shutil.rmtree(index_dir, ignore_errors=True)


# ---------------------------------------------------------------------------
# Retrieval
# ---------------------------------------------------------------------------

def _matches_scope(node, semester: str, subject: str) -> bool:
    """Return True when a retrieved node belongs to the active semester + subject."""
    metadata = getattr(node.node, "metadata", {}) or {}
    return (
        metadata.get("semester") == semester
        and metadata.get("subject") == subject
    )


def retrieve_chunks(
    index,
    query: str,
    semester: str,
    subject: str,
    top_k: int = DEFAULT_TOP_K,
) -> list[ChunkDict]:
    """
    Embed the query, retrieve top-k chunks from the index scoped to the
    given semester + subject.

    API-BUG-4: Returns list[ChunkDict] (TypedDict) instead of list[dict] so
    type errors surface at the source rather than silently at Pydantic
    serialization time.

    API-BUG-5: Primary path uses MetadataFilters so the vector store does the
    scope-gating internally — no risk of the over-fetch window failing to reach
    chunks that belong to the active subject.  Falls back to the generous
    over-fetch + manual filter approach when MetadataFilters are unavailable
    (e.g., older llama-index-core builds).
    """
    if index is None:
        return []

    try:
        if METADATA_FILTERS_AVAILABLE:
            # Primary: vector store filters by metadata at query time.
            filters = MetadataFilters(
                filters=[
                    MetadataFilter(key="semester", value=semester),
                    MetadataFilter(key="subject", value=subject),
                ]
            )
            retriever = index.as_retriever(similarity_top_k=top_k, filters=filters)
            nodes = retriever.retrieve(query)
        else:
            # Fallback: over-fetch (generous 10× window) + manual scope filter.
            retriever = index.as_retriever(similarity_top_k=max(top_k * 10, 50))
            nodes = retriever.retrieve(query)
            nodes = [n for n in nodes if _matches_scope(n, semester, subject)]

        return [
            ChunkDict(
                text=node.node.get_content(),
                score=round(float(node.score), 4) if node.score is not None else None,
                filename=(getattr(node.node, "metadata", {}) or {}).get("filename"),
                subject=(getattr(node.node, "metadata", {}) or {}).get("subject"),
            )
            for node in nodes[:top_k]
        ]
    except Exception as err:
        print(f"Retrieval error: {err}")
        return []


# ---------------------------------------------------------------------------
# Context formatting (for system prompt injection)
# ---------------------------------------------------------------------------

def format_context(chunks: list[ChunkDict]) -> str:
    """Format retrieved chunks as a numbered block for the LLM system prompt."""
    if not chunks:
        return "(No relevant context retrieved for this query.)"
    parts = []
    for i, chunk in enumerate(chunks, start=1):
        score_info = f"  [relevance: {chunk['score']}]" if chunk.get("score") else ""
        source = chunk.get("filename") or chunk.get("subject") or "Unknown source"
        parts.append(
            f"--- Chunk {i} | {source}{score_info} ---\n{chunk['text']}\n"
        )
    return "\n".join(parts)


# ---------------------------------------------------------------------------
# PDF text extraction
# NOTE: read_pdf_text() is currently unused by the API routes — retained for
# future use (CQ-4). load_embed_model() removed — was dead code (API-BUG-6).
# ---------------------------------------------------------------------------

def read_pdf_text(pdf_bytes: bytes) -> str:
    """Extract plain text from PDF bytes. Returns empty string on failure."""
    if not PYPDF2_AVAILABLE:
        return ""
    text_parts = []
    try:
        reader = PyPDF2.PdfReader(io.BytesIO(pdf_bytes))
        for page in reader.pages:
            try:
                page_text = page.extract_text()
                if page_text:
                    # Sanitise surrogate pairs that can appear in some PDFs
                    filtered = page_text.encode("utf-16", "surrogatepass").decode(
                        "utf-16", "ignore"
                    )
                    text_parts.append(filtered)
            except Exception:
                pass
    except Exception:
        return pdf_bytes.decode("utf-8", errors="ignore") if pdf_bytes else ""
    return "\n".join(text_parts)