rag_engine.py

from __future__ import annotations
import html
import os
import re
import threading
from dataclasses import dataclass

IS_HF_SPACE = bool(os.getenv("SPACE_ID"))

if IS_HF_SPACE:
    import spaces
else:
    
    class _SpacesShim:
        @staticmethod
        def GPU(duration: int = 60, **_kwargs):
            def decorator(fn):
                return fn
            return decorator

    spaces = _SpacesShim()

import faiss
import numpy as np
import pandas as pd
import torch
from huggingface_hub import InferenceClient, hf_hub_download
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer

from config import (
    EMBEDDING_MODEL_ID,
    FAR_DATASET_FILE,
    FAR_DATASET_REPO,
    FAR_SOURCE_URL,
    HF_INFERENCE_MODEL,
    HF_TOKEN,
    LLM_BACKEND,
    LOCAL_MODEL_ID,
    MAX_CONTEXT_CHARS,
    SYSTEM_PROMPT,
)

from document_loader import Chunk, infer_far_section


# search results
@dataclass
class SearchHit:
    """
    One retrieved RAG passage

    chunk:
        Original document chunk

    score:
        FAISS cosine-similarity score
    """

    chunk: Chunk
    score: float


# Knowledge base
class KnowledgeBase:
    """
    in-memory FAISS vector knowledge base
    """

    def __init__(self) -> None:
        self.chunks: list[Chunk] = []
        self.index: faiss.IndexFlatIP | None = None
        self.dimension: int | None = None

    def add(self, chunks: list[Chunk]) -> int:
        """
        Embed and add chunks to the FAISS index

        Returns the number of chunks added
        """

        if not chunks:
            return 0

        embeddings = embed_passages(
            [chunk.text for chunk in chunks]
        )

        # building the FAISS index the first time documents are added
        if self.index is None:
            self.dimension = int(embeddings.shape[1])

            self.index = faiss.IndexFlatIP(
                self.dimension
            )

        # prevent accidental mixing of embeddings with different vector dimensions
        if embeddings.shape[1] != self.dimension:
            raise ValueError(
                "Embedding dimension changed during this session"
            )

        self.index.add(embeddings)

        self.chunks.extend(chunks)

        return len(chunks)

    def clear(self) -> None:
        """Remove all indexed documents"""

        self.chunks.clear()
        self.index = None
        self.dimension = None

    def search(
        self,
        query: str,
        top_k: int = 6,
    ) -> list[SearchHit]:
        """
        Retrieving the most relevant passages for a question
        """

        if self.index is None or not self.chunks:
            return []

        query = query.strip()

        if not query:
            return []

        # convert question into an embedding
        query_vec = embed_query(query)

        # never request more results than exist
        k = min(
            max(1, int(top_k)),
            len(self.chunks),
        )

        # retrieve extra candidates because we perform lightweight deduplication afterward
        candidate_count = min(
            k * 2,
            len(self.chunks),
        )

        scores, indices = self.index.search(
            query_vec,
            candidate_count,
        )

        hits: list[SearchHit] = []

        # avoid repeatedly returning effectively identical source/page/section combinations
        seen: set[
            tuple[str, int | None, str | None]
        ] = set()

        for score, idx in zip(
            scores[0],
            indices[0],
        ):
            if idx < 0:
                continue

            chunk = self.chunks[int(idx)]

            fingerprint = (
                chunk.source_name,
                chunk.page,
                chunk.section,
            )

            if fingerprint in seen:
                continue

            seen.add(fingerprint)

            hits.append(
                SearchHit(
                    chunk=chunk,
                    score=float(score),
                )
            )

            if len(hits) >= k:
                break

        return hits

    @property
    def document_count(self) -> int:
        """no. of unique source documents"""

        return len(
            {
                chunk.source_name
                for chunk in self.chunks
            }
        )

    @property
    def chunk_count(self) -> int:
        """total number of chunks in the knowledge base"""

        return len(self.chunks)


# Global model caches

_EMBEDDER: SentenceTransformer | None = None
_EMBEDDER_LOCK = threading.Lock()

_LOCAL_MODEL = None
_LOCAL_TOKENIZER = None
_LOCAL_MODEL_LOCK = threading.Lock()


# Embedding model

def get_embedding_device() -> str:
    """
    return the device using embedding model
    """

    if IS_HF_SPACE:
        return "cuda"

    return "cuda" if torch.cuda.is_available() else "cpu"


