src/backend/rag.py

"""Retrieval-augmented generation: the one shared entry point.

Both the ``@app.api(name="chat")`` endpoint and the Gradio UI call ``answer()``
so there is a single grounding code path.

Per the hardware decision, query embedding **and** generation run inside one
``@spaces.GPU`` function (``_rag_generate``) — a single GPU allocation per turn,
not two separate bursts. FAISS search is CPU work but runs inside the same
allocation (it's cheap and avoids a second GPU round-trip). The corpus index is
prebuilt offline and only loaded here.
"""

from __future__ import annotations

import spaces

from .indexing.embedder import Embedder
from .indexing.store import VectorStore
from .llm.client import LLMClient, get_llm_client
from .llm.prompt import build_grounded_prompt
from .schemas import ChatMessage, ChatResponse, Chunk

_store: VectorStore | None = None
_llm: LLMClient | None = None

# Retrieval is run per knowledge source then fused, so each source contributes
# regardless of the other's similarity-score scale (directive = long English
# legal text; IDA = short salary summaries). Sources to retrieve from.
_SOURCES = ("directive", "lonstatistik")
K_PER_SOURCE = 5  # top-k pulled from each source before fusion
K_FINAL = 6       # chunks kept after fusion (passed to the LLM)
RRF_C = 60        # Reciprocal Rank Fusion constant (standard default)


def _rrf_fuse(
    ranked_lists: list[list[tuple[Chunk, float]]], k_final: int, c: int = RRF_C
) -> list[Chunk]:
    """Fuse per-source ranked lists with Reciprocal Rank Fusion.

    Rank-based and score-scale-agnostic: each chunk scores ``sum 1/(c + rank)``
    across the lists it appears in. Our sources are disjoint, so this cleanly
    interleaves the lists by rank. Returns the top ``k_final`` chunks.
    """
    score: dict[str, float] = {}
    keep: dict[str, Chunk] = {}
    for lst in ranked_lists:
        for rank, (chunk, _s) in enumerate(lst):
            score[chunk.id] = score.get(chunk.id, 0.0) + 1.0 / (c + rank + 1)
            keep[chunk.id] = chunk
    ordered = sorted(keep, key=lambda cid: score[cid], reverse=True)
    return [keep[cid] for cid in ordered[:k_final]]


def _get_store() -> VectorStore:
    """Load the persisted FAISS index once (lazy singleton).

    The corpus is embedded **offline** (``scripts/build_index.py`` on ZeroGPU)
    and persisted, so runtime never re-embeds it — only the query is embedded.
    Raises if no index is present.
    """
    global _store
    if _store is None:
        store = VectorStore()
        if not store.load():
            raise RuntimeError(
                "No FAISS index found at data/processed/index/. Build it first "
                "with `scripts/build_index.py` (on a GPU machine / the ZeroGPU "
                "Space) and commit the result."
            )
        _store = store
    return _store


def _citations(chunks: list[Chunk]) -> list[str]:
    """Dedup citation labels, preserving retrieval order."""
    return list(dict.fromkeys(c.citation for c in chunks if c.citation))


def _get_llm() -> LLMClient:
    global _llm
    if _llm is None:
        _llm = get_llm_client("default")
    return _llm


@spaces.GPU(duration=300)
def _embed_corpus(texts: list[str]):
    """Embed the full corpus on GPU — used by ``build_index()``."""
    return Embedder.get().embed_passages(texts)


