app.py

"""user interface for RAG"""

from __future__ import annotations

import os
import threading
from typing import Any

import gradio as gr
import spaces

from config import SETTINGS
from rag_engine import RAGEngine


ENGINE = RAGEngine()
ENGINE_LOCK = threading.Lock()


CSS = r"""
:root {
  --surface: rgba(11, 17, 31, .76);
  --surface-2: rgba(20, 29, 50, .72);
  --line: rgba(148, 163, 184, .16);
  --muted: #9ba9bd;
  --text: #f7f9fc;
  --accent: #8b5cf6;
  --accent-2: #22d3ee;
}

.gradio-container {
  max-width: 1480px !important;
  margin: 0 auto !important;
  color: var(--text) !important;
  background:
    radial-gradient(
      circle at 8% 8%,
      rgba(139, 92, 246, .20),
      transparent 29%
    ),
    radial-gradient(
      circle at 91% 12%,
      rgba(34, 211, 238, .14),
      transparent 26%
    ),
    linear-gradient(
      145deg,
      #050811 0%,
      #080d19 46%,
      #0b1020 100%
    ) !important;
  min-height: 100vh;
}

.main-shell {
  padding: 28px 24px 44px;
}

.hero {
  position: relative;
  overflow: hidden;
  border: 1px solid var(--line);
  background:
    linear-gradient(
      135deg,
      rgba(19, 27, 48, .94),
      rgba(10, 15, 29, .84)
    );
  border-radius: 24px;
  padding: 30px 32px;
  box-shadow: 0 28px 80px rgba(0, 0, 0, .28);
  margin-bottom: 18px;
}

.hero::after {
  content: "";
  position: absolute;
  width: 330px;
  height: 330px;
  right: -110px;
  top: -160px;
  background:
    radial-gradient(
      circle,
      rgba(34, 211, 238, .22),
      transparent 65%
    );
}

.eyebrow {
  color: #b8a6ff;
  font-size: 12px;
  font-weight: 800;
  letter-spacing: .16em;
  text-transform: uppercase;
}

.hero h1 {
  margin: 8px 0 7px;
  font-size: clamp(32px, 5vw, 57px);
  line-height: 1.02;
  letter-spacing: -.045em;
}

.hero p {
  max-width: 850px;
  color: #b8c3d4;
  font-size: 16px;
  line-height: 1.65;
  margin: 0;
}

.badges {
  display: flex;
  flex-wrap: wrap;
  gap: 9px;
  margin-top: 18px;
}

.badge {
  border: 1px solid var(--line);
  background: rgba(255, 255, 255, .035);
  padding: 7px 11px;
  border-radius: 999px;
  color: #cbd5e1;
  font-size: 12px;
}

.badge strong {
  color: #fff;
}

.panel,
.gr-panel,
.block {
  border-color: var(--line) !important;
}

.app-panel {
  background: var(--surface) !important;
  border: 1px solid var(--line) !important;
  border-radius: 20px !important;
  box-shadow: 0 18px 48px rgba(0, 0, 0, .20);
}

.sidebar-card {
  background: var(--surface-2);
  border: 1px solid var(--line);
  border-radius: 18px;
  padding: 18px;
  margin-bottom: 14px;
}

.sidebar-card h3 {
  margin: 0 0 8px;
  font-size: 14px;
}

.sidebar-card p {
  margin: 0;
  color: var(--muted);
  font-size: 13px;
  line-height: 1.55;
}

#chatbot {
  min-height: 520px;
}

#chatbot .message {
  border-radius: 16px !important;
}

#prompt textarea {
  font-size: 15px !important;
  line-height: 1.5 !important;
}

#send-button {
  min-width: 115px;
  font-weight: 800;
}

.source-list {
  display: grid;
  gap: 10px;
}

.source-card {
  border: 1px solid var(--line);
  border-radius: 14px;
  background: rgba(255, 255, 255, .025);
  overflow: hidden;
}

.source-card summary {
  cursor: pointer;
  list-style: none;
  display: grid;
  grid-template-columns: 28px 1fr auto;
  gap: 9px;
  align-items: center;
  padding: 12px 13px;
}

.source-card summary::-webkit-details-marker {
  display: none;
}

.source-number {
  display: grid;
  place-items: center;
  width: 24px;
  height: 24px;
  border-radius: 8px;
  background:
    linear-gradient(
      135deg,
      var(--accent),
      var(--accent-2)
    );
  color: white;
  font-size: 11px;
  font-weight: 900;
}

.source-title {
  color: #e5eaf2;
  font-size: 13px;
  font-weight: 700;
  overflow: hidden;
  text-overflow: ellipsis;
  white-space: nowrap;
}

.source-score {
  font-size: 11px;
  color: #8cdff0;
  font-variant-numeric: tabular-nums;
}

.source-body {
  padding: 0 14px 12px 50px;
  color: #b7c2d4;
  font-size: 12px;
  line-height: 1.65;
}

.source-meta {
  padding: 9px 14px;
  border-top: 1px solid var(--line);
  color: #75859c;
  font-size: 10px;
}

.empty-state {
  color: var(--muted);
  padding: 18px;
  text-align: center;
  border: 1px dashed var(--line);
  border-radius: 14px;
}

.footer-note {
  color: #718097;
  font-size: 11px;
  text-align: center;
  margin-top: 16px;
}

button.primary {
  background:
    linear-gradient(
      135deg,
      #7c3aed,
      #0891b2
    ) !important;
  border: none !important;
}

.accordion {
  background: rgba(255, 255, 255, .02) !important;
  border-color: var(--line) !important;
}

@media (max-width: 800px) {
  .main-shell {
    padding: 14px 10px 28px;
  }

  .hero {
    padding: 23px 20px;
    border-radius: 18px;
  }

  #chatbot {
    min-height: 430px;
  }
}
"""