def _load_embedder() -> SentenceTransformer:
    """
    embedding model
    """

    device = get_embedding_device()

    print(
        f"[embedding] Loading {EMBEDDING_MODEL_ID} on {device}",
        flush=True,
    )

    model = SentenceTransformer(
        EMBEDDING_MODEL_ID,
        device=device,
    )

    print(
        "[embedding] Model ready.",
        flush=True,
    )

    return model


if IS_HF_SPACE:
    _EMBEDDER = _load_embedder()


def get_embedder() -> SentenceTransformer:
    """
    return the shared embedding model

    """

    global _EMBEDDER

    if _EMBEDDER is not None:
        return _EMBEDDER

    with _EMBEDDER_LOCK:
        if _EMBEDDER is None:
            _EMBEDDER = _load_embedder()

    return _EMBEDDER


def _encode(
    texts: list[str],
    *,
    batch_size: int = 16,
) -> np.ndarray:
    """
    encode text into normalized float32 vectors for FAISS
    """

    if not texts:
        raise ValueError(
            "Cannot embed an empty list of texts"
        )

    model = get_embedder()

    vectors = model.encode(
        texts,
        batch_size=int(batch_size),
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )

    vectors = vectors.astype("float32")

    # return ordinary CPU/NumPy data across the ZeroGPU process boundary
    return np.ascontiguousarray(vectors)


@spaces.GPU(duration=120)
def embed_passages(
    texts: list[str],
) -> np.ndarray:
    """
    embed document/FAR passages

    """

    return _encode(
        texts,
        batch_size=16,
    )


@spaces.GPU(duration=30)
def embed_query(
    query: str,
) -> np.ndarray:
    """
    Embed a retrieval query on ZeroGPU
    """

    clean = query.strip()

    if not clean:
        raise ValueError(
            "Query cannot be empty"
        )

    model = get_embedder()

    prompts = getattr(
        model,
        "prompts",
        {},
    )

    if "query" in prompts:
        vectors = model.encode(
            [clean],
            prompt_name="query",
            batch_size=1,
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False,
        )

        vectors = vectors.astype(
            "float32"
        )

        return np.ascontiguousarray(
            vectors
        )

    instructed = (
        "Represent this sentence for searching "
        f"relevant passages: {clean}"
    )

    return _encode(
        [instructed],
        batch_size=1,
    )


# FAR Hugging Face dataset

def _normalize_chunk_id(
    raw_id,
    fallback: int,
) -> str:
    """
    Convert CSV Chunk ID values into clean strings.

    Examples:
        12.0 -> "12"
        "FAR-52.212" -> "FAR-52.212"
    """

    if pd.isna(raw_id):
        return str(fallback)

    value = str(raw_id).strip()

    try:
        number = float(value)

        if number.is_integer():
            return str(int(number))

    except ValueError:
        pass

    return value or str(fallback)


def load_far_dataset(
    max_rows: int = 2500,
) -> list[Chunk]:
    """
    Download and load the FAR chunk corpus from hf
    """

    max_rows = int(max_rows)

    if max_rows <= 0:
        raise ValueError(
            "max_rows must be greater than zero."
        )

    print(
        "[FAR] Downloading starter corpus...",
        flush=True,
    )

    print(
        f"[FAR] Repository: {FAR_DATASET_REPO}",
        flush=True,
    )

    print(
        f"[FAR] File: {FAR_DATASET_FILE}",
        flush=True,
    )

    try:

        local_path = hf_hub_download(
            repo_id=FAR_DATASET_REPO,
            filename=FAR_DATASET_FILE,
            repo_type="dataset",
            token=HF_TOKEN,
        )

    except Exception as exc:

        raise RuntimeError(
            "Unable to download the FAR dataset from "
            "Hugging Face.\n\n"
            f"Repository: {FAR_DATASET_REPO}\n"
            f"File: {FAR_DATASET_FILE}\n"
            f"Error: {exc}"
        ) from exc

    print(
        f"[FAR] Downloaded to: {local_path}",
        flush=True,
    )

    try:

        frame = pd.read_csv(
            local_path
        )

    except Exception as exc:

        raise RuntimeError(
            "The FAR CSV was downloaded but could "
            f"not be read: {exc}"
        ) from exc

    required_column = "Chunk Text"

    if required_column not in frame.columns:

        raise ValueError(
            "Expected the FAR chunk file to contain "
            f"'{required_column}'. "
            f"Found columns: {list(frame.columns)}"
        )

    # remove records without usable text
    frame = frame.dropna(
        subset=[required_column]
    )

    # limit initial corpus size for local testing
    frame = frame.head(
        max_rows
    )

    chunks: list[Chunk] = []

    for row_number, (_, row) in enumerate(
        frame.iterrows(),
        start=1,
    ):

        text = str(
            row[required_column]
        ).strip()

        if not text:
            continue

        raw_id = row.get(
            "Chunk ID",
            row_number,
        )

        chunk_id = _normalize_chunk_id(
            raw_id,
            row_number,
        )

        # try to identify the FAR section from the chunk text itself
        section = infer_far_section(
            text
        )

        chunk = Chunk(
            chunk_id=chunk_id,
            text=text,
            source_name=(
                "Federal Acquisition Regulation"
            ),
            source_type=(
                "Hugging Face FAR dataset"
            ),
            section=section,
            url=FAR_SOURCE_URL,
            metadata={
                "dataset_repo": FAR_DATASET_REPO,
                "dataset_file": FAR_DATASET_FILE,
                "source": "Federal Acquisition Regulation",
            },
        )

        chunks.append(chunk)

    if not chunks:
        raise RuntimeError(
            "The FAR dataset loaded successfully, "
            "but no usable text chunks were found."
        )

    print(
        f"[FAR] Loaded {len(chunks):,} chunks.",
        flush=True,
    )

    return chunks


