app.py

import json
import os
import re
from dataclasses import dataclass, field
from datetime import date
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple, Union

import gradio as gr
import requests
from bs4 import BeautifulSoup
from duckduckgo_search import DDGS
from huggingface_hub import InferenceClient


# --- Model configuration ---------------------------------------------------
# Our own DeepResearch model. When QUEST_BASE_URL is configured in Space
# Secrets, the app will route requests to that dedicated HF Inference Endpoint
# instead of the shared HF Inference API.
QUEST_MODEL_ID = "osunlp/Quest-4B"
QUEST_BASE_URL = os.getenv("QUEST_BASE_URL", "").strip()
# Endpoints built from the TGI image expose a single-model OpenAI route; the
# model name passed to chat_completion is usually "tgi". vLLM endpoints usually
# want the original repo id. QUEST_ENDPOINT_MODEL overrides this if needed.
QUEST_ENDPOINT_MODEL = os.getenv("QUEST_ENDPOINT_MODEL", "tgi").strip() or "tgi"

# Shared HF Inference API fallbacks (free, rate-limited). These are used when
# the user picks one of these from the Model dropdown; they do NOT go through
# the private endpoint.
FREE_FALLBACK_MODELS = [
    "Qwen/Qwen3-8B",
    "google/gemma-3-12b-it",
    "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
    "Qwen/Qwen2.5-7B-Instruct",
    "meta-llama/Llama-3.1-8B-Instruct",
]

# Quest-4B shows up first when the endpoint is wired; otherwise we still list
# it so you can see what the target model is, but it will only work after the
# QUEST_BASE_URL secret is configured.
DEFAULT_MODEL_CHOICES = [QUEST_MODEL_ID] + FREE_FALLBACK_MODELS
DEFAULT_MODEL = os.getenv(
    "DEFAULT_MODEL",
    QUEST_MODEL_ID if QUEST_BASE_URL else FREE_FALLBACK_MODELS[0],
)

PAPER_URL = os.getenv("PAPER_URL", "#")
CODE_URL = os.getenv("CODE_URL", "#")
DATASET_URL = os.getenv("DATASET_URL", "#")
MODEL_URL = os.getenv("MODEL_URL", "#")


# --- System prompt ---------------------------------------------------------
# Full QUEST SYSTEM_PROMPT (mirrors inference/prompt.py in the research repo)
# so that Quest-4B sees the exact tool schema it was trained with. Other
# models still follow this schema just fine in practice.
QUEST_SYSTEM_PROMPT = """You are a deep research assistant. Your core function is to conduct thorough, multi-source investigations into any topic. You must handle both broad, open-domain inquiries and queries within specialized academic fields. For every request, synthesize information from credible, diverse sources to deliver a comprehensive, accurate, and objective response. When you have gathered sufficient information and are ready to provide the definitive response, you must enclose the entire final answer within <answer></answer> tags.

# Tools

You may call one or more functions to assist with the user query.

You are provided with function signatures within <tools></tools> XML tags:
<tools>
{"type": "function", "function": {"name": "search", "description": "Perform Google web searches then returns a string of the top search results. Accepts multiple queries.", "parameters": {"type": "object", "properties": {"query": {"type": "array", "items": {"type": "string", "description": "The search query."}, "minItems": 1, "description": "The list of search queries."}}, "required": ["query"]}}}
{"type": "function", "function": {"name": "visit", "description": "Visit webpage(s) and return the summary of the content.", "parameters": {"type": "object", "properties": {"url": {"type": "array", "items": {"type": "string"}, "description": "The URL(s) of the webpage(s) to visit. Can be a single URL or an array of URLs."}, "goal": {"type": "string", "description": "The specific information goal for visiting webpage(s)."}}, "required": ["url", "goal"]}}}
</tools>

# Using prev_state (Research State Summary)

If you see a "RESEARCH STATE SUMMARY (prev_state)" section in the user message, it contains a compressed summary of previous research progress. Use it to avoid repeating searches/visits that have already been executed, use verified information directly in your answer, and follow up on uncertain claims only when needed.

For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
<tool_call>
{"name": <function-name>, "arguments": <args-json-object>}
</tool_call>

Current date: """


def build_system_prompt() -> str:
    return QUEST_SYSTEM_PROMPT + date.today().isoformat()


TOOL_RESPONSE_TEMPLATE = """<tool_response>
{payload}
</tool_response>"""

SEARCH_CACHE: Dict[str, Dict[str, Any]] = {}
VISIT_CACHE: Dict[str, Dict[str, Any]] = {}
ASSETS_DIR = Path(__file__).resolve().parent / "assets"
LOGO_PATH = str(ASSETS_DIR / "quest-logo.png")
OSU_NLP_LOGO_PATH = str(ASSETS_DIR / "osu-nlp-logo.png")
OSU_NLP_URL = "https://nlp.osu.edu/"

# Gradio 5 default themes pull greys from the neutral palette; Base + explicit overrides keep
# the whole UI on a strict blue/white scheme without unwanted slate panels.
APP_THEME = gr.themes.Base(
    primary_hue=gr.themes.colors.blue,
    secondary_hue=gr.themes.colors.sky,
    neutral_hue=gr.themes.colors.blue,
    font=[
        gr.themes.GoogleFont("Plus Jakarta Sans"),
        "ui-sans-serif",
        "system-ui",
        "sans-serif",
    ],
    font_mono=[
        gr.themes.GoogleFont("JetBrains Mono"),
        "ui-monospace",
        "monospace",
    ],
).set(
    body_background_fill="#f5f9ff",
    body_text_color="#0f2744",
    body_text_color_subdued="#4a6a8c",
    color_accent="*primary_600",
    color_accent_soft="#dbeafe",
    background_fill_primary="#ffffff",
    background_fill_secondary="#f0f6ff",
    border_color_primary="#dbeafe",
    border_color_accent="*primary_500",
    block_background_fill="#ffffff",
    block_border_width="0px",
    block_border_color="transparent",
    block_shadow="0 1px 2px rgba(15,39,68,0.04), 0 10px 30px rgba(37,99,235,0.06)",
    block_radius="18px",
    block_label_background_fill="transparent",
    block_label_border_width="0px",
    block_label_text_color="#3b5b7a",
    block_label_text_weight="600",
    block_title_text_color="#0f2744",
    block_title_text_weight="700",
    block_title_border_width="0px",
    panel_background_fill="transparent",
    panel_border_width="0px",
    panel_border_color="transparent",
    input_background_fill="#ffffff",
    input_background_fill_focus="#ffffff",
    input_border_color="#bfdbfe",
    input_border_color_focus="*primary_500",
    input_border_width="1px",
    input_radius="14px",
    input_shadow="none",
    input_shadow_focus="0 0 0 3px rgba(59,130,246,0.18)",
    code_background_fill="#f0f7ff",
    slider_color="*primary_500",
    button_primary_background_fill="linear-gradient(135deg,#1d4ed8 0%,#3b82f6 100%)",
    button_primary_background_fill_hover="linear-gradient(135deg,#1e40af 0%,#2563eb 100%)",
    button_primary_text_color="#ffffff",
    button_primary_border_color="transparent",
    button_primary_shadow="0 8px 22px rgba(37,99,235,0.28)",
    button_secondary_background_fill="#ffffff",
    button_secondary_background_fill_hover="#eff6ff",
    button_secondary_text_color="*primary_700",
    button_secondary_border_color="#bfdbfe",
    button_cancel_background_fill="#ffffff",
    button_cancel_background_fill_hover="#eff6ff",
    button_cancel_text_color="*primary_700",
    button_cancel_border_color="#bfdbfe",
    table_border_color="#dbeafe",
    table_even_background_fill="#fafdff",
    table_odd_background_fill="#ffffff",
)

