app.py

"""
RAG-Based AI Assistant — Multi-Query + RRF + Reranking + Citation Agent
=======================================================================
Upload one or more PDFs in the sidebar and ask questions about their contents.

Query-time pipeline:
    question
      └─► LLM writes N paraphrases              (multi-query rewriting)
            └─► N parallel vector searches
                  └─► RRF merges ranked lists    (fusion)
                        └─► cross-encoder        (reranking)
                              └─► Citation Agent (answer + inline [N] citations)
                                    └─► 📚 Sources panel (chunk + filename + page)

Run with:
    streamlit run app.py
"""

import os
import re
import hashlib
import tempfile
from collections import defaultdict

import streamlit as st

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq
from langchain_core.output_parsers import StrOutputParser
from sentence_transformers import CrossEncoder


# ---------------------------------------------------------------------------
# Groq API key — never hard-coded here.
#   Local:     .streamlit/secrets.toml  →  GROQ_API_KEY = "gsk_..."
#   HF Space:  add GROQ_API_KEY as a repository secret
# ---------------------------------------------------------------------------
def load_api_key() -> str:
    try:
        if "GROQ_API_KEY" in st.secrets:
            return st.secrets["GROQ_API_KEY"]
    except Exception:
        pass
    return os.environ.get("GROQ_API_KEY", "")


GROQ_API_KEY = load_api_key()

N_PARAPHRASES = 5    # query paraphrases generated per question
K_PER_QUERY   = 3    # chunks retrieved per query before RRF
TOP_N_RERANK  = 15   # chunks passed to the LLM after reranking


# ---------------------------------------------------------------------------
# Page setup
# ---------------------------------------------------------------------------
st.set_page_config(page_title="RAG-Based AI Assistant", page_icon="🤖", layout="centered")
st.title("🤖 RAG-Based AI Assistant")
st.caption(
    "Multi-Query · RRF · Cross-Encoder Reranking · Citation Agent · Powered by Groq"
)


# ---------------------------------------------------------------------------
# Sidebar
# ---------------------------------------------------------------------------
with st.sidebar:
    st.header("Documents")

    uploaded_files = st.file_uploader(
        "Upload PDF(s)",
        type=["pdf"],
        accept_multiple_files=True,
        help="The assistant answers questions based on these files.",
    )

    model_name = st.selectbox(
        "Groq model",
        ["llama-3.3-70b-versatile", "llama-3.1-8b-instant", "openai/gpt-oss-20b"],
        index=0,
    )

    if st.button("Clear chat history"):
        st.session_state.messages = []
        st.rerun()


# ---------------------------------------------------------------------------
# Cached models (loaded once per session)
# ---------------------------------------------------------------------------
@st.cache_resource(show_spinner=False)
def load_embeddings() -> HuggingFaceEmbeddings:
    return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")


@st.cache_resource(show_spinner="Loading reranker model (first run only)…")
def load_reranker() -> CrossEncoder:
    return CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")


# ---------------------------------------------------------------------------
# PDF ingestion + vector store
# ---------------------------------------------------------------------------
def _files_signature(files) -> str:
    h = hashlib.sha256()
    for f in files:
        h.update(f.name.encode())
        h.update(f.getvalue())
    return h.hexdigest()


def _load_pdfs(files) -> list:
    """Load PDFs and stamp every chunk with the original filename."""
    docs = []
    for f in files:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            tmp.write(f.getvalue())
            tmp_path = tmp.name
        try:
            loaded = PyPDFLoader(tmp_path).load()
            for doc in loaded:
                # PyPDFLoader sets `source` to the temp path — override with the
                # real filename so citations are human-readable.
                doc.metadata["filename"] = f.name
            docs.extend(loaded)
        finally:
            os.remove(tmp_path)
    return docs


@st.cache_resource(show_spinner="Indexing your PDFs…")
def build_vectorstore(file_hash: str, _files):
    """Index PDFs into Chroma. Re-runs only when the uploaded files change."""
    docs = _load_pdfs(_files)
    splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=200)
    chunks = splitter.split_documents(docs)
    return Chroma.from_documents(documents=chunks, embedding=load_embeddings())


# ---------------------------------------------------------------------------
# Multi-query retrieval
# ---------------------------------------------------------------------------
def multi_query_retrieve(
    question: str,
    vectorstore,
    n: int,
    k: int,
    api_key: str,
    model: str,
) -> tuple[list[str], list[list]]:
    os.environ["GROQ_API_KEY"] = api_key
    prompt = ChatPromptTemplate.from_template(
        """You are an AI assistant helping improve document retrieval.
Generate {n} different versions of the question below.
Each version should use different vocabulary or frame the intent from a different angle.
Return ONLY the questions — one per line, no numbering, no extra text.

Original question: {question}"""
    )
    llm = ChatGroq(model=model, temperature=0.4)
    raw = (prompt | llm | StrOutputParser()).invoke({"question": question, "n": n})
    paraphrases = [q.strip() for q in raw.strip().split("\n") if q.strip()]

    retriever = vectorstore.as_retriever(search_kwargs={"k": k})
    all_queries = [question] + paraphrases   # always include the original
    results_per_query = [retriever.invoke(q) for q in all_queries]
    return paraphrases, results_per_query