# RAG context construction
def build_context(
    hits: list[SearchHit],
    max_chars: int = MAX_CONTEXT_CHARS,
) -> str:
    """
    convert retrieved passages into the context supplied
    to the language model

    each passage receives a stable source identifier:
        [S1]
        [S2]
        [S3]
        ...
    """

    blocks: list[str] = []

    used_chars = 0

    for source_number, hit in enumerate(
        hits,
        start=1,
    ):

        chunk = hit.chunk

        header = [
            f"[S{source_number}]",
            f"Source: {chunk.source_name}",
            f"Citation: {chunk.citation}",
        ]

        if chunk.url:
            header.append(
                f"Official/source URL: {chunk.url}"
            )

        block = (
            "\n".join(header)
            + "\nRetrieved text:\n"
            + chunk.text.strip()
        )

        # stop adding context once the configured character budget has been reached
        if (
            used_chars + len(block) > max_chars
            and blocks
        ):
            break

        blocks.append(block)

        used_chars += len(block)

    return "\n\n---\n\n".join(
        blocks
    )


# Local LM

def _get_local_model():
    """
    loading the local Hugging Face language model once

    (CUDA is used automatically when available)
    """

    if IS_HF_SPACE:
        raise RuntimeError(
            "Local LLM fallback is disabled on this ZeroGPU Space. "
            "Use LLM_BACKEND=hf_api with HF_TOKEN configured as a Space Secret."
        )

    global _LOCAL_MODEL
    global _LOCAL_TOKENIZER

    if (
        _LOCAL_MODEL is not None
        and _LOCAL_TOKENIZER is not None
    ):
        return (
            _LOCAL_MODEL,
            _LOCAL_TOKENIZER,
        )

    with _LOCAL_MODEL_LOCK:

        if (
            _LOCAL_MODEL is not None
            and _LOCAL_TOKENIZER is not None
        ):
            return (
                _LOCAL_MODEL,
                _LOCAL_TOKENIZER,
            )

        print(
            f"[LLM] Loading tokenizer: "
            f"{LOCAL_MODEL_ID}",
            flush=True,
        )

        tokenizer = AutoTokenizer.from_pretrained(
            LOCAL_MODEL_ID,
            token=HF_TOKEN,
        )

        if tokenizer.pad_token is None:
            tokenizer.pad_token = (
                tokenizer.eos_token
            )

        if torch.cuda.is_available():

            dtype = (
                torch.bfloat16
                if torch.cuda.is_bf16_supported()
                else torch.float16
            )

            device = "cuda"

        else:

            dtype = torch.float32
            device = "cpu"

        print(
            f"[LLM] Loading model: "
            f"{LOCAL_MODEL_ID}",
            flush=True,
        )

        print(
            f"[LLM] Device: {device}",
            flush=True,
        )

        print(
            f"[LLM] Dtype: {dtype}",
            flush=True,
        )

        model = AutoModelForCausalLM.from_pretrained(
            LOCAL_MODEL_ID,
            dtype=dtype,
            low_cpu_mem_usage=True,
            token=HF_TOKEN,
        )

        model = model.to(
            device
        )

        model.eval()

        _LOCAL_MODEL = model
        _LOCAL_TOKENIZER = tokenizer

        print(
            "[LLM] Local model ready.",
            flush=True,
        )

    return (
        _LOCAL_MODEL,
        _LOCAL_TOKENIZER,
    )