CUSTOM_CSS = """
/* Gradio 5 uses versioned root classes (gradio-container-5-29-0). Match all of them and
   replace every neutral grey surface with white / soft-blue tints. */

html, body, gradio-app, [class*="gradio-container"] {
  background: #f5f9ff !important;
}

/* HF Space iframe wraps Gradio in <gradio-app>. Force every wrapper to stretch to
   the full viewport from the FIRST paint so the page doesn't visibly grow after
   the first answer arrives. */
html, body { width: 100% !important; min-height: 100vh !important; margin: 0 !important; }
gradio-app {
  display: block !important;
  width: 100% !important;
  min-height: 100vh !important;
  margin-left: auto !important;
  margin-right: auto !important;
}
gradio-app > .gradio-container,
gradio-app > div {
  display: block !important;
  width: 100% !important;
  margin-left: auto !important;
  margin-right: auto !important;
}

[class*="gradio-container"] {
  max-width: 1480px !important;
  width: 100% !important;
  min-width: 320px !important;
  margin-left: auto !important;
  margin-right: auto !important;
  padding-left: 28px !important;
  padding-right: 28px !important;
  color: #0f2744;
  box-sizing: border-box !important;
}

/* Lock the inner two-column row so the right panel doesn't shrink before content
   arrives, then snap back when results appear. */
[class*="gradio-container"] .layout-gap {
  width: 100% !important;
}
[class*="gradio-container"] .layout-gap > .gr-column,
[class*="gradio-container"] .layout-gap > div {
  min-width: 0 !important;
}

/* Reserve vertical space for the Result/Record area so the first answer doesn't
   trigger a visible vertical jump either. */
[class*="gradio-container"] .gradio-markdown,
[class*="gradio-container"] [data-testid="markdown"] {
  min-height: 220px !important;
}
[class*="gradio-container"] .codemirror-wrapper,
[class*="gradio-container"] .cm-editor {
  min-height: 220px !important;
}

/* Prevent Code / Markdown / Tabs from pushing the page wider than the container.
   Every wrapper in the chain is locked to max-width: 100%; only the innermost
   .cm-scroller scrolls horizontally. */
[class*="gradio-container"] .gradio-code,
[class*="gradio-container"] .gradio-markdown,
[class*="gradio-container"] .prose,
[class*="gradio-container"] .markdown,
[class*="gradio-container"] [data-testid="markdown"],
[class*="gradio-container"] .tabs,
[class*="gradio-container"] .tabitem,
[class*="gradio-container"] .tab-content {
  max-width: 100% !important;
  width: 100% !important;
  min-width: 0 !important;
  word-wrap: break-word !important;
  overflow-wrap: anywhere !important;
}
[class*="gradio-container"] .codemirror-wrapper {
  max-width: 100% !important;
  width: 100% !important;
  min-width: 0 !important;
  border-radius: 14px !important;
  overflow: hidden !important;
}
[class*="gradio-container"] .cm-editor {
  max-width: 100% !important;
  width: 100% !important;
  min-width: 0 !important;
  overflow: hidden !important;
}
[class*="gradio-container"] .cm-scroller {
  max-width: 100% !important;
  width: 100% !important;
  overflow-x: auto !important;
}
[class*="gradio-container"] .cm-content,
[class*="gradio-container"] .cm-line {
  max-width: 100% !important;
  white-space: pre-wrap !important;
  word-break: break-word !important;
}
[class*="gradio-container"] .prose pre,
[class*="gradio-container"] .markdown pre {
  max-width: 100% !important;
  overflow-x: auto !important;
  white-space: pre-wrap !important;
}

/* Some Gradio builds wrap the app in `.main` or `.app`; center those too. */
[class*="gradio-container"] .main,
[class*="gradio-container"] .app,
[class*="gradio-container"] .contain {
  margin-left: auto !important;
  margin-right: auto !important;
}

[class*="gradio-container"] *::selection { background: rgba(37,99,235,0.18); }

/* --- TOP BANNER --- */
.top-banner {
  align-items: center !important;
  padding: 28px 0 16px 0;
  margin-bottom: 8px;
  gap: 0 !important;
}
.top-banner .banner-side { min-width: 0; }
.top-banner .banner-center {
  display: flex !important;
  flex-direction: column !important;
  align-items: center !important;
  gap: 10px !important;
}
.banner-subtitle {
  color: #4a6a8c;
  font-size: 15px;
  font-weight: 500;
  letter-spacing: 0.01em;
  text-align: center;
  margin: 0;
}

/* Shared logo chrome: both Quest and OSU NLP get the same rounded frame so they
   look like a pair. Height is fixed so their visual weight matches. */
.banner-quest-logo,
.banner-quest-logo .image-container,
.banner-quest-logo .image-frame,
.banner-quest-logo > div,
.banner-quest-logo button,
.osu-nlp-logo,
.osu-nlp-logo .image-container,
.osu-nlp-logo .image-frame,
.osu-nlp-logo > div,
.osu-nlp-logo button {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
}
.banner-quest-logo .icon-button-wrapper,
.banner-quest-logo [aria-label*="hare" i],
.banner-quest-logo [aria-label*="ownload" i],
.banner-quest-logo [aria-label*="ullscreen" i],
.osu-nlp-logo .icon-button-wrapper,
.osu-nlp-logo [aria-label*="hare" i],
.osu-nlp-logo [aria-label*="ownload" i],
.osu-nlp-logo [aria-label*="ullscreen" i] {
  display: none !important;
}
/* Both logos share the same fixed height so they visually line up. Width is
   auto so each logo keeps its natural aspect ratio. */
.banner-quest-logo img,
.osu-nlp-logo img {
  height: 140px !important;
  width: auto !important;
  max-width: 100% !important;
  object-fit: contain !important;
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  border-radius: 0 !important;
  padding: 0 !important;
  margin: 0 auto !important;
  display: block !important;
}
.banner-right {
  display: flex !important;
  justify-content: flex-end !important;
  align-items: center !important;
}
.osu-nlp-logo img { margin-left: auto !important; }
.osu-nlp-logo {
  cursor: pointer;
  transition: transform .15s ease;
}
.osu-nlp-logo:hover img { transform: translateY(-1px); }

/* --- LEFT/RIGHT layout --- */
.layout-gap { gap: 24px !important; align-items: flex-start; }
.right-stack > * { margin-bottom: 12px; }
.action-row { gap: 10px !important; margin-top: 12px; }
.action-row button { min-width: 0; flex: 1; }

.hero-heading {
  font-size: 1.1rem !important;
  font-weight: 700 !important;
  letter-spacing: 0.005em !important;
  text-transform: none !important;
  color: #0f2744 !important;
  margin-bottom: 14px !important;
}

/* --- SECTION CARDS --- */
/* `section-card` becomes a real white rounded card with soft blue shadow, no grey. */
.section-card {
  background: #ffffff !important;
  border: 1px solid rgba(191,219,254,0.55) !important;
  border-radius: 20px !important;
  box-shadow: 0 1px 2px rgba(15,39,68,0.03), 0 16px 40px rgba(37,99,235,0.07) !important;
  padding: 18px !important;
}
/* `no-frame` opts a card out of all chrome (used for logo + icon grid). */
.no-frame {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  padding: 0 !important;
}

.section-heading {
  font-size: 0.72rem;
  font-weight: 700;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: #2563eb;
  margin: 0 0 12px 0;
}

/* --- ICON GRID (Paper / Code / Dataset / Model) --- */
.icon-grid {
  display: grid;
  grid-template-columns: repeat(2, minmax(0, 1fr));
  gap: 12px;
  width: 100%;
  margin: 0;
  box-sizing: border-box;
}
.icon-link {
  display: flex;
  align-items: center;
  justify-content: center;
  gap: 8px;
  padding: 14px 10px;
  border-radius: 16px;
  text-decoration: none !important;
  color: #1d4ed8 !important;
  background: #ffffff;
  font-weight: 600;
  font-size: 15px;
  white-space: nowrap;
  border: 1px solid rgba(191,219,254,0.7);
  box-shadow: 0 1px 2px rgba(15,39,68,0.03), 0 10px 28px rgba(37,99,235,0.07);
  transition: transform .15s ease, box-shadow .15s ease, border-color .15s ease;
}
.icon-link:hover {
  transform: translateY(-1px);
  border-color: #93c5fd;
  box-shadow: 0 4px 18px rgba(37,99,235,0.14);
}

/* --- BUTTONS --- */
[class*="gradio-container"] button.primary,
[class*="gradio-container"] .gr-button-primary {
  background: linear-gradient(135deg,#1d4ed8 0%,#3b82f6 100%) !important;
  color: #ffffff !important;
  border: none !important;
  box-shadow: 0 8px 22px rgba(37,99,235,0.28) !important;
}
[class*="gradio-container"] button.primary:hover,
[class*="gradio-container"] .gr-button-primary:hover {
  background: linear-gradient(135deg,#1e40af 0%,#2563eb 100%) !important;
  box-shadow: 0 10px 26px rgba(37,99,235,0.34) !important;
}
[class*="gradio-container"] button.secondary,
[class*="gradio-container"] button.stop,
[class*="gradio-container"] .gr-button-secondary,
[class*="gradio-container"] .gr-button-stop {
  background: #ffffff !important;
  color: #1d4ed8 !important;
  border: 1px solid #bfdbfe !important;
  box-shadow: 0 1px 2px rgba(15,39,68,0.04) !important;
}
[class*="gradio-container"] button.secondary:hover,
[class*="gradio-container"] button.stop:hover,
[class*="gradio-container"] .gr-button-secondary:hover,
[class*="gradio-container"] .gr-button-stop:hover {
  background: #eff6ff !important;
  border-color: #93c5fd !important;
}

/* --- KILL DEFAULT GREY BLOCKS / FORMS / PANELS --- */
[class*="gradio-container"] .gr-group,
[class*="gradio-container"] fieldset,
[class*="gradio-container"] .gr-box,
[class*="gradio-container"] .gr-panel,
[class*="gradio-container"] .form,
[class*="gradio-container"] .gr-form,
[class*="gradio-container"] .container {
  background: transparent !important;
}

/* Inside our white section-cards, every nested Gradio wrapper must be FLAT
   (no border / shadow / background), otherwise the auto-form Gradio inserts
   around consecutive Dropdown+Slider components shows up as a "card inside card".
   We override the relevant CSS variables locally + add explicit overrides. */
.section-card {
  --block-shadow: none;
  --block-shadow-dark: none;
  --block-background-fill: transparent;
  --block-border-color: transparent;
  --block-border-width: 0px;
  --panel-background-fill: transparent;
  --panel-border-width: 0px;
  --background-fill-secondary: transparent;
  /* Gradio 5's auto-form uses `background: var(--border-color-primary)`; making
     that transparent inside a section-card eliminates the inner blue rectangle
     around Dropdown + Slider groups. Inputs use --input-border-color instead,
     so they keep their blue border. */
  --border-color-primary: transparent;
  /* Allow dropdown popups / overflowing content to escape the card. */
  overflow: visible !important;
}
/* Catch the auto-form and its immediate wrappers regardless of svelte hash:
   they are direct descendant divs of a section-card that have `display:flex`
   from Gradio defaults. Strip all visual chrome, but keep spacing. */
.section-card > div,
.section-card > div > div,
.section-card > div > div > div {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  overflow: visible !important;
}
.section-card .block,
.section-card .form,
.section-card .gr-form,
.section-card .gr-block,
.section-card .gr-panel,
.section-card .gr-group,
.section-card .gradio-dropdown,
.section-card .gradio-slider,
.section-card .gradio-textbox,
.section-card .gradio-markdown,
.section-card .gradio-code {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  overflow: visible !important;
}

/* Auto-form (consecutive form components) needs to lay its kids out vertically. */
.section-card .form,
.section-card .gr-form {
  display: flex !important;
  flex-direction: column !important;
  gap: 14px !important;
}

/* gr.Row must STAY horizontal even inside a section-card (Run / Stop / Clear). */
[class*="gradio-container"] .section-card .row,
[class*="gradio-container"] .section-card [class*="row"] {
  display: flex !important;
  flex-direction: row !important;
  flex-wrap: wrap !important;
  gap: 10px !important;
}
.action-row {
  display: flex !important;
  flex-direction: row !important;
  gap: 10px !important;
  margin-top: 12px;
}
.action-row > * { flex: 1 1 0; min-width: 0; }

/* Direct children of a card get vertical rhythm without extra chrome. */
.section-card > * + * { margin-top: 14px; }

/* --- INPUTS / TEXTAREA --- */
[class*="gradio-container"] textarea,
[class*="gradio-container"] input:not([type="checkbox"]):not([type="radio"]):not([type="range"]) {
  background: #ffffff !important;
  border: 1px solid #bfdbfe !important;
  box-shadow: none !important;
  border-radius: 14px !important;
  color: #0f2744 !important;
}
[class*="gradio-container"] textarea::placeholder,
[class*="gradio-container"] input::placeholder { color: #7591b3 !important; }
[class*="gradio-container"] textarea:focus,
[class*="gradio-container"] input:focus {
  border-color: #3b82f6 !important;
  box-shadow: 0 0 0 3px rgba(59,130,246,0.2) !important;
  outline: none !important;
}

/* --- DROPDOWN --- */
/* The visible "Model" pill is the [data-testid="dropdown"] wrap; it's the only
   thing that should carry the blue border. The <input> inside it must be
   transparent/borderless or we get "border inside border" nesting. */
[class*="gradio-container"] [data-testid="dropdown"] {
  background: #ffffff !important;
  border: 1px solid #bfdbfe !important;
  border-radius: 14px !important;
  box-shadow: none !important;
  padding: 2px 4px !important;
}
[class*="gradio-container"] [data-testid="dropdown"] > .wrap,
[class*="gradio-container"] [data-testid="dropdown"] .secondary-wrap,
[class*="gradio-container"] [data-testid="dropdown"] .wrap-inner,
[class*="gradio-container"] [data-testid="dropdown"] input,
[class*="gradio-container"] [data-testid="dropdown"] .input-container,
[class*="gradio-container"] [class*="dropdown"] .wrap {
  background: transparent !important;
  border: none !important;
  box-shadow: none !important;
  border-radius: 12px !important;
}
[class*="gradio-container"] .options ul,
[class*="gradio-container"] .options {
  background: #ffffff !important;
  border: 1px solid #bfdbfe !important;
  border-radius: 12px !important;
  box-shadow: 0 12px 30px rgba(37,99,235,0.12) !important;
}
[class*="gradio-container"] .options li[aria-selected="true"],
[class*="gradio-container"] .options li:hover {
  background: #eff6ff !important;
  color: #1d4ed8 !important;
}

/* Small "info / help" labels under inputs */
[class*="gradio-container"] .info,
[class*="gradio-container"] [data-testid*="info"],
[class*="gradio-container"] .gr-info {
  color: #6b86a6 !important;
  background: transparent !important;
  font-size: 12px !important;
}

/* --- SLIDERS (Gradio uses native input[type=range], not noUi) --- */
[class*="gradio-container"] input[type="range"] {
  -webkit-appearance: none;
  appearance: none;
  width: 100%;
  height: 6px;
  background: #dbeafe;
  border-radius: 999px;
  outline: none;
  box-shadow: none !important;
  border: none !important;
}
[class*="gradio-container"] input[type="range"]::-webkit-slider-runnable-track {
  height: 6px;
  background: linear-gradient(90deg,#2563eb var(--val,50%), #dbeafe var(--val,50%));
  border-radius: 999px;
}
[class*="gradio-container"] input[type="range"]::-webkit-slider-thumb {
  -webkit-appearance: none;
  appearance: none;
  width: 18px;
  height: 18px;
  border-radius: 50%;
  background: #ffffff;
  border: 2px solid #2563eb;
  box-shadow: 0 2px 6px rgba(37,99,235,0.25);
  margin-top: -6px;
  cursor: pointer;
}
[class*="gradio-container"] input[type="range"]::-moz-range-track {
  height: 6px;
  background: #dbeafe;
  border-radius: 999px;
}
[class*="gradio-container"] input[type="range"]::-moz-range-progress {
  height: 6px;
  background: #2563eb;
  border-radius: 999px;
}
[class*="gradio-container"] input[type="range"]::-moz-range-thumb {
  width: 16px;
  height: 16px;
  border-radius: 50%;
  background: #ffffff;
  border: 2px solid #2563eb;
  box-shadow: 0 2px 6px rgba(37,99,235,0.25);
}
/* The legacy noUi slider, kept just in case */
[class*="gradio-container"] .noUi-target { background: #dbeafe !important; border: none !important; box-shadow: none !important; }
[class*="gradio-container"] .noUi-connect { background: #2563eb !important; }
[class*="gradio-container"] .noUi-handle  { background: #ffffff !important; border: 2px solid #2563eb !important; box-shadow: 0 2px 8px rgba(37,99,235,0.2) !important; }

/* --- TABS --- */
[class*="gradio-container"] .tabs,
[class*="gradio-container"] .tab-container,
[class*="gradio-container"] .tab-wrapper {
  background: transparent !important;
}
[class*="gradio-container"] .tab-container::after {
  background: rgba(37,99,235,0.18) !important;
}
[class*="gradio-container"] .tab-wrapper button {
  color: #4a6a8c !important;
  font-weight: 600 !important;
}
[class*="gradio-container"] .tab-wrapper button.selected {
  color: #1d4ed8 !important;
}
[class*="gradio-container"] .tab-wrapper button.selected::after {
  background: #2563eb !important;
}

/* --- MARKDOWN / PROSE --- */
[class*="gradio-container"] .gr-markdown,
[class*="gradio-container"] .prose,
[class*="gradio-container"] .markdown { color: #0f2744 !important; }
[class*="gradio-container"] .gr-markdown a,
[class*="gradio-container"] .prose a { color: #1d4ed8 !important; }

/* --- CODE BLOCK (Trace tab) --- */
[class*="gradio-container"] .codemirror-wrapper,
[class*="gradio-container"] .cm-editor,
[class*="gradio-container"] .cm-scroller,
[class*="gradio-container"] .cm-gutters,
[class*="gradio-container"] .cm-content {
  background: #f3f8ff !important;
  color: #0f2744 !important;
  border: none !important;
}
[class*="gradio-container"] .cm-gutters {
  border-right: 1px solid #dbeafe !important;
  color: #6b86a6 !important;
}

/* --- GLOBAL ROUNDED CORNERS: kill any leftover right-angle frames --- */
[class*="gradio-container"] .block,
[class*="gradio-container"] .form,
[class*="gradio-container"] .gr-box,
[class*="gradio-container"] .gr-panel,
[class*="gradio-container"] .gr-group,
[class*="gradio-container"] [data-testid="textbox"],
[class*="gradio-container"] [data-testid="dropdown"],
[class*="gradio-container"] .tabitem,
[class*="gradio-container"] .tab-content,
[class*="gradio-container"] .gradio-markdown,
[class*="gradio-container"] .gradio-code {
  border-radius: 16px !important;
}
[class*="gradio-container"] button {
  border-radius: 12px !important;
}

/* --- "Try Examples" preset buttons --- */
.inline-example-title {
  font-size: 0.72rem;
  font-weight: 700;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: #2563eb;
  margin: 0 0 4px 0;
}
.example-note { color: #4a6a8c; font-size: 12px; margin: 0 0 10px 0; }
.example-buttons { display: grid; gap: 10px; margin-top: 4px; }

/* Each Example button shows "<emoji>  <Category> — <query text>". Left-align the
   label so multi-line examples read like cards, and tint the category prefix in
   blue via a CSS gradient (more robust across Gradio versions than inline HTML). */
[class*="gradio-container"] .example-btn {
  text-align: left !important;
  justify-content: flex-start !important;
  white-space: normal !important;
  line-height: 1.5 !important;
  padding: 14px 16px !important;
  font-size: 14px !important;
  color: #0f2744 !important;
  background: linear-gradient(180deg, #ffffff 0%, #f7fbff 100%) !important;
  border: 1px solid #bfdbfe !important;
  border-radius: 14px !important;
  box-shadow: 0 1px 2px rgba(15,39,68,0.03) !important;
}
[class*="gradio-container"] .example-btn:hover {
  background: #eff6ff !important;
  border-color: #93c5fd !important;
  box-shadow: 0 4px 14px rgba(37,99,235,0.12) !important;
}
[class*="gradio-container"] .example-btn > * {
  color: inherit !important;
  white-space: normal !important;
  display: inline !important;
}

/* --- gr.Examples component (currently unused but defensively styled) --- */
[class*="gradio-container"] [data-testid="block-examples"] {
  background: #f0f6ff !important;
  border: 1px solid #dbeafe !important;
  border-radius: 16px !important;
  padding: 12px !important;
  box-shadow: none !important;
}
[class*="gradio-container"] [data-testid="block-examples"] table,
[class*="gradio-container"] [data-testid="block-examples"] thead,
[class*="gradio-container"] [data-testid="block-examples"] tbody,
[class*="gradio-container"] [data-testid="block-examples"] tr,
[class*="gradio-container"] [data-testid="block-examples"] td {
  border: none !important;
  background: transparent !important;
}
[class*="gradio-container"] [data-testid="block-examples"] button {
  background: #ffffff !important;
  color: #1d4ed8 !important;
  border: 1px solid #bfdbfe !important;
  border-radius: 12px !important;
  box-shadow: none !important;
  font-size: 13px !important;
}
[class*="gradio-container"] [data-testid="block-examples"] button:hover {
  background: #eff6ff !important;
  border-color: #93c5fd !important;
}

/* Hide the small "footer" branding so nothing grey leaks in below the app */
footer { display: none !important; }
"""


