gradio_app.py

"""Gradio entrypoint for ProfillyBot on Hugging Face Spaces (ZeroGPU).

Compatible with Gradio 6.x (messages format is default; css goes to launch()).
Shows answer + retrieval debug panel (query + retrieved chunks).
"""

from __future__ import annotations

import logging
import os
import sys
from pathlib import Path

# Ensure project root is on path
ROOT = Path(__file__).parent.resolve()
sys.path.insert(0, str(ROOT))

import gradio as gr

from src.build_vectorstore import build_index
from src.config_loader import get_config
from src.llm_handler import get_llm_handler
from src.rag_pipeline import RAGPipeline

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger(__name__)

# Optional ZeroGPU decorator (no-op locally when `spaces` is unavailable)
try:
    import spaces

    gpu = spaces.GPU(duration=120)
except Exception:  # noqa: BLE001

    def gpu(fn):  # type: ignore[misc]
        return fn


_rag_pipeline: RAGPipeline | None = None


def _indexes_ready() -> bool:
    chroma_db = ROOT / "chroma_db" / "chroma.sqlite3"
    bm25_dir = ROOT / "bm25_index"
    has_bm25 = bm25_dir.exists() and any(bm25_dir.iterdir())
    return chroma_db.exists() or has_bm25


def _ensure_indexes() -> None:
    """Build retrieval indexes on CPU if they are missing."""
    if _indexes_ready():
        logger.info("Retrieval indexes already present")
        return

    docs_dir = ROOT / "data" / "documents"
    if not docs_dir.exists() or not any(docs_dir.iterdir()):
        logger.warning("No documents found; skipping index build")
        return

    logger.info("Building retrieval indexes (CPU)...")
    build_index(documents_dir=str(docs_dir), force_rebuild=False)


def get_pipeline() -> RAGPipeline:
    """Lazily initialize the RAG pipeline (embeddings/retrieval on CPU)."""
    global _rag_pipeline
    if _rag_pipeline is not None:
        return _rag_pipeline

    _ensure_indexes()

    if os.getenv("SPACE_ID") or os.getenv("SYSTEM") == "spaces":
        provider = "transformers"
    else:
        provider = get_config().get("llm.provider", "transformers")

    handler = get_llm_handler(provider_override=provider)
    _rag_pipeline = RAGPipeline(llm_handler=handler)
    logger.info(
        "ProfillyBot ready — provider=%s model=%s",
        handler.get_provider(),
        handler.get_model(),
    )
    return _rag_pipeline


def _message_text(content: object) -> str | None:
    """Normalize Gradio 5/6 message content to a plain string."""
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts: list[str] = []
        for part in content:
            if isinstance(part, str):
                parts.append(part)
            elif isinstance(part, dict):
                if part.get("type") == "text":
                    parts.append(str(part.get("text", "")))
                elif "text" in part:
                    parts.append(str(part["text"]))
        return "".join(parts) if parts else None
    return None


@gpu
def answer_with_retrieval(
    message: str, history: list[dict] | None
) -> tuple[list[dict], str, str]:
    """Generate answer and retrieval panel for Gradio UI."""
    empty_retrieval = "_Ask a question to see retrieved chunks._"
    if not message or not message.strip():
        history = history or []
        return history, "Please ask a question about the profile.", empty_retrieval

    pipeline = get_pipeline()
    question = message.strip()

    chat_history: list[dict] = []
    if history:
        for turn in history:
            role = turn.get("role")
            text = _message_text(turn.get("content"))
            if role in {"user", "assistant"} and text:
                chat_history.append({"role": role, "content": text})

    llm = pipeline.llm_handler.get_llm()
    if hasattr(llm, "_ensure_loaded"):
        llm._ensure_loaded()

    result = pipeline.get_answer_with_retrieval(question, chat_history=chat_history)
    answer = result["answer"]
    retrieval_panel = result["retrieval_panel"]

    new_history = list(history or [])
    new_history.append({"role": "user", "content": question})
    new_history.append({"role": "assistant", "content": answer})
    return new_history, "", retrieval_panel


CHAT_CSS = """
.contain textarea,
footer textarea,
.input-container textarea,
textarea[data-testid="textbox"] {
  overflow-y: hidden !important;
  resize: none !important;
  scrollbar-width: none !important;
  -ms-overflow-style: none !important;
}
.contain textarea::-webkit-scrollbar,
footer textarea::-webkit-scrollbar,
.input-container textarea::-webkit-scrollbar,
textarea[data-testid="textbox"]::-webkit-scrollbar {
  display: none !important;
  width: 0 !important;
  height: 0 !important;
}
"""


def build_demo() -> gr.Blocks:
    config = get_config()
    profile_name = config.get("profile.name", "the candidate")
    profile_title = config.get("profile.title", "Professional")
    model_name = config.get("llm.hf_model", "Qwen/Qwen3.5-4B")
    greeting = config.format_template(
        config.get(
            "profile.greeting",
            f"Hi! I'm ProfillyBot. Ask me about {profile_name}'s background.",
        )
    )
    examples = [config.format_template(q) for q in config.get("ui.example_questions", [])]

    with gr.Blocks(title="ProfillyBot") as demo:
        gr.Markdown(
            f"# ProfillyBot\n"
            f"**{profile_name}** — {profile_title}\n\n"
            f"{greeting}\n\n"
            f"_{model_name} on Hugging Face ZeroGPU + hybrid RAG._"
        )

        with gr.Row():
            with gr.Column(scale=3):
                chatbot = gr.Chatbot(label="Chat", height=460)
                msg = gr.Textbox(
                    placeholder="Type a message...",
                    lines=1,
                    max_lines=1,
                    show_label=False,
                    autofocus=True,
                )
                with gr.Row():
                    send = gr.Button("Send", variant="primary")
                    clear = gr.Button("Clear")
                if examples:
                    gr.Examples(examples=examples, inputs=msg)

            with gr.Column(scale=2):
                retrieval_md = gr.Markdown(
                    value="_Ask a question to see retrieved chunks._",
                    label="Retrieval",
                )

        send.click(
            answer_with_retrieval,
            inputs=[msg, chatbot],
            outputs=[chatbot, msg, retrieval_md],
        )
        msg.submit(
            answer_with_retrieval,
            inputs=[msg, chatbot],
            outputs=[chatbot, msg, retrieval_md],
        )
        clear.click(
            lambda: ([], "", "_Ask a question to see retrieved chunks._"),
            outputs=[chatbot, msg, retrieval_md],
        )

    return demo


demo = build_demo()

demo.queue(max_size=20).launch(
    server_name="0.0.0.0",
    server_port=7860,
    css=CHAT_CSS,
)