src/generation/rag_pipeline.py

"""
src/generation/rag_pipeline.py
--------------------------------
Scientific RAG pipeline with:
- grounded answer generation
- source citation
- hallucination control (out-of-corpus refusal)
- optional reranking stage
Compatible with lab_s6 rag_pipeline.py API.
"""

import time
from typing import Optional

# ---------------------------------------------------------------------------
# Prompt templates
# ---------------------------------------------------------------------------

SCIENTIFIC_RAG_PROMPT = """\
You are an expert scientific assistant. Answer ONLY from the provided context.
If the context does not contain the answer, reply EXACTLY:
"I cannot answer this question from the available scientific corpus."

Do NOT hallucinate or add information beyond the context.

### Context
{context}

### Question
{question}

### Answer (cite papers in [Title, Year] format when possible)
"""

NO_RAG_PROMPT = """\
You are a knowledgeable AI assistant. Answer the following question to the best of your ability.

### Question
{question}

### Answer
"""


# ---------------------------------------------------------------------------
# Citation helper
# ---------------------------------------------------------------------------

def format_sources(metadatas: list[dict], distances: list[float], rerank_scores: list[float] = None) -> list[dict]:
    """Format retrieved chunk metadata into source cards."""
    sources = []
    for i, (meta, dist) in enumerate(zip(metadatas, distances)):
        score_label = f"{rerank_scores[i]:.3f}" if rerank_scores else f"{1 - dist:.3f}"
        sources.append({
            "rank": i + 1,
            "paper_id": meta.get("paper_id", ""),
            "title": meta.get("title", "Unknown"),
            "authors": meta.get("authors_str", ""),
            "year": meta.get("year", ""),
            "section": meta.get("section", ""),
            "score": score_label,
            "chunk_preview": "",   # filled below
        })
    return sources


# ---------------------------------------------------------------------------
# Core pipeline functions
# ---------------------------------------------------------------------------

def answer_without_rag(question: str, llm, max_new_tokens: int = 400) -> str:
    """Baseline: direct LLM generation with no retrieval context."""
    prompt = NO_RAG_PROMPT.format(question=question)
    return llm.generate(prompt, max_new_tokens=max_new_tokens)


def answer_with_rag(
    question: str,
    embedder,
    store,
    llm,
    k: int = 5,
    max_new_tokens: int = 400,
    reranker=None,
    reranker_k: int = 20,
    metadata_filter: Optional[dict] = None,
) -> dict:
    """
    Full scientific RAG pipeline.

    Parameters
    ----------
    question      : user question
    embedder      : Embedder instance
    store         : ScientificChromaStore instance
    llm           : LLM backend (generate method)
    k             : final number of chunks in the context
    max_new_tokens: generation budget
    reranker      : CrossEncoderReranker instance (optional)
    reranker_k    : first-stage retrieval size when reranking
    metadata_filter: Chroma 'where' dict for filtering

    Returns
    -------
    dict with keys:
        question, retrieved, metadatas, distances, rerank_scores,
        sources, prompt, answer, latency
    """
    timings = {}

    # Stage 1: Embed question
    t0 = time.time()
    q_vec = embedder.encode_one(question)
    timings["embed_s"] = round(time.time() - t0, 3)

    # Stage 2: Dense retrieval
    t0 = time.time()
    first_k = reranker_k if reranker else k
    res = store.query(q_vec, k=first_k, where=metadata_filter)
    timings["retrieve_s"] = round(time.time() - t0, 3)

    rerank_scores = None

    # Stage 3 (optional): Cross-encoder reranking
    if reranker and res["documents"]:
        t0 = time.time()
        ranked = reranker.rerank(question, res, top_k=k)
        timings["rerank_s"] = round(time.time() - t0, 3)
        from src.retrieval.reranker import ranked_to_result
        res = ranked_to_result(ranked)
        rerank_scores = res.get("rerank_scores")
    else:
        # Keep only top-k from dense retrieval
        for key in ["documents", "metadatas", "distances"]:
            res[key] = res[key][:k]

    # Stage 4: Build grounded prompt with context
    context_parts = []
    for i, (doc, meta) in enumerate(zip(res["documents"], res["metadatas"])):
        citation = f"[{meta.get('title', 'Unknown')} ({meta.get('year', '?')})]"
        context_parts.append(f"Source {i+1} {citation}:\n{doc}")
    context = "\n\n---\n\n".join(context_parts)
    prompt = SCIENTIFIC_RAG_PROMPT.format(context=context, question=question)

    # Stage 5: Generate
    t0 = time.time()
    answer = llm.generate(prompt, max_new_tokens=max_new_tokens)
    timings["generate_s"] = round(time.time() - t0, 3)
    timings["total_s"] = round(sum(timings.values()), 3)

    # Format source cards
    sources = format_sources(res["metadatas"], res["distances"], rerank_scores)
    for i, src in enumerate(sources):
        src["chunk_preview"] = res["documents"][i][:200] + "..."

    return {
        "question": question,
        "retrieved": res["documents"],
        "metadatas": res["metadatas"],
        "distances": res["distances"],
        "rerank_scores": rerank_scores,
        "sources": sources,
        "prompt": prompt,
        "answer": answer,
        "latency": timings,
        "used_reranker": reranker is not None,
    }


def run_comparison(
    question: str,
    embedder,
    store,
    llm,
    reranker=None,
    k: int = 5,
) -> dict:
    """
    Run all three experiment conditions for one question:
    Exp1: No RAG, Exp2: RAG, Exp3: RAG + Reranker

    Returns dict with keys: no_rag, rag, rag_reranker
    """
    result = {
        "question": question,
        "no_rag": answer_without_rag(question, llm),
        "rag": answer_with_rag(question, embedder, store, llm, k=k),
    }
    if reranker:
        result["rag_reranker"] = answer_with_rag(
            question, embedder, store, llm, k=k, reranker=reranker
        )
    return result