@dataclass
class AgentState:
    searched_queries: List[str] = field(default_factory=list)
    visited_urls: List[str] = field(default_factory=list)
    searched_query_set: Set[str] = field(default_factory=set)
    visited_url_set: Set[str] = field(default_factory=set)
    trusted_notes: List[str] = field(default_factory=list)
    trace: List[Dict[str, Any]] = field(default_factory=list)


# Accept a variety of placeholder-only answers: a bare ellipsis (ASCII `...`
# or unicode `…`), a single interpunct, and any whitespace-only content. These
# show up when the model echoes a literal `<answer>...</answer>` template
# from the prompt instead of producing a real answer.
_PLACEHOLDER_ANSWER_RE = re.compile(r"^[\s.\u2026\u00b7]*$")

# Pipe-table separator line, e.g. `| --- | :---: |`. The outer pipes are
# optional in some GFM dialects, so we accept both.
_TABLE_SEPARATOR_RE = re.compile(
    r"^\s*\|?\s*:?-{2,}:?(?:\s*\|\s*:?-{2,}:?)+\s*\|?\s*$"
)


def strip_think_blocks(text: str) -> str:
    """Remove any <think>...</think> reasoning blocks.

    Quest-4B (Qwen3 family) emits `<think>` reasoning before the final
    answer. When the endpoint is deployed without a reasoning parser, the raw
    tags leak into chat completion `content`; stripping them here keeps the
    extracted answer clean for Markdown rendering.
    """
    return re.sub(
        r"<think>.*?</think>", "", text, flags=re.DOTALL | re.IGNORECASE
    )


