backend/rag.py

"""RAG engine: PyMuPDF4LLM -> chunks -> HF-API embeddings -> LanceDB.

Designed for HF Spaces Free CPU Basic:
    * No local embedding/LLM weights are downloaded.
    * LanceDB lives on ephemeral disk locally, but its files are mirrored into
        the linked HF Dataset so cold starts can restore the vector cache instead
        of rebuilding from PDFs every time.
"""
from __future__ import annotations

import logging
from pathlib import Path
from threading import Lock
from typing import Iterable, List, Optional

import pymupdf4llm
from llama_index.core import Document, StorageContext, VectorStoreIndex, Settings
from llama_index.core.node_parser import MarkdownNodeParser
from llama_index.core.schema import MetadataMode, NodeWithScore, TextNode
from llama_index.vector_stores.lancedb import LanceDBVectorStore

from config import EMBED_MODEL, HF_TOKEN, LANCEDB_PATH
from hf_embedding import SyncHuggingFaceInferenceEmbedding

logger = logging.getLogger(__name__)

TABLE_NAME = "documents"

# Filename is stored on every node so we can delete by source.
FILENAME_KEY = "source_filename"


class RagEngine:
    """Thread-safe singleton wrapper around the LanceDB-backed index."""

    def __init__(self) -> None:
        self._lock = Lock()
        self._index: Optional[VectorStoreIndex] = None
        self._vector_store: Optional[LanceDBVectorStore] = None

        if not HF_TOKEN:
            logger.warning(
                "HF_TOKEN is not set — embedding calls will fail. "
                "Set it in .env before uploading or querying documents."
            )
            Settings.embed_model = None  # type: ignore[assignment]
        else:
            Settings.embed_model = SyncHuggingFaceInferenceEmbedding(
                model_name=EMBED_MODEL,
                token=HF_TOKEN,
            )
        Settings.llm = None  # We call the LLM ourselves in main.py for streaming.
        Settings.node_parser = MarkdownNodeParser()

        Path(LANCEDB_PATH).mkdir(parents=True, exist_ok=True)
        self._vector_store = LanceDBVectorStore(
            uri=LANCEDB_PATH,
            table_name=TABLE_NAME,
            mode="overwrite" if not self._table_exists() else "append",
        )

    # ---------- internals ----------
    def _table_exists(self) -> bool:
        import lancedb

        try:
            db = lancedb.connect(LANCEDB_PATH)
            return TABLE_NAME in db.table_names()
        except Exception:  # noqa: BLE001
            return False

    def _ensure_index(self) -> VectorStoreIndex:
        if self._index is None:
            storage = StorageContext.from_defaults(vector_store=self._vector_store)
            if self._table_exists():
                self._index = VectorStoreIndex.from_vector_store(
                    vector_store=self._vector_store,
                    storage_context=storage,
                )
            else:
                self._index = VectorStoreIndex.from_documents(
                    [], storage_context=storage
                )
        return self._index

    @staticmethod
    def _pdf_to_documents(pdf_path: Path, filename: str) -> List[Document]:
        try:
            md_text = pymupdf4llm.to_markdown(str(pdf_path))
        except Exception as exc:  # noqa: BLE001
            logger.warning("Failed to parse %s: %s", filename, exc)
            return []
        if not md_text.strip():
            return []
        return [
            Document(
                text=md_text,
                metadata={FILENAME_KEY: filename},
                excluded_llm_metadata_keys=[FILENAME_KEY],
                excluded_embed_metadata_keys=[FILENAME_KEY],
            )
        ]

    # ---------- public API ----------
    def index_pdf(self, pdf_path: Path, filename: str) -> int:
        """Parse a single PDF and insert its chunks. Returns # nodes added."""
        with self._lock:
            # First, drop any existing nodes for this filename (re-uploads).
            self.delete_by_filename(filename, _locked=True)
            docs = self._pdf_to_documents(pdf_path, filename)
            if not docs:
                return 0
            index = self._ensure_index()
            nodes = Settings.node_parser.get_nodes_from_documents(docs)
            for n in nodes:
                n.metadata[FILENAME_KEY] = filename
            index.insert_nodes(nodes)
            return len(nodes)

    def index_many(self, items: Iterable[tuple[Path, str]]) -> int:
        total = 0
        for path, name in items:
            total += self.index_pdf(path, name)
        return total

    def delete_by_filename(self, filename: str, *, _locked: bool = False) -> int:
        """Remove all nodes with metadata.source_filename == filename."""
        def _do() -> int:
            import lancedb

            if not self._table_exists():
                return 0
            db = lancedb.connect(LANCEDB_PATH)
            tbl = db.open_table(TABLE_NAME)
            # LanceDB stores metadata as a struct column called "metadata".
            # Filter accesses nested fields with dot syntax.
            try:
                before = tbl.count_rows()
                tbl.delete(f"metadata.{FILENAME_KEY} = '{filename}'")
                removed = before - tbl.count_rows()
            except Exception as exc:  # noqa: BLE001
                logger.warning("delete_by_filename failed: %s", exc)
                removed = 0
            # Force the index to be rebuilt next query.
            self._index = None
            return removed

        if _locked:
            return _do()
        with self._lock:
            return _do()

    def retrieve(self, query: str, top_k: int = 4) -> List[NodeWithScore]:
        with self._lock:
            if not self._table_exists():
                return []
            index = self._ensure_index()
            retriever = index.as_retriever(similarity_top_k=top_k)
            return retriever.retrieve(query)

    @staticmethod
    def format_context(nodes: List[NodeWithScore]) -> str:
        if not nodes:
            return ""
        chunks = []
        for i, n in enumerate(nodes, 1):
            src = n.node.metadata.get(FILENAME_KEY, "unknown")
            text = n.node.get_content(metadata_mode=MetadataMode.NONE).strip()
            chunks.append(f"[{i}] (source: {src})\n{text}")
        return "\n\n---\n\n".join(chunks)