# HF hosted inference

def _generate_hf_api(
    messages: list[dict],
    temperature: float,
    max_new_tokens: int,
) -> str:
    """
    generate an answer through Hugging Face Inference Providers

    """

    if not HF_TOKEN:
        raise RuntimeError(
            "HF_TOKEN is not configured for the "
            "Hugging Face Inference backend."
        )

    client = InferenceClient(
        api_key=HF_TOKEN,
        provider="together",
    )

    response = client.chat_completion(
        model=HF_INFERENCE_MODEL,
        messages=messages,
        max_tokens=int(max_new_tokens),
        temperature=max(
            0.7,
            float(temperature),
        ),
        top_p=0.8,
        presence_penalty=1.5,
        extra_body={
            "top_k": 20,
            "chat_template_kwargs": {
                "enable_thinking": False,
            },
        },
    )

    message = response.choices[0].message
    content = message.content

    if not content:
        reasoning = getattr(
            message,
            "reasoning",
            None,
        )

        if reasoning:
            raise RuntimeError(
                "The model returned reasoning but no final answer. "
                "Increase the answer token limit or verify that "
                "thinking mode is disabled by the selected provider."
            )

        raise RuntimeError(
            "The Hugging Face model returned an empty response."
        )

    return content.strip()


# Local generation

def _generate_local(
    messages: list[dict],
    temperature: float,
    max_new_tokens: int,
) -> str:
    """
    generate an answer using the locally loaded model
    """

    model, tokenizer = (
        _get_local_model()
    )

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
    )

    device = next(
        model.parameters()
    ).device

    inputs = {
        key: value.to(device)
        for key, value
        in inputs.items()
    }

    temperature = float(
        temperature
    )

    do_sample = (
        temperature > 0
    )

    generation_kwargs = {
        "max_new_tokens": int(
            max_new_tokens
        ),
        "do_sample": do_sample,
        "pad_token_id": (
            tokenizer.eos_token_id
        ),
        "eos_token_id": (
            tokenizer.eos_token_id
        ),
    }

    if do_sample:

        generation_kwargs.update(
            temperature=max(
                temperature,
                1e-5,
            ),
            top_p=0.85, # param in ml model that controls the randomness and creativity of text generation, tells the model to add up the probabilities of the most likely next words until they reach 85%, the model only picks from this group of words, cutting out the 15% of weird or rare words
        )

    with torch.inference_mode():

        output = model.generate(
            **inputs,
            **generation_kwargs,
        )

    generated_tokens = output[0][
        inputs["input_ids"].shape[1]:
    ]

    answer = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True,
    )

    return answer.strip()


# backend selection

def choose_backend() -> str:
    """
    choosing the generation backend
    """

    configured = (
        LLM_BACKEND
        .strip()
        .lower()
    )

    if IS_HF_SPACE:
        if not HF_TOKEN:
            raise RuntimeError(
                "HF_TOKEN is required on the Hugging Face Space. "
                "Configure it as a Space Secret and use "
                "LLM_BACKEND=hf_api."
            )

        return "hf_api"

    if configured in {
        "hf_api",
        "local",
    }:
        return configured

    if HF_TOKEN:
        return "hf_api"

    return "local"


# citation validation

def _extract_citation_ids(
    answer: str,
) -> list[str]:
    """
    Extract source IDs such as:
        [S1]
        [S2]
    """

    return re.findall(
        r"\[S(\d+)\]",
        answer,
    )


# main RAG answer function