def decode_escaped_whitespace(text: str) -> str:
    """Decode literal `\\n`/`\\t`/`\\r` sequences back to real whitespace.

    Some OpenAI-compatible servers (and some vLLM builds when a tokenizer's
    chat template escapes control characters) return `choices[0].message.content`
    with newlines stored as the two-character backslash+n sequence rather than
    as a real newline. That breaks Markdown rendering because a pipe table on
    a single line is not a table — it is just a sentence with `|` in it, which
    is exactly the symptom we saw with:

        \\n| Color | Hex |\\n|---|---|\\n| Red | #FF0000 |...

    We only decode when the escapes dominate (at least 3 of them, and at
    least as many as the real newlines in the text). That keeps us from
    corrupting legitimate backslash-n pairs that happen to appear in a code
    sample the model produced.
    """
    if not text:
        return text
    escaped_newlines = text.count("\\n")
    if escaped_newlines == 0 and "\\t" not in text and "\\r" not in text:
        return text
    real_newlines = text.count("\n")
    if escaped_newlines < max(3, real_newlines + 1):
        return text
    # Preserve real backslashes so that `\\\\n` (an actual `\n` the model
    # wrote) doesn't get collapsed to a newline.
    sentinel = "\x00__BS__\x00"
    out = text.replace("\\\\", sentinel)
    out = out.replace("\\n", "\n").replace("\\r", "\r").replace("\\t", "\t")
    out = out.replace(sentinel, "\\")
    return out