def build_index() -> str:
    """Build the FAISS index in-process (admin trigger for ZeroGPU Spaces).

    Reuses the same ingest → embed → persist pipeline as
    ``scripts/build_index.py``, but runs through the live Gradio process so
    ``@spaces.GPU`` can allocate a real GPU on ZeroGPU.  Resets the cached
    store so the next chat query picks up the new index.
    """
    global _store
    from .indexing.ingest import build_corpus, load_directive_chunks, load_ida_chunks
    from .indexing.store import DEFAULT_INDEX_DIR

    n_dir = len(load_directive_chunks())
    n_ida = len(load_ida_chunks())
    chunks = build_corpus()

    vectors = _embed_corpus([c.text for c in chunks])
    dim = vectors.shape[1] if vectors.size else 0

    store = VectorStore()
    store.add(chunks, vectors)
    store.persist()

    _store = None

    msg = (
        f"Index built: {n_dir} directive + {n_ida} IDA = {len(chunks)} chunks "
        f"(dim={dim}) → {DEFAULT_INDEX_DIR}"
    )

    import os
    space_id = os.environ.get("SPACE_ID")
    if space_id:
        from huggingface_hub import HfApi

        HfApi(token=os.environ.get("HF_TOKEN")).upload_folder(
            folder_path=str(DEFAULT_INDEX_DIR),
            path_in_repo="data/processed/index",
            repo_id=space_id,
            repo_type="space",
            commit_message="Persist FAISS index (built via admin UI)",
        )
        msg += f"\nCommitted to {space_id} — index will survive restarts."
    else:
        msg += "\nNot on HF Spaces — commit data/processed/index/ manually."

    return msg


@spaces.GPU(duration=120)
def _warmup_gpu() -> bool:
    """Load + exercise the embedder and LLM inside one GPU allocation.

    Triggered early (when the user starts the onboarding wizard) so the weights
    are downloaded and the models resident by the time the chat is reached — the
    first real question then answers without the cold-start wait.
    """
    Embedder.get().embed_query("warm up")
    _get_store()  # FAISS load is CPU/cheap, but prime it too
    _get_llm().chat([ChatMessage(role="user", content="Hi")], max_new_tokens=1)
    return True


def warmup() -> dict:
    """Pre-warm the chat models; never raises (fire-and-forget from the UI)."""
    try:
        _warmup_gpu()
        return {"status": "ready"}
    except Exception as exc:  # noqa: BLE001 — warmup is best-effort
        return {"status": "error", "detail": str(exc)}


@spaces.GPU(duration=180)
def _retrieve_and_stream(query: str, lang: str, k: int, context: str | None = None):
    """Embed query, retrieve per source + fuse, and stream generation.

    One GPU allocation. Retrieves the top ``K_PER_SOURCE`` from each knowledge
    source independently, fuses with Reciprocal Rank Fusion, and keeps the top
    ``k`` fused chunks — so both the Directive and IDA statistics can surface
    even though their similarity scores live on different scales. ``context`` is
    the user's compact dashboard summary, injected into the prompt so answers can
    be specific to them. Yields ``(partial_reply, chunks)`` as tokens arrive so
    the UI renders incrementally.
    """
    qvec = Embedder.get().embed_query(query)
    store = _get_store()
    ranked = [store.search_by_source(qvec, s, K_PER_SOURCE) for s in _SOURCES]
    chunks = _rrf_fuse(ranked, k_final=k)
    prompt = build_grounded_prompt(query, chunks, lang, context=context)
    acc = ""
    for piece in _get_llm().stream(prompt):
        acc += piece
        yield acc, chunks


def answer_stream(
    messages: list[ChatMessage], lang: str = "en", k: int = K_FINAL,
    context: str | None = None,
):
    """Stream grounded answers as ``ChatResponse`` snapshots (growing reply)."""
    query = next((m.content for m in reversed(messages) if m.role == "user"), "")
    if not query.strip():
        yield ChatResponse(reply="Please ask a question.", citations=[])
        return
    for acc, chunks in _retrieve_and_stream(query, lang, k, context):
        yield ChatResponse(reply=acc, citations=_citations(chunks))


def answer(
    messages: list[ChatMessage], lang: str = "en", k: int = K_FINAL,
    context: str | None = None,
) -> ChatResponse:
    """Non-streaming answer (collapses the stream to its final snapshot)."""
    resp = ChatResponse(reply="", citations=[])
    for resp in answer_stream(messages, lang=lang, k=k, context=context):
        pass
    return resp