HEAD = """
<meta name="theme-color" content="#070b15">
<meta
  name="description"
  content="Hybrid retrieval-augmented generation with Hugging Face models and datasets."
>
"""


def get_engine() -> RAGEngine:
    """Initialize and return the shared RAG engine."""

    if not ENGINE.ready:
        with ENGINE_LOCK:
            if not ENGINE.ready:
                ENGINE.initialize()

    return ENGINE


@spaces.GPU(duration=60)
def run_chat(
    message: str,
    history: list[dict[str, Any]] | None,
    top_k: int,
    dense_weight: float,
    use_reranker: bool,
    temperature: float,
    max_tokens: int,
):
    """Process a question through the RAG pipeline."""

    clean_message = (message or "").strip()
    history = list(history or [])

    if not clean_message:
        return (
            history,
            "<div class='empty-state'>Enter a question first.</div>",
            {},
            "",
        )

    engine = get_engine()

    prior_history = [
        {
            "role": str(item.get("role")),
            "content": str(item.get("content", "")),
        }
        for item in history
        if isinstance(item, dict)
    ]

    answer, results = engine.answer(
        query=clean_message,
        history=prior_history,
        top_k=int(top_k),
        dense_weight=float(dense_weight),
        use_reranker=bool(use_reranker),
        temperature=float(temperature),
        max_tokens=int(max_tokens),
    )

    history.extend(
        [
            {
                "role": "user",
                "content": clean_message,
            },
            {
                "role": "assistant",
                "content": answer,
            },
        ]
    )

    return (
        history,
        engine.render_sources(results),
        engine.diagnostics(results),
        "",
    )


def clear_all():
    """Clear the conversation and retrieved evidence."""

    return (
        [],
        "<div class='empty-state'>No sources retrieved yet.</div>",
        {},
        "",
    )