def _is_placeholder_answer(text: str) -> bool:
    return bool(_PLACEHOLDER_ANSWER_RE.match(text or ""))


def ensure_markdown_table_blank_lines(text: str) -> str:
    """Insert a blank line before any pipe-table header row.

    GitHub-Flavored Markdown requires a pipe table to be preceded by a
    paragraph break; otherwise the header row is folded into the previous
    paragraph and the whole table renders as raw text. Models sometimes glue
    the table directly under a sentence (e.g. "Here's the comparison: | Col
    ..."), so we fix that up defensively.
    """
    lines = text.split("\n")
    out: List[str] = []
    for idx, line in enumerate(lines):
        is_header = (
            "|" in line
            and idx + 1 < len(lines)
            and _TABLE_SEPARATOR_RE.match(lines[idx + 1]) is not None
        )
        if is_header and out and out[-1].strip() != "":
            out.append("")
        out.append(line)
    return "\n".join(out)


def extract_answer(text: str) -> Optional[str]:
    """Return the content of the first `<answer>...</answer>` block.

    Tries two strategies, in order, and discards placeholder-only content
    (bare ellipses) that the model sometimes echoes from the prompt:

    1. Well-formed `<answer>...</answer>` block.
    2. Truncated `<answer>...` with no closing tag (tokens ran out);
       in that case we take everything after the opening tag.
    """
    # Decode escaped whitespace on the whole output first so the <answer>
    # regex can actually match the opening and closing tags across lines.
    decoded = decode_escaped_whitespace(text or "")
    cleaned = strip_think_blocks(decoded)

    full_match = re.search(
        r"<answer>\s*(.*?)\s*</answer>",
        cleaned,
        flags=re.DOTALL | re.IGNORECASE,
    )
    if full_match is not None:
        candidate = decode_escaped_whitespace(full_match.group(1).strip())
        if candidate and not _is_placeholder_answer(candidate):
            return candidate
        # Closed block was a placeholder / empty: fail fast. Do NOT fall
        # through to the open-ended strategy, or it would re-match the same
        # tag and incorrectly capture `...</answer>` as the answer.
        return None

    open_match = re.search(
        r"<answer>\s*(.*)$", cleaned, flags=re.DOTALL | re.IGNORECASE
    )
    if open_match is not None:
        candidate = decode_escaped_whitespace(open_match.group(1).strip())
        if candidate and not _is_placeholder_answer(candidate):
            return candidate

    return None


