app.py

"""
University Admissions RAG Chatbot — Main Gradio Application.

Run locally:
    python app.py

Hugging Face Spaces:
    Set HF_TOKEN in Spaces Secrets. The app auto-launches on port 7860.
"""

from __future__ import annotations

import os
import shutil
import tempfile
from pathlib import Path
from typing import Optional

import gradio as gr
from langchain_community.embeddings import HuggingFaceEmbeddings
from config import cfg
from logging_config import setup_logging, get_logger
from model_loader import GraniteModelLoader
from rag_engine import RAGEngine
from utils import (
    sanitize_input,
    format_sources,
    format_history_for_context,
    export_chat_history,
    ensure_dir,
)

# ── Bootstrap ─────────────────────────────────────────────────────────────────

setup_logging(log_dir=cfg.app.log_dir)
logger = get_logger(__name__)

ensure_dir(cfg.app.data_dir)
ensure_dir(cfg.app.log_dir)
ensure_dir(cfg.retrieval.index_path)

rag = RAGEngine()
granite = GraniteModelLoader()

# Try to load a persisted index on startup
_index_loaded_at_start = rag.load_index()


# ── Core pipeline ─────────────────────────────────────────────────────────────

def ensure_model_loaded() -> tuple[bool, str]:
    """Load the Granite model if not already loaded. Returns (success, message)."""
    if granite.is_loaded:
        return True, "Model already loaded."
    token = cfg.hf_token
    if not token:
        msg = (
            "⚠️ HF_TOKEN environment variable is not set. "
            "Please set it and restart the application."
        )
        logger.error(msg)
        return False, msg
    try:
        granite.load(token=token)
        return True, "Model loaded successfully."
    except Exception as exc:
        msg = f"❌ Model loading failed: {exc}"
        logger.exception("Model loading error.")
        return False, msg


def chat_fn(
    user_message: str,
    history: list[tuple[str, str]],
    temperature: float,
    max_tokens: int,
    top_p: float,
    retrieval_k: int,
) -> tuple[list[tuple[str, str]], str]:
    """
    Core chat function called by the Gradio interface.
    Returns updated (history, source_text).
    """
    user_message = sanitize_input(user_message)

    if not user_message:
        return history, "⚠️ Please enter a question."

    # Ensure model is loaded
    ok, err_msg = ensure_model_loaded()
    if not ok:
        history = history + [(user_message, err_msg)]
        return history, ""

    # Retrieval
    sources_text = ""
    retrieved_context = ""
    if rag.is_ready:
        try:
            docs = rag.retrieve(user_message, top_k=retrieval_k)
            if docs:
                retrieved_context = "\n\n---\n\n".join(
                    f"[Source: {Path(d.metadata.get('source', 'unknown')).name}]\n{d.page_content}"
                    for d in docs
                )
                sources_text = format_sources(docs)
            else:
                retrieved_context = "No relevant documents found in the knowledge base."
        except Exception as exc:
            logger.error("Retrieval error: %s", exc)
            retrieved_context = "Retrieval system is unavailable."
    else:
        retrieved_context = (
            "The knowledge base is not yet built. "
            "Please upload documents and click 'Build Knowledge Base'."
        )

    # Conversation history
    conv_history = format_history_for_context(history, max_turns=cfg.app.max_history_turns)

    # Build prompt and generate
    prompt = GraniteModelLoader.build_prompt(
        query=user_message,
        retrieved_context=retrieved_context,
        conversation_history=conv_history,
    )

    try:
        response = granite.generate(
            prompt=prompt,
            max_new_tokens=int(max_tokens),
            temperature=temperature,
            top_p=top_p,
        )
    except Exception as exc:
        logger.exception("Generation error.")
        response = f"❌ Generation error: {exc}"

    if sources_text:
        response = f"{response}\n\n**Sources:**\n{sources_text}"

    history = history + [(user_message, response)]
    return history, ""   # second value clears the input box