def build_app() -> gr.Blocks:
    """Build and return the Gradio application."""

    token_state = "configured" if SETTINGS.hf_token else "missing"

    hero = f"""
    <div class="hero">
      <div class="eyebrow">RAG</div>

      <h1>RAG with Wikipedia</h1>

      <p>
        Ask an RAG model questions trained on
        <code>rag-datasets/rag-mini-wikipedia</code>.
      </p>

      <div class="badges">
        <span class="badge">
          <strong>Generator</strong>
          Qwen3 4B Instruct
        </span>

        <span class="badge">
          <strong>Embeddings</strong>
          MiniLM-L6
        </span>

        <span class="badge">
          <strong>Dataset</strong>
          rag-mini-wikipedia
        </span>

        <span class="badge">
          <strong>HF token</strong>
          {token_state}
        </span>
      </div>
    </div>
    """

    with gr.Blocks(title="RAG Wikipedia") as demo:
        with gr.Column(elem_classes=["main-shell"]):
            gr.HTML(hero)

            with gr.Row(equal_height=False):
                with gr.Column(
                    scale=8,
                    min_width=520,
                    elem_classes=["app-panel"],
                ):
                    chatbot = gr.Chatbot(
                        value=[],
                        label="Grounded conversation",
                        elem_id="chatbot",
                        height=565,
                        buttons=["copy"],
                        placeholder=(
                            "Ask a factual question about the indexed "
                            "Wikipedia passages."
                        ),
                    )

                    with gr.Row():
                        prompt = gr.Textbox(
                            label="Question",
                            placeholder=(
                                "Example: What did the Legal Tender Act "
                                "of 1862 establish?"
                            ),
                            lines=2,
                            max_lines=5,
                            elem_id="prompt",
                            scale=8,
                        )

                        send = gr.Button(
                            "Ask",
                            variant="primary",
                            elem_id="send-button",
                            scale=1,
                        )

                    with gr.Row():
                        clear = gr.Button("Clear conversation")

                    gr.Examples(
                        examples=[
                            (
                                "What did the Legal Tender Act of 1862 "
                                "establish?"
                            ),
                            "What is Amedeo Avogadro most noted for?",
                            "What is the study of beetles called?",
                            (
                                "What happened to the Celsius temperature "
                                "scale in 1745?"
                            ),
                        ],
                        inputs=prompt,
                        label="Suggested questions",
                    )

                with gr.Column(
                    scale=5,
                    min_width=360,
                ):
                    gr.HTML(
                        """
                        <div class="sidebar-card">
                          <h3>Retrieved evidence</h3>

                          <p>
                            Answers use numbered citations that correspond
                            to the passages below. Open a source to inspect
                            the exact evidence.
                          </p>
                        </div>
                        """
                    )

                    sources = gr.HTML(
                        (
                            "<div class='empty-state'>"
                            "No sources retrieved yet."
                            "</div>"
                        ),
                        label="Sources",
                        elem_classes=["app-panel"],
                    )

                    with gr.Accordion(
                        "Retrieval controls",
                        open=True,
                        elem_classes=["accordion"],
                    ):
                        top_k = gr.Slider(
                            minimum=3,
                            maximum=8,
                            value=5,
                            step=1,
                            label="Sources returned",
                        )

                        dense_weight = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            value=0.72,
                            step=0.05,
                            label="Semantic retrieval weight",
                            info=(
                                "The remaining weight is assigned to "
                                "TF-IDF lexical retrieval."
                            ),
                        )

                        use_reranker = gr.Checkbox(
                            value=True,
                            label="Use neural reranker",
                        )

                    with gr.Accordion(
                        "Generation controls",
                        open=False,
                        elem_classes=["accordion"],
                    ):
                        temperature = gr.Slider(
                            minimum=0.0,
                            maximum=1.0,
                            value=0.2,
                            step=0.05,
                            label="Temperature",
                        )

                        max_tokens = gr.Slider(
                            minimum=200,
                            maximum=1200,
                            value=700,
                            step=50,
                            label="Maximum output tokens",
                        )

                    with gr.Accordion(
                        "Diagnostics",
                        open=False,
                        elem_classes=["accordion"],
                    ):
                        diagnostics = gr.JSON(
                            value={},
                            label="Pipeline diagnostics",
                        )

            gr.HTML(
                """
                <div class="footer-note">
                  Retrieval runs locally. Generation uses Hugging Face
                  Inference Providers and requires an HF token with
                  available credits.
                </div>
                """
            )

        inputs = [
            prompt,
            chatbot,
            top_k,
            dense_weight,
            use_reranker,
            temperature,
            max_tokens,
        ]

        outputs = [
            chatbot,
            sources,
            diagnostics,
            prompt,
        ]

        send.click(
            fn=run_chat,
            inputs=inputs,
            outputs=outputs,
        )

        prompt.submit(
            fn=run_chat,
            inputs=inputs,
            outputs=outputs,
        )

        clear.click(
            fn=clear_all,
            outputs=outputs,
        )

    return demo


if __name__ == "__main__":
    app = build_app()

    app.queue(
        default_concurrency_limit=4,
    ).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        show_error=True,
        theme=gr.themes.Base(),
        css=CSS,
        head=HEAD,
    )