# ---------------------------------------------------------------------------
# RRF fusion
# ---------------------------------------------------------------------------
def reciprocal_rank_fusion(results_lists: list, k: int = 60) -> list:
    """
    score(doc) = Σ  1 / (k + rank_i(doc))   across all lists i
    Higher score → appeared near the top in more lists.
    """
    scores: dict[str, float] = defaultdict(float)
    doc_map: dict[str, object] = {}
    for ranked_list in results_lists:
        for rank, doc in enumerate(ranked_list, start=1):
            key = doc.page_content
            scores[key] += 1.0 / (k + rank)
            doc_map[key] = doc
    return [doc_map[key] for key in sorted(scores, key=scores.__getitem__, reverse=True)]


# ---------------------------------------------------------------------------
# Cross-encoder reranking
# ---------------------------------------------------------------------------
def rerank(question: str, docs: list, reranker: CrossEncoder, top_n: int) -> list:
    if not docs:
        return docs
    pairs = [[question, doc.page_content] for doc in docs]
    scores = reranker.predict(pairs)
    ranked = sorted(zip(scores, docs), key=lambda x: x[0], reverse=True)
    return [doc for _, doc in ranked[:top_n]]


# ---------------------------------------------------------------------------
# Citation Agent
# ---------------------------------------------------------------------------
def format_context_with_citations(docs: list) -> str:
    """
    Number each chunk and embed its source metadata so the LLM can cite by [N].

    Produces:
        [1] report.pdf · page 3
        <chunk text>

        [2] report.pdf · page 7
        <chunk text>
        ...
    """
    parts = []
    for i, doc in enumerate(docs, start=1):
        filename = doc.metadata.get("filename", "unknown")
        page = doc.metadata.get("page", 0) + 1   # PyPDFLoader is 0-indexed
        header = f"[{i}] {filename} · page {page}"
        parts.append(f"{header}\n{doc.page_content}")
    return "\n\n".join(parts)


def scope_check(question: str, context: str, api_key: str, model: str) -> bool:
    """
    Gate that decides whether the citation agent should run.
    Returns True if the retrieved context can answer the question, False otherwise.
    A dedicated YES/NO call so the answer chain prompt is not affected.
    """
    os.environ["GROQ_API_KEY"] = api_key
    prompt = ChatPromptTemplate.from_template(
        "You are a relevance checker for a document QA system.\n"
        "Given the context below and a question, decide whether the context contains\n"
        "sufficient information to answer the question.\n"
        "Reply with ONLY the word YES or NO — no punctuation, no explanation.\n\n"
        "Context:\n{context}\n\nQuestion: {question}"
    )
    llm = ChatGroq(model=model, temperature=0)
    result = (prompt | llm | StrOutputParser()).invoke(
        {"context": context, "question": question}
    )
    return result.strip().upper().startswith("YES")



def scope_check(question: str, context: str, api_key: str, model: str) -> bool:
    """
    Gate that decides whether the Citation Agent should run.
    Asks the LLM whether the retrieved context is sufficient to answer
    the question. Returns True (in scope) or False (out of scope).
    Kept as a separate call so the answer prompt is not affected.
    """
    os.environ["GROQ_API_KEY"] = api_key
    prompt = ChatPromptTemplate.from_template(
        "You are a relevance checker for a document QA system.\n"
        "Given the context below and a question, decide whether the context\n"
        "contains sufficient information to answer the question.\n"
        "Reply with ONLY the word YES or NO — no punctuation, no explanation.\n\n"
        "Context:\n{context}\n\nQuestion: {question}"
    )
    llm = ChatGroq(model=model, temperature=0)
    result = (prompt | llm | StrOutputParser()).invoke(
        {"context": context, "question": question}
    )
    return result.strip().upper().startswith("YES")


def build_answer_chain(api_key: str, model: str):
    """
    Answer chain that is instructed to add inline [N] citations.
    The context is pre-formatted with chunk numbers so the LLM can reference them.
    """
    os.environ["GROQ_API_KEY"] = api_key
    prompt = ChatPromptTemplate.from_template(
        """You are an assistant for question-answering tasks.
Use ONLY the numbered context chunks below to answer the question.
After every sentence or claim you draw from the context, add an inline citation [N]
where N is the chunk number. Every factual statement must be cited.
If the answer cannot be found in the context, say you don't know.
Be concise.

{context}

Question: {question}"""
    )
    llm = ChatGroq(model=model, temperature=0)
    return prompt | llm | StrOutputParser()