def parse_tool_call(text: str) -> Tuple[Optional[str], Optional[Dict[str, Any]], Optional[str]]:
    cleaned = strip_think_blocks(text or "")
    match = re.search(r"<tool_call>\s*(.*?)\s*</tool_call>", cleaned, flags=re.DOTALL | re.IGNORECASE)
    if not match:
        return None, None, None
    payload = match.group(1).strip()
    try:
        data = json.loads(payload)
    except json.JSONDecodeError:
        return None, None, "Invalid JSON in <tool_call> block."

    name = data.get("name")
    arguments = data.get("arguments", {})
    if not isinstance(name, str) or not isinstance(arguments, dict):
        return None, None, "Invalid tool format. Expect name(str) and arguments(dict)."
    return name, arguments, None


def _run_search_single(query: str, max_results: int) -> Dict[str, Any]:
    if not query.strip():
        return {"ok": False, "error": "Search query cannot be empty."}
    cache_key = f"{query.strip().lower()}::{max_results}"
    if cache_key in SEARCH_CACHE:
        return {**SEARCH_CACHE[cache_key], "cached": True}

    rows: List[Dict[str, str]] = []
    with DDGS() as ddgs:
        for item in ddgs.text(query, max_results=max_results):
            rows.append(
                {
                    "title": item.get("title", ""),
                    "href": item.get("href", ""),
                    "body": item.get("body", ""),
                }
            )
    payload = {"ok": True, "query": query, "results": rows, "cached": False}
    SEARCH_CACHE[cache_key] = payload
    return payload


def run_search(query: Union[str, List[str]], max_results: int = 5) -> Dict[str, Any]:
    """Runs one or more queries through DuckDuckGo.

    QUEST's schema passes `query` as an array of strings, while the simpler
    starter schema used a single string. We accept both shapes.
    """
    if isinstance(query, list):
        sub_results: List[Dict[str, Any]] = []
        for q in query:
            if not isinstance(q, str) or not q.strip():
                continue
            sub_results.append(_run_search_single(q, max_results))
        return {"ok": True, "queries": query, "results": sub_results}
    return _run_search_single(str(query or "").strip(), max_results)


def _clean_html_to_text(html: str, max_chars: int) -> str:
    soup = BeautifulSoup(html, "html.parser")
    for tag in soup(["script", "style", "noscript"]):
        tag.decompose()
    text = soup.get_text(separator=" ", strip=True)
    text = re.sub(r"\s+", " ", text)
    return text[:max_chars]


def _run_visit_single(url: str, max_chars: int, goal: str = "") -> Dict[str, Any]:
    if not url.strip():
        return {"ok": False, "error": "URL cannot be empty."}
    cache_key = f"{url.strip()}::{max_chars}"
    if cache_key in VISIT_CACHE:
        return {**VISIT_CACHE[cache_key], "cached": True, "goal": goal}
    try:
        resp = requests.get(
            url,
            timeout=20,
            headers={"User-Agent": "Mozilla/5.0 (compatible; DeepResearchSpace/1.0)"},
        )
        resp.raise_for_status()
        content_type = resp.headers.get("content-type", "")
        if "text/html" in content_type or "<html" in resp.text[:200].lower():
            text = _clean_html_to_text(resp.text, max_chars=max_chars)
        else:
            text = resp.text[:max_chars]
        payload = {"ok": True, "url": url, "content": text, "cached": False, "goal": goal}
        VISIT_CACHE[cache_key] = payload
        return payload
    except Exception as exc:
        return {"ok": False, "url": url, "error": str(exc), "goal": goal}


def run_visit(
    url: Union[str, List[str]],
    max_chars: int = 6000,
    goal: str = "",
) -> Dict[str, Any]:
    """Fetches one or more URLs. Accepts string or list (QUEST schema)."""
    if isinstance(url, list):
        sub_results: List[Dict[str, Any]] = []
        for u in url:
            if not isinstance(u, str) or not u.strip():
                continue
            sub_results.append(_run_visit_single(u, max_chars, goal))
        return {"ok": True, "goal": goal, "results": sub_results}
    return _run_visit_single(str(url or "").strip(), max_chars, goal)


def _build_client_for_model(model: str) -> Tuple[InferenceClient, str, List[str]]:
    """Returns (client, primary_model_id, fallback_model_ids).

    When the user picks the Quest model and QUEST_BASE_URL is configured, the
    InferenceClient is pointed at the dedicated endpoint; otherwise we hit the
    shared HF Inference API and let the starter fall back across free models.
    """
    token = os.getenv("HF_TOKEN")
    if model == QUEST_MODEL_ID and QUEST_BASE_URL:
        client = InferenceClient(
            base_url=QUEST_BASE_URL,
            token=token,
            timeout=120,
        )
        return client, QUEST_ENDPOINT_MODEL, []
    client = InferenceClient(token=token, timeout=60)
    fallbacks = [m for m in FREE_FALLBACK_MODELS if m != model]
    return client, model, fallbacks


def call_model(
    client: InferenceClient,
    messages: List[Dict[str, str]],
    preferred_model: str,
    candidate_models: List[str],
    temperature: float,
    max_new_tokens: int,
) -> Tuple[str, str]:
    model_order: List[str] = []
    for m in [preferred_model] + candidate_models:
        if m and m not in model_order:
            model_order.append(m)

    last_error = None
    for model_name in model_order:
        try:
            completion = client.chat_completion(
                model=model_name,
                messages=messages,
                temperature=temperature,
                max_tokens=max_new_tokens,
            )
            return completion.choices[0].message.content or "", model_name
        except Exception as exc:
            last_error = exc
            continue
    raise RuntimeError(f"All model candidates failed. Last error: {last_error}")