def upload_files_fn(files: list) -> str:
    """Copy uploaded temp files into the data directory."""
    if not files:
        return "⚠️ No files selected."

    ensure_dir(cfg.app.data_dir)
    saved: list[str] = []
    skipped: list[str] = []

    for file_obj in files:
        # Gradio provides a temp path as a string or NamedTempFile
        src = file_obj if isinstance(file_obj, str) else file_obj.name
        fname = Path(src).name
        ext = Path(fname).suffix.lower()

        if ext not in cfg.app.allowed_extensions:
            skipped.append(f"{fname} (unsupported type)")
            continue

        size_mb = os.path.getsize(src) / (1024 * 1024)
        if size_mb > cfg.app.max_file_size_mb:
            skipped.append(f"{fname} (exceeds {cfg.app.max_file_size_mb} MB)")
            continue

        dest = os.path.join(cfg.app.data_dir, fname)
        shutil.copy2(src, dest)
        saved.append(fname)
        logger.info("Uploaded: %s → %s", fname, dest)

    msg_parts: list[str] = []
    if saved:
        msg_parts.append(f"✅ Saved {len(saved)} file(s): {', '.join(saved)}")
    if skipped:
        msg_parts.append(f"⚠️ Skipped {len(skipped)}: {', '.join(skipped)}")

    return "\n".join(msg_parts) if msg_parts else "No files processed."


def build_kb_fn(chunk_size: int, chunk_overlap: int) -> str:
    """Load all documents from data/ and build the FAISS index."""
    data_dir = cfg.app.data_dir
    files = [
        os.path.join(data_dir, f)
        for f in os.listdir(data_dir)
        if Path(f).suffix.lower() in cfg.app.allowed_extensions
    ]

    if not files:
        return (
            f"⚠️ No supported documents found in '{data_dir}'. "
            "Please upload files first."
        )

    rag.update_splitter(int(chunk_size), int(chunk_overlap))

    try:
        docs = rag.load_documents(files)
        if not docs:
            return "⚠️ No content could be extracted from the uploaded documents."

        chunks = rag.chunk_documents(docs)
        rag.build_index(chunks)

        return (
            f"✅ Knowledge base built successfully!\n"
            f"   Documents: {rag.doc_count}\n"
            f"   Chunks:    {rag.chunk_count}"
        )
    except Exception as exc:
        logger.exception("Knowledge base build error.")
        return f"❌ Failed to build knowledge base: {exc}"


def reload_kb_fn() -> str:
    """Reload index from disk."""
    success = rag.load_index()
    if success:
        return (
            f"✅ Knowledge base reloaded.\n"
            f"   Documents: {rag.doc_count}\n"
            f"   Chunks:    {rag.chunk_count}"
        )
    return "⚠️ No saved index found. Please build the knowledge base first."


def get_kb_status() -> tuple[str, str, str]:
    """Return (status, doc_count, chunk_count) for the UI."""
    if rag.is_ready:
        status = "🟢 Ready"
        doc_c = str(rag.doc_count)
        chunk_c = str(rag.chunk_count)
    else:
        status = "🔴 Not initialised"
        doc_c = "0"
        chunk_c = "0"
    return status, doc_c, chunk_c


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


def export_fn(history: list[tuple[str, str]]) -> gr.File:
    transcript = export_chat_history(history)
    tmp = tempfile.NamedTemporaryFile(
        mode="w", suffix=".txt", delete=False, prefix="chat_export_"
    )
    tmp.write(transcript)
    tmp.flush()
    return tmp.name


def list_uploaded_files() -> str:
    data_dir = cfg.app.data_dir
    if not os.path.isdir(data_dir):
        return "No files uploaded yet."
    files = [
        f for f in os.listdir(data_dir)
        if Path(f).suffix.lower() in cfg.app.allowed_extensions
    ]
    if not files:
        return "No files uploaded yet."
    return "\n".join(f"📄 {f}" for f in sorted(files))


# ── Gradio UI ─────────────────────────────────────────────────────────────────

CSS = """
#header { text-align: center; padding: 1rem 0; }
#header h1 { font-size: 2rem; margin-bottom: 0.25rem; }
#chatbot { height: 480px; }
.panel-box { border: 1px solid #e0e0e0; border-radius: 8px; padding: 1rem; }
footer { display: none !important; }
"""

