app.py

# gradio_app.py
import os
import gradio as gr

from backend import (
    create_vector_store,
    upload_single_pdf,
    answer_with_retrieval,
)

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise RuntimeError("OPENAI_API_KEY not set")


def upload_pdf(file_path, vector_store_id):
    if not file_path:
        yield vector_store_id, "⚠️ **No file uploaded.**"
        return

    if not file_path.lower().endswith(".pdf"):
        yield vector_store_id, "❌ **Only PDF files are supported.**"
        return

    # Create vector store if missing
    if vector_store_id is None:
        store = create_vector_store("gradio_session_store")
        vector_store_id = store["id"]
        yield vector_store_id, "📦 **Vector store created**"

    try:
        # 🔑 THIS IS THE IMPORTANT PART
        for message in upload_single_pdf(file_path, vector_store_id):
            yield vector_store_id, message

    except Exception as e:
        yield vector_store_id, f"❌ **Upload failed:** `{str(e)}`"


def ask_question(question, vector_store_id):
    # 1️⃣ Immediate feedback
    yield "🧠 **Thinking…**", ""

    if not vector_store_id:
        yield "❗ **Upload a PDF first.**", ""
        return

    if not question.strip():
        yield "❗ **Ask a real question.**", ""
        return

    try:
        result = answer_with_retrieval(
            question,
            [vector_store_id],
        )

        answer = result.get("answer", "")
        files = result.get("files", [])

        sources_md = (
            "\n".join([f"- `{f}`" for f in files])
            if files
            else "_No sources returned._"
        )

        yield answer, sources_md

    except Exception as e:
        yield f"❌ **Error:** `{str(e)}`", ""


# ========= UI =========

with gr.Blocks(title="RAG PDF Chatbot") as demo:

    gr.Markdown(
        """
        # 📄 RAG PDF Chatbot  
        ### Upload → Process → Ask → Answer

        You will see live feedback while documents are processed and answers are generated.
        """
    )

    vector_store_state = gr.State(None)

    # Upload
    gr.Markdown("## 📤 Upload Document")

    with gr.Row():
        pdf_input = gr.File(
            label="Upload PDF",
            file_types=[".pdf"],
            type="filepath",
        )
        upload_btn = gr.Button("🚀 Upload", variant="primary")

    upload_status = gr.Markdown()

    upload_btn.click(
        upload_pdf,
        inputs=[pdf_input, vector_store_state],
        outputs=[vector_store_state, upload_status],
    )

    gr.Markdown("---")

    # Question
    gr.Markdown("## 💬 Ask a Question (from first two pages)")

    question_input = gr.Textbox(
        placeholder="e.g. What is the main conclusion of this document?",
        lines=3,
        label="Your Question",
    )

    ask_btn = gr.Button("🧠 Ask", variant="secondary")

    # Answer
    gr.Markdown("## 📘 Answer")
    answer_output = gr.Markdown()
    gr.Markdown("### 📎 Sources")
    sources_output = gr.Markdown()

    ask_btn.click(
        ask_question,
        inputs=[question_input, vector_store_state],
        outputs=[answer_output, sources_output],
    )


if __name__ == "__main__":
    demo.launch()