def answer_question(
    question: str,
    hits: list[SearchHit],
    temperature: float = 0.15,
    max_new_tokens: int = 700,
) -> tuple[str, dict]:
    """
    Produce a grounded answer from retrieved source passages
    """

    question = question.strip()

    if not question:

        return (
            "Please enter a policy question.",
            {
                "backend": None,
                "retrieved_chunks": 0,
            },
        )

    if not hits:

        return (
            "No indexed source passages were "
            "available for this question",
            {
                "backend": None,
                "retrieved_chunks": 0,
            },
        )

    context = build_context(
        hits
    )

    user_prompt = f"""
Question:
{question}

Retrieved sources:
{context}

Instructions:
Answer the question using only the retrieved sources

Cite factual claims inline using the source identifiers
provided above, for example [S1] or [S2]

Do not invent source identifiers

If the retrieved material does not contain enough information
to answer the question, state that explicitly

Write the answer now.
""".strip()

    messages = [
        {
            "role": "system",
            "content": SYSTEM_PROMPT,
        },
        {
            "role": "user",
            "content": user_prompt,
        },
    ]

    backend = choose_backend()

    
    # Hosted Hugging Face generation

    if backend == "hf_api":

        try:

            answer = _generate_hf_api(
                messages,
                temperature,
                max_new_tokens,
            )

            model_used = (
                HF_INFERENCE_MODEL
            )

        except Exception as exc:

            # if user explicitly selected hf_api, surface the API failure
            if (
                LLM_BACKEND
                .strip()
                .lower()
                == "hf_api"
            ):
                raise

            print(
                "[LLM] Hugging Face API failed. "
                "Falling back to local model.",
                flush=True,
            )

            print(
                f"[LLM] API error: {exc}",
                flush=True,
            )

            answer = _generate_local(
                messages,
                temperature,
                max_new_tokens,
            )

            model_used = (
                LOCAL_MODEL_ID
            )

            backend = (
                "local fallback after HF API "
                f"error: {type(exc).__name__}"
            )

    
    # Local generation

    else:

        answer = _generate_local(
            messages,
            temperature,
            max_new_tokens,
        )

        model_used = (
            LOCAL_MODEL_ID
        )

    # Citation validation

    valid_ids = {
        str(i)
        for i in range(
            1,
            len(hits) + 1,
        )
    }

    cited_ids = (
        _extract_citation_ids(
            answer
        )
    )

    invalid_ids = sorted(
        {
            citation_id
            for citation_id
            in cited_ids
            if citation_id
            not in valid_ids
        }
    )

    if invalid_ids:

        answer += (
            "\n\n**Citation check:** "
            "The model emitted an invalid source "
            "marker. Verify the retrieved source "
            "panel before relying on that statement."
        )

    diagnostics = {
        "backend": backend,
        "model": model_used,
        "embedding_model": (
            EMBEDDING_MODEL_ID
        ),
        "embedding_device": (
            get_embedding_device()
        ),
        "retrieved_chunks": len(
            hits
        ),
        "citations_emitted": sorted(
            set(cited_ids),
            key=lambda x: int(x),
        ),
        "invalid_citations": (
            invalid_ids
        ),
        "context_characters": len(
            context
        ),
    }

    return (
        answer,
        diagnostics,
    )


# render retrieved source cards

def render_sources(
    hits: list[SearchHit],
) -> str:
    """
    Render retrieved passages as HTML cards for Gradio.
    """

    if not hits:

        return (
            "<div class='status-card'>"
            "No sources retrieved."
            "</div>"
        )

    cards: list[str] = []

    for source_number, hit in enumerate(
        hits,
        start=1,
    ):

        chunk = hit.chunk

        excerpt = (
            chunk.text
            .strip()
            .replace("\n", " ")
        )

        if len(excerpt) > 520:

            excerpt = (
                excerpt[:517]
                .rstrip()
                + "..."
            )

        link = ""

        if chunk.url:

            safe_url = html.escape(
                chunk.url,
                quote=True,
            )

            link = (
                f"<a href='{safe_url}' "
                "target='_blank' "
                "rel='noopener noreferrer'>"
                "Open source"
                "</a>"
            )

        safe_citation = html.escape(
            chunk.citation
        )

        safe_excerpt = html.escape(
            excerpt
        )

        cards.append(
            f"""
            <div class="source-card">
                <div class="source-top">
                    <span class="source-id">
                        S{source_number}
                    </span>

                    <span class="score">
                        score {hit.score:.3f}
                    </span>
                </div>

                <div class="source-title">
                    {safe_citation}
                </div>

                <div class="source-excerpt">
                    {safe_excerpt}
                </div>

                <div class="source-link">
                    {link}
                </div>
            </div>
            """
        )

    return (
        "<div class='sources-grid'>"
        + "".join(cards)
        + "</div>"
    )