def build_research_agent(
    question: str,
    model: str,
    max_turns: int,
    max_search_results: int,
    temperature: float,
) -> Tuple[str, str]:
    client, primary_model, fallback_models = _build_client_for_model(model)
    # Display label: the real HF repo id is nicer than the TGI shim name.
    display_primary = model if (model == QUEST_MODEL_ID) else primary_model
    state = AgentState()
    used_model = display_primary

    messages: List[Dict[str, str]] = [
        {"role": "system", "content": build_system_prompt()},
        {"role": "user", "content": question},
    ]

    final_answer: Optional[str] = None

    for turn in range(1, max_turns + 1):
        if state.trusted_notes and turn > 1 and turn % 3 == 0:
            summary_lines = "\n".join(f"- {n}" for n in state.trusted_notes[-6:])
            messages.append(
                {
                    "role": "user",
                    "content": f"RESEARCH STATE SUMMARY\n{summary_lines}\nUse this summary to avoid repeating work.",
                }
            )

        raw_output, endpoint_model = call_model(
            client=client,
            messages=messages,
            preferred_model=primary_model,
            candidate_models=fallback_models,
            temperature=temperature,
            max_new_tokens=int(os.getenv("QUEST_MAX_NEW_TOKENS", "4096")),
        )
        model_output = raw_output
        # Preserve the human-friendly model id for the trace even if the
        # endpoint ignores the "model" param and returns the TGI shim name.
        used_model = display_primary if endpoint_model == primary_model == QUEST_ENDPOINT_MODEL else endpoint_model
        messages.append({"role": "assistant", "content": model_output})
        state.trace.append({"turn": turn, "assistant": model_output})

        extracted_answer = extract_answer(model_output)
        if extracted_answer:
            final_answer = extracted_answer
            break

        tool_name, tool_args, tool_err = parse_tool_call(model_output)
        if tool_err:
            tool_response = {"ok": False, "error": tool_err}
        elif not tool_name:
            # No explicit tool call and no final answer: force finalization.
            # IMPORTANT: do not write the literal characters `<answer>...</answer>`
            # here. Some models (notably the Qwen3 family that Quest-4B is
            # built on) will echo the template verbatim, which means the
            # extracted answer ends up being the three-dot placeholder `...`
            # and the user sees an empty-looking result.
            messages.append(
                {
                    "role": "user",
                    "content": (
                        "You did not call a tool and did not produce a final "
                        "answer. Please now write your best final answer, "
                        "wrapped between an opening <answer> tag and a "
                        "closing </answer> tag. Put the real answer text "
                        "between those tags; do not write a literal ellipsis "
                        "or other placeholder. If the question asks for "
                        "tabular data, use GitHub-Flavored Markdown pipe "
                        "tables (`| col1 | col2 |` + `|---|---|`) and put a "
                        "blank line before the first row so the table renders."
                    ),
                }
            )
            continue
        else:
            if tool_name == "search":
                raw_query = tool_args.get("query", "")
                queries: List[str]
                if isinstance(raw_query, list):
                    queries = [str(q).strip() for q in raw_query if str(q).strip()]
                else:
                    queries = [str(raw_query).strip()] if str(raw_query).strip() else []
                max_results = int(tool_args.get("max_results", max_search_results))
                max_results = max(1, min(max_results, 10))

                per_query: List[Dict[str, Any]] = []
                for q in queries:
                    if q in state.searched_query_set:
                        per_query.append({
                            "ok": True,
                            "query": q,
                            "cached": True,
                            "note": "Already searched; reusing cached result.",
                            "results": [],
                        })
                        continue
                    state.searched_queries.append(q)
                    state.searched_query_set.add(q)
                    single = _run_search_single(q, max_results)
                    per_query.append(single)
                    if single.get("ok"):
                        first_titles = [r.get("title", "") for r in single.get("results", [])[:2]]
                        if first_titles:
                            state.trusted_notes.append(
                                f"Searched '{q}' and found leads: {', '.join(t for t in first_titles if t)}"
                            )
                tool_response = (
                    per_query[0]
                    if len(per_query) == 1
                    else {"ok": True, "queries": queries, "results": per_query}
                )
            elif tool_name == "visit":
                raw_url = tool_args.get("url", "")
                urls: List[str]
                if isinstance(raw_url, list):
                    urls = [str(u).strip() for u in raw_url if str(u).strip()]
                else:
                    urls = [str(raw_url).strip()] if str(raw_url).strip() else []
                goal = str(tool_args.get("goal", "")).strip()
                max_chars = int(tool_args.get("max_chars", 6000))
                max_chars = max(500, min(max_chars, 20000))

                per_url: List[Dict[str, Any]] = []
                for u in urls:
                    if u in state.visited_url_set:
                        per_url.append({
                            "ok": True,
                            "url": u,
                            "cached": True,
                            "note": "Already visited; reusing cached result.",
                        })
                        continue
                    state.visited_urls.append(u)
                    state.visited_url_set.add(u)
                    single = _run_visit_single(u, max_chars, goal)
                    per_url.append(single)
                    if single.get("ok"):
                        snippet = str(single.get("content", ""))[:180]
                        if snippet:
                            state.trusted_notes.append(
                                f"Visited {u} and extracted key context: {snippet}"
                            )
                tool_response = (
                    per_url[0]
                    if len(per_url) == 1
                    else {"ok": True, "goal": goal, "results": per_url}
                )
            else:
                tool_response = {"ok": False, "error": f"Unknown tool: {tool_name}"}

        state.trace.append({"turn": turn, "tool": tool_name, "tool_response": tool_response})
        messages.append(
            {
                "role": "user",
                "content": TOOL_RESPONSE_TEMPLATE.format(
                    payload=json.dumps(tool_response, ensure_ascii=False)
                ),
            }
        )

    if final_answer is None:
        final_answer = (
            "I could not finish a complete research answer within the configured turns. "
            "Try increasing max turns or switching to a stronger model."
        )
    else:
        final_answer = ensure_markdown_table_blank_lines(final_answer)

    citations = "\n".join(f"- {url}" for url in sorted(set(state.visited_urls)))
    final_answer = f"**Model used:** `{used_model}`\n\n{final_answer}"
    if citations:
        final_answer = f"{final_answer}\n\n### Visited Sources\n{citations}"

    trace_text = json.dumps(
        {
            "used_model": used_model,
            "searched_queries": state.searched_queries,
            "visited_urls": state.visited_urls,
            "trusted_notes": state.trusted_notes[-10:],
            "trace": state.trace,
        },
        ensure_ascii=False,
        indent=2,
    )
    return final_answer, trace_text