def extract_cited_chunks(answer: str, docs: list) -> list[dict]:
    """
    Parse [N] markers out of the answer and return the matching chunks
    with human-readable source metadata.

    Falls back to all chunks if the LLM produced no citation markers,
    so the Sources panel is never empty.
    """
    cited_ids = sorted({int(m) for m in re.findall(r"\[(\d+)\]", answer)})
    # Keep only IDs that are actually in range.
    cited_ids = [i for i in cited_ids if 1 <= i <= len(docs)]
    if not cited_ids:
        cited_ids = list(range(1, len(docs) + 1))   # graceful fallback

    return [
        {
            "id": i,
            "filename": docs[i - 1].metadata.get("filename", "unknown"),
            "page": docs[i - 1].metadata.get("page", 0) + 1,
            "text": docs[i - 1].page_content,
        }
        for i in cited_ids
    ]


def render_citations(cited_chunks: list, expanded: bool = False) -> None:
    """Render the Sources panel below the answer."""
    if not cited_chunks:
        return
    with st.expander("📚 Sources", expanded=expanded):
        for chunk in cited_chunks:
            st.markdown(
                f"**[{chunk['id']}]** &nbsp; `{chunk['filename']}` &nbsp;·&nbsp; page {chunk['page']}"
            )
            # Show a readable snippet — truncate long chunks.
            snippet = chunk["text"][:450].strip()
            if len(chunk["text"]) > 450:
                snippet += "…"
            st.caption(snippet)
            if chunk is not cited_chunks[-1]:
                st.divider()


# ---------------------------------------------------------------------------
# Chat state + replay
# ---------------------------------------------------------------------------
if "messages" not in st.session_state:
    st.session_state.messages = []

for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        if msg.get("paraphrases"):
            with st.expander("🔀 Query paraphrases used for retrieval"):
                for i, q in enumerate(msg["paraphrases"], 1):
                    st.markdown(f"**{i}.** {q}")
        st.markdown(msg["content"])
        if msg.get("cited_chunks"):
            render_citations(msg["cited_chunks"], expanded=False)

ready = bool(uploaded_files)
if not ready:
    st.info("Upload one or more PDFs in the sidebar to start asking questions.")

user_input = st.chat_input("Ask a question about your PDFs…", disabled=not ready)


# ---------------------------------------------------------------------------
# Handle a new question
# ---------------------------------------------------------------------------
if user_input:
    if not GROQ_API_KEY:
        st.error(
            "No Groq API key found. "
            "Set GROQ_API_KEY in .streamlit/secrets.toml (local) "
            "or as a repository secret (HF Space)."
        )
        st.stop()

    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    try:
        vectorstore = build_vectorstore(_files_signature(uploaded_files), uploaded_files)
        answer_chain = build_answer_chain(GROQ_API_KEY, model_name)

        with st.chat_message("assistant"):

            # ── Step 1: multi-query retrieval ────────────────────────────
            with st.spinner("Generating paraphrases and retrieving chunks…"):
                paraphrases, results_per_query = multi_query_retrieve(
                    question=user_input,
                    vectorstore=vectorstore,
                    n=N_PARAPHRASES,
                    k=K_PER_QUERY,
                    api_key=GROQ_API_KEY,
                    model=model_name,
                )
            with st.expander("🔀 Query paraphrases used for retrieval", expanded=True):
                for i, q in enumerate(paraphrases, 1):
                    st.markdown(f"**{i}.** {q}")

            # ── Step 2: RRF fusion ───────────────────────────────────────
            fused_docs = reciprocal_rank_fusion(results_per_query)

            # ── Step 3: cross-encoder reranking ──────────────────────────
            with st.spinner("Reranking chunks…"):
                final_docs = rerank(
                    question=user_input,
                    docs=fused_docs,
                    reranker=load_reranker(),
                    top_n=TOP_N_RERANK,
                )

            # ── Step 4: Scope-check gate ──────────────────────────────────
            context = format_context_with_citations(final_docs)
            with st.spinner("Checking question scope…"):
                in_scope = scope_check(user_input, context, GROQ_API_KEY, model_name)

            # ── Step 5: Answer ──────────────────────────────────────
            answer = st.write_stream(
                answer_chain.stream({"context": context, "question": user_input})
            )

            # ── Step 6: Citation Agent — runs only when in scope ───────
            if in_scope:
                cited_chunks = extract_cited_chunks(answer, final_docs)
                render_citations(cited_chunks, expanded=True)
            else:
                cited_chunks = None

        st.session_state.messages.append({
            "role": "assistant",
            "content": answer,
            "paraphrases": paraphrases,
            "cited_chunks": cited_chunks,
        })

    except Exception as e:
        st.error(f"Something went wrong: {e}")