streamlit_app.py

"""
RAG PDF Q&A — upload a PDF, ask questions, get answers grounded in the document.

pypdf text extraction -> paragraph chunking -> MiniLM embeddings ->
FAISS cosine retrieval (top-4) -> answer from Groq, sources shown.
Runs free on Streamlit Community Cloud (CPU).
"""

import os
import re
import numpy as np
import streamlit as st
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
import faiss
from groq import Groq

EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
GROQ_MODEL = os.environ.get("GROQ_MODEL", st.secrets.get("GROQ_MODEL", "openai/gpt-oss-120b"))
CHUNK_CHARS, CHUNK_OVERLAP, TOP_K = 900, 150, 4

SYSTEM = ("Answer strictly from the provided context. "
          "If the answer is not in the context, say you can't find it. Be concise.")


def _groq_keys() -> list[str]:
    keys = []
    for name in ("GROQ_API_KEY", "GROQ_API_KEY2", "GROQ_API_KEY3", "GROQ_API_KEY4", "GROQ_API_KEY5"):
        v = os.environ.get(name) or st.secrets.get(name, "")
        if v:
            keys.append(v)
    return keys


GROQ_KEYS = _groq_keys()


@st.cache_resource
def get_embedder():
    return SentenceTransformer(EMBED_MODEL)


def extract_text(file) -> str:
    reader = PdfReader(file)
    return "\n".join((p.extract_text() or "") for p in reader.pages)


def chunk_text(text: str) -> list[str]:
    text = re.sub(r"[ \t]+", " ", text)
    paras = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
    chunks, buf = [], ""
    for p in paras:
        if len(buf) + len(p) + 1 <= CHUNK_CHARS:
            buf = f"{buf}\n{p}".strip()
        else:
            if buf:
                chunks.append(buf)
            if len(p) <= CHUNK_CHARS:
                buf = p
            else:
                for i in range(0, len(p), CHUNK_CHARS - CHUNK_OVERLAP):
                    chunks.append(p[i:i + CHUNK_CHARS])
                buf = ""
    if buf:
        chunks.append(buf)
    return chunks


def build_index(chunks):
    vecs = get_embedder().encode(chunks, convert_to_numpy=True, normalize_embeddings=True)
    index = faiss.IndexFlatIP(vecs.shape[1])
    index.add(vecs.astype(np.float32))
    return index


def retrieve(query, chunks, index):
    q = get_embedder().encode([query], convert_to_numpy=True, normalize_embeddings=True).astype(np.float32)
    scores, idx = index.search(q, min(TOP_K, len(chunks)))
    return [(chunks[i], float(scores[0][r])) for r, i in enumerate(idx[0])]


def _chat(key, messages):
    client = Groq(api_key=key)
    kwargs = dict(model=GROQ_MODEL, messages=messages, max_tokens=1024, temperature=0.2)
    if GROQ_MODEL.startswith("openai/gpt-oss"):
        kwargs["reasoning_effort"] = "low"
    return client.chat.completions.create(**kwargs).choices[0].message.content.strip()


def answer(query, chunks, index):
    hits = retrieve(query, chunks, index)
    context = "\n\n---\n\n".join(c for c, _ in hits)
    messages = [
        {"role": "system", "content": SYSTEM},
        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"},
    ]
    errors = []
    for i, key in enumerate(GROQ_KEYS, 1):
        try:
            return _chat(key, messages), hits
        except Exception as e:
            errors.append(f"key {i}: {e}")
    raise RuntimeError("all Groq keys failed:\n" + "\n".join(errors))


st.set_page_config(page_title="RAG PDF Q&A", page_icon="📄")
st.title("📄 RAG PDF Q&A")
st.caption(f"MiniLM embeddings → FAISS cosine retrieval → Groq `{GROQ_MODEL}`. Answers cite their sources.")

if not GROQ_KEYS:
    st.warning(
        "Set `GROQ_API_KEY` (and optionally `GROQ_API_KEY2`..`5`) in the app's Secrets "
        "(Manage app → Settings → Secrets). Free keys: https://console.groq.com/keys"
    )
    st.stop()

up = st.file_uploader("PDF (text-based, no OCR)", type="pdf")

if up:
    if st.session_state.get("name") != up.name:
        with st.spinner("Extracting and indexing…"):
            text = extract_text(up)
            if len(text.strip()) < 200:
                st.error("Could not extract text — this PDF looks scanned. No OCR in this demo.")
                st.stop()
            chunks = chunk_text(text)
            st.session_state.update(name=up.name, chunks=chunks, index=build_index(chunks))
        st.success(f"Indexed {len(st.session_state['chunks'])} chunks from {len(text):,} characters.")

    q = st.text_input("Question", placeholder="What does chapter 2 say about …?")
    if q:
        with st.spinner("Retrieving + answering…"):
            try:
                out, hits = answer(q, st.session_state["chunks"], st.session_state["index"])
            except Exception as e:
                st.error(f"LLM call failed: {e}")
                st.stop()
        st.markdown("### Answer")
        st.write(out)
        with st.expander("Retrieved sources"):
            for i, (c, s) in enumerate(hits, 1):
                st.markdown(f"**[{i}] score {s:.2f}**")
                st.write(c[:800] + ("…" if len(c) > 800 else ""))