def run_ui(
    question: str,
    model: str,
    max_turns: int,
    max_search_results: int,
    temperature: float,
):
    if not question.strip():
        return "Please input a question.", "{}"
    if not os.getenv("HF_TOKEN"):
        warning = (
            "HF_TOKEN is not configured in Space Secrets. "
            "Go to Settings -> Secrets -> add `HF_TOKEN`, then retry."
        )
        return warning, json.dumps({"error": warning}, ensure_ascii=False, indent=2)
    if model == QUEST_MODEL_ID and not QUEST_BASE_URL:
        warning = (
            f"`{QUEST_MODEL_ID}` is private and not available via the free HF Inference API. "
            "Create a dedicated HF Inference Endpoint for it (https://ui.endpoints.huggingface.co/), "
            "then set `QUEST_BASE_URL` in Space Secrets to the endpoint's `/v1/` URL. "
            "In the meantime you can pick one of the open-weights models in the dropdown."
        )
        return warning, json.dumps({"error": warning}, ensure_ascii=False, indent=2)
    try:
        return build_research_agent(
            question=question,
            model=model,
            max_turns=max_turns,
            max_search_results=max_search_results,
            temperature=temperature,
        )
    except Exception as exc:
        return f"Error: {exc}", json.dumps({"error": str(exc)}, ensure_ascii=False, indent=2)


EXAMPLES = [
    {
        "category": "Fixed facts",
        "icon": "🎯",
        "text": "Who wrote the novel 1984, and when was it first published?",
    },
    {
        "category": "Time-varying",
        "icon": "📈",
        "text": "Who is the current CEO of Tesla, and what is the company's latest stock price?",
    },
    {
        "category": "Multi-constraints",
        "icon": "🧩",
        "text": "Find a 2-day Tokyo itinerary under $250 focused on museums and vegetarian food.",
    },
    {
        "category": "Long-form research report",
        "icon": "📚",
        "text": "Write a short guide comparing electric cars vs hybrid cars for a daily commuter, covering cost, range, and maintenance.",
    },
]


def _example_label(ex: Dict[str, str]) -> str:
    return f"{ex['icon']}  {ex['category']} — {ex['text']}"


with gr.Blocks(
    title="Quest · Deep Research by OSU NLP",
    theme=APP_THEME,
    css=CUSTOM_CSS,
    fill_width=True,
) as demo:
    # --- Top banner with the Quest logo centered and the OSU NLP mark on the right ---
    with gr.Row(elem_classes="top-banner"):
        with gr.Column(scale=1, elem_classes="banner-side"):
            pass
        with gr.Column(scale=4, elem_classes="banner-center"):
            gr.Image(
                value=LOGO_PATH,
                show_label=False,
                container=False,
                interactive=False,
                show_download_button=False,
                show_fullscreen_button=False,
                show_share_button=False,
                elem_classes="banner-quest-logo",
            )
        with gr.Column(scale=1, elem_classes="banner-side banner-right"):
            gr.Image(
                value=OSU_NLP_LOGO_PATH,
                show_label=False,
                container=False,
                interactive=False,
                show_download_button=False,
                show_fullscreen_button=False,
                show_share_button=False,
                elem_classes="osu-nlp-logo",
            )

    # --- Main two-column layout ---
    with gr.Row(elem_classes="layout-gap"):
        with gr.Column(scale=6, min_width=420):
            with gr.Group(elem_classes="section-card"):
                gr.HTML('<div class="section-heading hero-heading">What can I search for you?</div>')
                question = gr.Textbox(
                    show_label=False,
                    placeholder="Ask anything you want to research in depth...",
                    lines=6,
                )
                with gr.Row(elem_classes="action-row"):
                    run_btn = gr.Button("Run Research", variant="primary", size="lg")
                    stop_btn = gr.Button("Stop", variant="stop", size="lg")
                    clear_btn = gr.Button("Clear", variant="secondary", size="lg")

            with gr.Group(elem_classes="section-card"):
                gr.HTML('<div class="section-heading">Try Examples</div>')
                gr.HTML(
                    '<div class="example-note">Each example shows the kind of query it represents. Click one to auto-fill.</div>'
                )
                with gr.Column(elem_classes="example-buttons"):
                    example_buttons = [
                        gr.Button(_example_label(ex), variant="secondary", elem_classes="example-btn")
                        for ex in EXAMPLES
                    ]

            with gr.Group(elem_classes="section-card"):
                with gr.Tabs():
                    with gr.TabItem("Result"):
                        answer = gr.Markdown(label="Final Answer")
                    with gr.TabItem("Record"):
                        trace = gr.Code(label="Execution Trace (JSON)", language="json")

        with gr.Column(scale=4, min_width=340, elem_classes="right-stack"):
            with gr.Group(elem_classes=["section-card", "no-frame"]):
                gr.HTML(
                    f"""
<div class="icon-grid">
  <a class="icon-link" href="{PAPER_URL}" target="_blank" rel="noopener noreferrer">📄 Paper</a>
  <a class="icon-link" href="{CODE_URL}" target="_blank" rel="noopener noreferrer">💻 Code</a>
  <a class="icon-link" href="{DATASET_URL}" target="_blank" rel="noopener noreferrer">🗂️ Dataset</a>
  <a class="icon-link" href="{MODEL_URL}" target="_blank" rel="noopener noreferrer">🧠 Model</a>
</div>
"""
                )

            with gr.Group(elem_classes="section-card"):
                gr.HTML('<div class="section-heading">Settings</div>')
                model = gr.Dropdown(
                    label="Model",
                    choices=DEFAULT_MODEL_CHOICES,
                    value=DEFAULT_MODEL if DEFAULT_MODEL in DEFAULT_MODEL_CHOICES else DEFAULT_MODEL_CHOICES[0],
                    allow_custom_value=True,
                )
                max_turns = gr.Slider(
                    label="Max Turns",
                    minimum=2,
                    maximum=20,
                    value=8,
                    step=1,
                )
                max_search_results = gr.Slider(
                    label="Search Results Per Query",
                    minimum=1,
                    maximum=10,
                    value=5,
                    step=1,
                )
                temperature = gr.Slider(
                    label="Temperature",
                    minimum=0.0,
                    maximum=1.5,
                    value=0.4,
                    step=0.1,
                )

    run_event = run_btn.click(
        fn=run_ui,
        inputs=[question, model, max_turns, max_search_results, temperature],
        outputs=[answer, trace],
    )
    for btn, ex in zip(example_buttons, EXAMPLES):
        btn.click(
            fn=(lambda text=ex["text"]: text),
            inputs=[],
            outputs=[question],
        )
    stop_btn.click(fn=None, cancels=[run_event])
    clear_btn.click(
        fn=lambda: ("", "", "{}"),
        inputs=[],
        outputs=[question, answer, trace],
    )


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