def build_ui() -> gr.Blocks:
    with gr.Blocks(
        title=cfg.app.title,
        css=CSS,
        theme=gr.themes.Soft(primary_hue="blue"),
    ) as demo:

        # ── Header ────────────────────────────────────────────────────────────
        with gr.Row(elem_id="header"):
            gr.Markdown(
                f"# {cfg.app.title}\n\n{cfg.app.description}"
            )

        # ── Main layout ───────────────────────────────────────────────────────
        with gr.Row():

            # Left column — Chat
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(
                    label="Conversation",
                    elem_id="chatbot",
                    bubble_full_width=False,
                    show_copy_button=True,
                )

                with gr.Row():
                    user_input = gr.Textbox(
                        placeholder="Ask about admissions, programs, fees, deadlines…",
                        label="Your question",
                        lines=2,
                        scale=5,
                    )
                    send_btn = gr.Button("Send 📨", variant="primary", scale=1)

                with gr.Row():
                    clear_btn    = gr.Button("🗑️ Clear Chat",     variant="secondary")
                    export_btn   = gr.Button("💾 Export History", variant="secondary")
                    export_file  = gr.File(label="Download transcript", visible=False)

            # Right column — Controls
            with gr.Column(scale=1, min_width=280):

                with gr.Accordion("⚙️ Generation Settings", open=True):
                    temp_slider = gr.Slider(
                        minimum=0.1, maximum=1.0, value=cfg.model.temperature,
                        step=0.05, label="Temperature",
                    )
                    max_tokens_slider = gr.Slider(
                        minimum=64, maximum=1024, value=cfg.model.max_new_tokens,
                        step=32, label="Max New Tokens",
                    )
                    top_p_slider = gr.Slider(
                        minimum=0.5, maximum=1.0, value=cfg.model.top_p,
                        step=0.05, label="Top-p",
                    )
                    retrieval_k_slider = gr.Slider(
                        minimum=1, maximum=10, value=cfg.retrieval.top_k,
                        step=1, label="Retrieval Top-k",
                    )

                with gr.Accordion("📚 Knowledge Base", open=True):
                    file_upload = gr.File(
                        label="Upload Documents (PDF, DOCX, TXT)",
                        file_count="multiple",
                        file_types=[".pdf", ".docx", ".txt"],
                    )

                    chunk_size_slider = gr.Slider(
                        minimum=200, maximum=2000, value=cfg.chunking.chunk_size,
                        step=50, label="Chunk Size",
                    )
                    chunk_overlap_slider = gr.Slider(
                        minimum=0, maximum=400, value=cfg.chunking.chunk_overlap,
                        step=10, label="Chunk Overlap",
                    )

                    with gr.Row():
                        upload_btn  = gr.Button("⬆️ Upload", variant="secondary")
                        build_btn   = gr.Button("🔨 Build KB", variant="primary")
                        reload_btn  = gr.Button("🔄 Reload", variant="secondary")

                    kb_status_box = gr.Textbox(
                        label="Knowledge Base Status",
                        value="\n".join(get_kb_status()),
                        interactive=False,
                        lines=3,
                    )
                    uploaded_files_box = gr.Textbox(
                        label="Uploaded Files",
                        value=list_uploaded_files(),
                        interactive=False,
                        lines=4,
                    )

                with gr.Accordion("ℹ️ System Information", open=False):
                    gr.Markdown(
                        f"**Generation model:** `{cfg.model.model_id}`\n\n"
                        f"**Embedding model:** `{cfg.embedding.model_name}`\n\n"
                        f"**Vector database:** FAISS (LangChain)\n\n"
                        f"**Framework:** Gradio {gr.__version__}"
                    )

        # ── Event wiring ──────────────────────────────────────────────────────

        gen_inputs = [
            chatbot,
            temp_slider,
            max_tokens_slider,
            top_p_slider,
            retrieval_k_slider,
        ]

        def _send(msg, hist, temp, mtok, tp, rk):
            return chat_fn(msg, hist, temp, mtok, tp, rk)

        send_btn.click(
            fn=_send,
            inputs=[user_input] + gen_inputs,
            outputs=[chatbot, user_input],
        )
        user_input.submit(
            fn=_send,
            inputs=[user_input] + gen_inputs,
            outputs=[chatbot, user_input],
        )

        clear_btn.click(
            fn=clear_chat_fn,
            outputs=[chatbot, user_input],
        )

        def _export(hist):
            path = export_fn(hist)
            return gr.update(value=path, visible=True)

        export_btn.click(
            fn=_export,
            inputs=[chatbot],
            outputs=[export_file],
        )

        def _upload(files):
            msg = upload_files_fn(files)
            return msg, list_uploaded_files()

        upload_btn.click(
            fn=_upload,
            inputs=[file_upload],
            outputs=[kb_status_box, uploaded_files_box],
        )

        def _build(cs, co):
            msg = build_kb_fn(cs, co)
            s, d, c = get_kb_status()
            return f"{msg}\n\nIndex status: {s} | Docs: {d} | Chunks: {c}"

        build_btn.click(
            fn=_build,
            inputs=[chunk_size_slider, chunk_overlap_slider],
            outputs=[kb_status_box],
        )

        def _reload():
            msg = reload_kb_fn()
            return msg, list_uploaded_files()

        reload_btn.click(
            fn=_reload,
            outputs=[kb_status_box, uploaded_files_box],
        )

    return demo


# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    demo = build_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        show_error=True,
        share=False,   # set True for a public Gradio link (Colab convenience)
    )