app.py

"""
app.py
Gradio UI for the RAG Document Q&A system.
"""

import os
import shutil
import tempfile
import logging

import gradio as gr

from rag_pipeline import RAGPipeline

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

pipeline = RAGPipeline()


# ──────────────────────────────────────────────
# Handlers
# ──────────────────────────────────────────────

def handle_upload(file):
    if file is None:
        return "⚠️ No file uploaded.", gr.update(interactive=False)

    try:
        ext = os.path.splitext(file.name)[1].lower()
        tmp_path = tempfile.mktemp(suffix=ext)
        shutil.copy(file.name, tmp_path)

        stats = pipeline.ingest_document(tmp_path)

        status = (
            f"✅ **Document ready!**\n\n"
            f"- 📄 File: `{stats['document']}`\n"
            f"- 🧩 Chunks: `{stats['chunks']}`\n"
            f"- 📐 Embedding dim: `{stats['embedding_dim']}`\n"
            f"- ⚡ Ingestion time: `{stats['ingestion_time_s']}s`\n\n"
            f"You can now ask questions below."
        )
        return status, gr.update(interactive=True)

    except Exception as e:
        logger.error(f"Upload error: {e}")
        return f"❌ Error processing document: {str(e)}", gr.update(interactive=False)


def handle_query(question, history):
    if not pipeline.is_ready:
        history.append((question, "⚠️ Please upload a document first."))
        return history, "", ""

    if not question.strip():
        return history, "", ""

    try:
        result = pipeline.query(question)

        answer = result["answer"]
        latency = result["latency_s"]

        # ✅ FIXED: use rank instead of score
        sources_md = "### 📚 Retrieved Sources\n\n"

        for src in result["sources"]:
            rank = src["rank"]
            excerpt = src["chunk"][:300] + ("..." if len(src["chunk"]) > 300 else "")

            sources_md += f"**Source {rank}**\n\n> {excerpt}\n\n---\n\n"

        history.append((question, answer))

        return history, sources_md, f"⚡ Response time: {latency}s"

    except Exception as e:
        logger.error(f"Query error: {e}")
        history.append((question, f"❌ Error: {str(e)}"))
        return history, "", ""


def clear_all():
    return [], "", ""


# ──────────────────────────────────────────────
# UI
# ──────────────────────────────────────────────

with gr.Blocks(
    title="RAG Document Q&A",
    theme=gr.themes.Soft(primary_hue="blue"),
    css="""
        .title { text-align: center; margin-bottom: 8px; }
        .subtitle { text-align: center; color: #64748b; margin-bottom: 24px; }
        footer { display: none !important; }
    """
) as demo:

    gr.Markdown("# 📄 RAG Document Q&A", elem_classes="title")
    gr.Markdown(
        "Upload a PDF or TXT document, then ask questions. "
        "Answers are grounded in your document — not hallucinated.",
        elem_classes="subtitle"
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📁 Upload Document")
            file_input = gr.File(
                label="Upload PDF or TXT (max 10 MB)",
                file_types=[".pdf", ".txt"],
            )
            upload_status = gr.Markdown("_No document loaded yet._")

        with gr.Column(scale=2):
            gr.Markdown("### 💬 Ask a Question")
            chatbot = gr.Chatbot(height=380, label="Conversation")

            with gr.Row():
                question_box = gr.Textbox(
                    placeholder="e.g. What is the main topic of this document?",
                    label="Your question",
                    scale=4,
                    interactive=False,
                )
                submit_btn = gr.Button("Ask 🚀", variant="primary", scale=1)

            latency_display = gr.Markdown("")
            clear_btn = gr.Button("🗑️ Clear conversation", variant="secondary")

    with gr.Accordion("📚 View Retrieved Sources", open=False):
        sources_display = gr.Markdown("_Ask a question to see retrieved sources._")

    gr.Markdown("### 💡 Example Questions")
    gr.Examples(
        examples=[
            ["What is the main topic of this document?"],
            ["Summarize the key points."],
            ["What conclusions are drawn?"],
            ["What data or evidence is mentioned?"],
        ],
        inputs=question_box,
    )

    # Events
    file_input.change(
        fn=handle_upload,
        inputs=file_input,
        outputs=[upload_status, question_box],
    )

    submit_btn.click(
        fn=handle_query,
        inputs=[question_box, chatbot],
        outputs=[chatbot, sources_display, latency_display],
    )

    question_box.submit(
        fn=handle_query,
        inputs=[question_box, chatbot],
        outputs=[chatbot, sources_display, latency_display],
    )

    clear_btn.click(
        fn=clear_all,
        outputs=[chatbot, sources_display, latency_display],
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
    )