app_gradio.py

"""
Gradio UI for the RAG Document Q&A pipeline.

Run:
    source venv/bin/activate
    python app_gradio.py

Opens at http://127.0.0.1:7860
"""

import os
import time
from pathlib import Path

import gradio as gr
from dotenv import load_dotenv

from ingestion.embedder import Embedder
from ingestion.pipeline import IngestionPipeline
from retrieval.index import VectorIndex
from retrieval.searcher import search
from generation.generator import Generator

load_dotenv()


# ---------------------------------------------------------------------------
# Startup — same order as main.py lifespan
# ---------------------------------------------------------------------------

print("Loading embedder model...")
embedder = Embedder()
print(f"Initialising FAISS index (dim={embedder.dimension})...")
index = VectorIndex(dimension=embedder.dimension)
print("Building ingestion pipeline...")
pipeline = IngestionPipeline(
    embedder=embedder,
    index=index,
    strategy="recursive_character",
    chunk_size=500,
    overlap=100,
)
print("Initialising Groq generator...")
generator = Generator()
print("All components ready.\n")


# ---------------------------------------------------------------------------
# Callbacks
# ---------------------------------------------------------------------------

def ingest_files(files) -> tuple[str, str]:
    """Ingest uploaded PDFs and return a results markdown string + updated stats."""
    if not files:
        return "No files selected.", _stats_line()

    lines: list[str] = []
    for file in files:
        path = Path(file if isinstance(file, str) else file.name)
        t0 = time.perf_counter()
        result = pipeline.ingest_pdf(path)
        elapsed_ms = round((time.perf_counter() - t0) * 1000)

        if result.error:
            lines.append(f"**{path.name}** — Error: {result.error}")
        else:
            lines.append(
                f"**{result.file}** — "
                f"{result.pages} pages · **{result.chunks} chunks** · {elapsed_ms} ms"
            )

    return "\n\n".join(lines), _stats_line()


def answer_question(question: str, history: list[dict], expand_query: bool):
    """Retrieve chunks then stream the answer token-by-token into the chat history."""
    if not question.strip():
        yield history, ""
        return

    if index.size == 0:
        history = history + [
            {"role": "user", "content": question},
            {
                "role": "assistant",
                "content": (
                    "No documents have been indexed yet. "
                    "Upload one or more PDFs in the **Upload PDFs** tab first."
                ),
            },
        ]
        yield history, ""
        return

    search_resp = search(
        query=question,
        embedder=embedder,
        index=index,
        k=8,
        expand_query=expand_query,
    )
    results = search_resp.chunks

    source_lines = [
        f"- **[Source {i}]** `{r.metadata.get('source', '?')}` "
        f"p. {r.metadata.get('page_num', '?')} "
        f"— relevance {r.score:.3f}"
        for i, r in enumerate(results, start=1)
    ]
    sources_block = "---\n**Retrieved context:**\n" + "\n".join(source_lines)

    history = history + [
        {"role": "user", "content": question},
        {"role": "assistant", "content": ""},
    ]
    yield history, ""

    accumulated = ""
    try:
        for chunk in generator.generate_answer_stream(question, results, max_score=search_resp.max_score):
            accumulated += chunk
            history[-1]["content"] = accumulated
            yield history, ""
    except Exception as exc:
        history[-1]["content"] = accumulated + f"\n\n*Error generating answer: {exc}*"
        yield history, ""
        return

    history[-1]["content"] = accumulated + "\n\n" + sources_block
    yield history, ""


def clear_chat() -> tuple[list, str]:
    return [], ""


def _stats_line() -> str:
    return f"Index: **{index.size} vectors** across **{pipeline.chunk_count} chunks ingested**"


# ---------------------------------------------------------------------------
# UI layout
# ---------------------------------------------------------------------------

with gr.Blocks(title="RAG Document Q&A", theme=gr.themes.Soft()) as demo:

    gr.Markdown(
        "# RAG Document Q&A\n"
        "Upload PDF documents, then ask questions. "
        "Answers are grounded in the documents and cite sources inline."
    )

    with gr.Tab("Upload PDFs"):
        gr.Markdown(
            "Select one or more PDF files. Each is extracted, chunked, embedded, "
            "and added to the shared FAISS index."
        )
        with gr.Row():
            with gr.Column(scale=2):
                file_input = gr.File(
                    label="PDF files",
                    file_types=[".pdf"],
                    file_count="multiple",
                    height=160,
                )
                ingest_btn = gr.Button("Ingest", variant="primary", size="lg")
            with gr.Column(scale=3):
                ingest_output = gr.Markdown(label="Ingestion results", value="Results will appear here.")
                index_stats = gr.Markdown(value=_stats_line())

        ingest_btn.click(
            fn=ingest_files,
            inputs=file_input,
            outputs=[ingest_output, index_stats],
        )
        gr.Markdown("*Note: If you change retrieval settings, re-upload your PDFs to re-index with new chunking parameters.*")

    with gr.Tab("Ask Questions"):
        gr.Markdown(
            "Type a question about your uploaded documents. "
            "Press **Enter** or click **Ask**. "
            "Answers include `[Source N]` citations linked to specific pages."
        )

        chatbot = gr.Chatbot(
            label="Conversation",
            height=500,
            type="messages",
            render_markdown=True,
            placeholder="Ingest some PDFs, then ask a question.",
        )

        with gr.Row():
            question_box = gr.Textbox(
                placeholder="e.g. What retrieval method does RAG use?",
                label="",
                lines=1,
                scale=5,
                submit_btn=False,
            )
            ask_btn = gr.Button("Ask", variant="primary", scale=1, min_width=80)
            
        with gr.Accordion("Advanced Retrieval Options", open=False):
            expand_cb = gr.Checkbox(
                label="Enable Adaptive Query Expansion (LLM rephrases your question only when initial search confidence is low)",
                value=True
            )

        clear_btn = gr.Button("Clear conversation", size="sm", variant="secondary")

        # Submit on button click or Enter key
        ask_btn.click(
            fn=answer_question,
            inputs=[question_box, chatbot, expand_cb],
            outputs=[chatbot, question_box],
        )
        question_box.submit(
            fn=answer_question,
            inputs=[question_box, chatbot, expand_cb],
            outputs=[chatbot, question_box],
        )
        clear_btn.click(fn=clear_chat, outputs=[chatbot, question_box])


if __name__ == "__main__":
    demo.launch(
        server_name=os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1"),
        server_port=int(os.environ.get("GRADIO_SERVER_PORT", "7860")),
        show_error=True,
    )