app.py

"""Gradio UI for the kNN spot-check viewer.

Type a hostname, get its nearest neighbours in the learned 128-d embedding space. Companion to the
2D viewer (``embedding-2d-viewer/``): that one shows *where* hosts land after a UMAP projection,
this one answers *who is actually nearest* in the real space.

The point of the tool is **walking** the space, not one lookup: click any neighbour to make it the
new query, with a back button and a breadcrumb trail. That turns "are these neighbours sensible?"
into "does the neighbourhood stay coherent two or three hops out?", which is the question that
actually catches a bad embedding.

Runs two ways, unchanged:

* **locally** — ``WGL_KNN_INDEX=data/knn/cc_deg8_poc python knn-viewer/app.py``;
* **as a HF Space** — the index arrives on a mounted volume at ``/data``, which is the default.

The index loads once at module scope, not per request: it is the whole cost of a query, and a
Space re-imports this module only on cold start.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

import gradio as gr

# Allow running straight from a repo checkout without installing the package.
_REPO_SRC = Path(__file__).resolve().parent.parent / "src"
if _REPO_SRC.is_dir() and str(_REPO_SRC) not in sys.path:
    sys.path.insert(0, str(_REPO_SRC))

from wgl.knn.index import KnnIndex  # noqa: E402
from wgl.release.identity import domain_key as to_domain_key  # noqa: E402

INDEX_PATH = Path(os.environ.get("WGL_KNN_INDEX", "/data"))
EXAMPLES = ["www.bundestag.de", "www.bild.de", "https://www.bahn.de/service", "portal.gov.cz"]
HEADERS = ["#", "host", "cosine", "registrable domain", "same site", "open"]
DATATYPES = ["number", "str", "number", "str", "str", "markdown"]

INDEX = KnnIndex.load(INDEX_PATH)
_META = INDEX.meta
_FOOTER = (
    f"**{len(INDEX):,} hosts** · `{_META.get('release_id', 'unknown release')}` · "
    f"model `{_META.get('model_space_version', 'unknown')}` · "
    f"index `{_META.get('index_factory', '?')}` · dim {_META.get('dim', '?')} · "
    f"built {str(_META.get('built_at', '?'))[:19]}"
)
_SUBSET_NOTE = ""
if _META.get("source_row_groups_per_shard") is not None:
    _SUBSET_NOTE = (
        f"\n\n⚠️ **This is a proof-of-concept subset**, not the full corpus: source rows "
        f"{_META.get('source_row_min'):,} to {_META.get('source_row_max'):,} of 52,913,544. "
        "Hosts outside that slice will read as *not in the corpus*."
    )


# ── Common Crawl look ─────────────────────────────────────────────────────────────────────────
# Same palette as the release figures (``src/wgl/viz/svg.py``) and the 2D viewer: brand blue
# #0050bd, navy ink #152a47, cool neutral greys, Libre Franklin + IBM Plex Mono. Gradio derives a
# lot of shades from its hue ramps, so both ramps are spelled out and the tokens that matter are
# then pinned explicitly; every colour has a ``_dark`` twin because the Space follows the browser.
_CC_BLUE = gr.themes.Color(
    c50="#e8f0fb",
    c100="#cfe0f6",
    c200="#9dc0ee",
    c300="#6ba1e5",
    c400="#4a90e8",
    c500="#0050bd",
    c600="#0047a8",
    c700="#003d92",
    c800="#00306f",
    c900="#152a47",
    c950="#0d1929",
)
_CC_SLATE = gr.themes.Color(
    c50="#f6f7fa",
    c100="#eef1f6",
    c200="#e8ecf2",
    c300="#dde2e8",
    c400="#8d99b0",
    c500="#6e7191",
    c600="#44536f",
    c700="#2a456d",
    c800="#1d3557",
    c900="#152a47",
    c950="#0d1929",
)
THEME = gr.themes.Base(
    primary_hue=_CC_BLUE,
    secondary_hue=_CC_BLUE,
    neutral_hue=_CC_SLATE,
    font=(
        gr.themes.GoogleFont("Libre Franklin"),
        "system-ui",
        "-apple-system",
        "Segoe UI",
        "sans-serif",
    ),
    font_mono=(
        gr.themes.GoogleFont("IBM Plex Mono"),
        "ui-monospace",
        "SFMono-Regular",
        "Menlo",
        "monospace",
    ),
).set(
    # page + panels: surface over a recessed plane, inverted in dark
    body_background_fill="#f6f7fa",
    body_background_fill_dark="#0d1929",
    body_text_color="#152a47",
    body_text_color_dark="#ffffff",
    body_text_color_subdued="#6e7191",
    body_text_color_subdued_dark="#8d99b0",
    background_fill_primary="#ffffff",
    background_fill_primary_dark="#152a47",
    background_fill_secondary="#f6f7fa",
    background_fill_secondary_dark="#0d1929",
    block_background_fill="#ffffff",
    block_background_fill_dark="#152a47",
    block_border_color="#dde2e8",
    block_border_color_dark="#2a456d",
    block_label_text_color="#44536f",
    block_label_text_color_dark="#c3ccdb",
    block_title_text_color="#152a47",
    block_title_text_color_dark="#ffffff",
    border_color_primary="#dde2e8",
    border_color_primary_dark="#2a456d",
    border_color_accent="#0050bd",
    border_color_accent_dark="#4a90e8",
    code_background_fill="#f6f7fa",
    code_background_fill_dark="#0d1929",
    # accents: brand blue in light, the lighter dark-theme blue on navy
    color_accent="#0050bd",
    link_text_color="#0050bd",
    link_text_color_dark="#4a90e8",
    link_text_color_hover="#0047a8",
    link_text_color_hover_dark="#2895f7",
    link_text_color_active="#0047a8",
    link_text_color_active_dark="#2895f7",
    link_text_color_visited="#6631ff",
    link_text_color_visited_dark="#9353ff",
    loader_color="#0050bd",
    loader_color_dark="#4a90e8",
    slider_color="#0050bd",
    slider_color_dark="#4a90e8",
    # inputs
    input_background_fill="#ffffff",
    input_background_fill_dark="#0d1929",
    input_border_color="#dde2e8",
    input_border_color_dark="#2a456d",
    input_border_color_focus="#0050bd",
    input_border_color_focus_dark="#4a90e8",
    input_placeholder_color="#6e7191",
    input_placeholder_color_dark="#8d99b0",
    # buttons
    button_primary_background_fill="#0050bd",
    button_primary_background_fill_hover="#0047a8",
    button_primary_background_fill_dark="#0050bd",
    button_primary_background_fill_hover_dark="#4a90e8",
    button_primary_border_color="#0050bd",
    button_primary_border_color_dark="#0050bd",
    button_primary_text_color="#ffffff",
    button_primary_text_color_dark="#ffffff",
    button_secondary_background_fill="#f6f7fa",
    button_secondary_background_fill_hover="#e8ecf2",
    button_secondary_background_fill_dark="#152a47",
    button_secondary_background_fill_hover_dark="#1d3557",
    button_secondary_border_color="#dde2e8",
    button_secondary_border_color_dark="#2a456d",
    button_secondary_text_color="#152a47",
    button_secondary_text_color_dark="#ffffff",
    # the neighbour table, the one dense surface in the app
    table_border_color="#dde2e8",
    table_border_color_dark="#2a456d",
    table_even_background_fill="#f6f7fa",
    table_even_background_fill_dark="#0d1929",
    table_odd_background_fill="#ffffff",
    table_odd_background_fill_dark="#152a47",
    table_row_focus="#e8ecf2",
    table_row_focus_dark="#1d3557",
    table_text_color="#152a47",
    table_text_color_dark="#ffffff",
    # status colours, shared by both themes
    error_text_color="#d03b3b",
    error_text_color_dark="#d03b3b",
)

# The few things the theme tokens do not reach: the brand rule under the title, the mono trail and
# the small print in the footer. Dark values ride on Gradio's own ``.dark`` class on the wrapper.
CSS = """
.gradio-container { max-width: 1180px !important; }
#title h1 { margin-bottom: .35rem; padding-bottom: .35rem; letter-spacing: -.01em;
  border-bottom: 3px solid #0050bd; }
.dark #title h1 { border-bottom-color: #4a90e8; }
#crumb, #footer { font-size: .82rem; color: #6e7191; }
.dark #crumb, .dark #footer { color: #8d99b0; }
#crumb code, #footer code { background: transparent; padding: 0; color: #44536f; }
.dark #crumb code, .dark #footer code { color: #c3ccdb; }
"""


def _table(result) -> list[list[object]]:  # noqa: ANN001 - QueryResult, kept loose for gradio
    """Render a query result as dataframe rows, with the free per-neighbour metadata."""
    query_domain = to_domain_key(result.query, reversed_host=False)
    rows: list[list[object]] = []
    for neighbor in result.neighbors:
        domain = to_domain_key(neighbor.host, reversed_host=False)
        rows.append(
            [
                neighbor.rank,
                neighbor.host,
                round(neighbor.score, 4),
                domain,
                "✓" if domain and domain == query_domain else "",
                f"[↗](https://{neighbor.host})",
            ]
        )
    return rows


def _breadcrumb(trail: list[str]) -> str:
    """Render the navigation trail."""
    if len(trail) < 2:
        return ""
    shown = trail[-6:]
    prefix = "… → " if len(trail) > len(shown) else ""
    return "**Trail:** " + prefix + " → ".join(f"`{h}`" for h in shown)


def search(host: str, k: int, trail: list[str]) -> tuple:
    """Run a lookup and return every output the UI needs.

    Args:
        host: Raw user input — a bare host, a domain, or a pasted URL.
        k: Number of neighbours to return.
        trail: The navigation history so far.

    Returns:
        ``(status, rows, neighbour_hosts, trail, breadcrumb, host_box)`` — the extra state lets a
        row click resolve back to a hostname, and keeps the textbox in sync while walking.
    """
    trail = list(trail or [])
    if not host or not host.strip():
        return "Type a hostname to begin.", [], [], trail, _breadcrumb(trail), host

    result = INDEX.query(host, k=int(k))
    if not result.found:
        return (
            f"❌ **`{result.query or host.strip()}` is not in the indexed corpus.**\n\n"
            "That is a normal answer, not an error — the index covers a degree-filtered "
            "subset of the Common Crawl host graph. Tried the exact host, its `www.` variant "
            "and its registrable domain.",
            [],
            [],
            trail,
            _breadcrumb(trail),
            host,
        )

    if not trail or trail[-1] != result.query:
        trail.append(result.query)
    via = {
        "exact": "exact match",
        "www": "matched the `www.` variant",
        "no-www": "matched without the `www.` prefix",
        "domain": "fell back to the registrable domain",
    }.get(result.resolved_via, result.resolved_via)
    status = (
        f"✅ **`{result.query}`** — {via}, row {result.row:,}. "
        "*Click any row to explore that host's neighbours.*"
    )
    hosts = [n.host for n in result.neighbors]
    return status, _table(result), hosts, trail, _breadcrumb(trail), result.query


def follow(evt: gr.SelectData, hosts: list[str], k: int, trail: list[str]) -> tuple:
    """Make the clicked neighbour the new query."""
    row = evt.index[0] if isinstance(evt.index, (list, tuple)) else evt.index
    if not hosts or row is None or row >= len(hosts):
        return gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip()
    return search(hosts[row], k, trail)


def go_back(k: int, trail: list[str]) -> tuple:
    """Step one hop back along the trail."""
    trail = list(trail or [])
    if len(trail) < 2:
        return gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip(), gr.skip()
    trail.pop()
    previous = trail.pop()
    return search(previous, k, trail)


with gr.Blocks(title="Web graph kNN spot-check", analytics_enabled=False) as demo:
    neighbor_state = gr.State([])
    trail_state = gr.State([])

    gr.Markdown("# Web graph kNN spot-check", elem_id="title")
    gr.Markdown(
        "Nearest neighbours of a host in the learned link-prediction embedding space "
        "(cosine similarity on L2-normalized vectors). **Click a neighbour to walk to it.**"
        + _SUBSET_NOTE
    )
    with gr.Row():
        host_box = gr.Textbox(
            label="Host or URL", placeholder="www.example.com", scale=4, autofocus=True
        )
        k_slider = gr.Slider(minimum=5, maximum=100, value=20, step=5, label="Neighbours", scale=1)
    with gr.Row():
        go = gr.Button("Search", variant="primary", scale=3)
        back = gr.Button("← Back", scale=1)
    status = gr.Markdown()
    crumb = gr.Markdown(elem_id="crumb")
    table = gr.Dataframe(
        headers=HEADERS,
        datatype=DATATYPES,
        label="Nearest neighbours — click a row to explore it",
        wrap=True,
        interactive=False,
    )
    gr.Examples(examples=[[e] for e in EXAMPLES], inputs=[host_box])
    gr.Markdown(_FOOTER, elem_id="footer")

    OUTPUTS = [status, table, neighbor_state, trail_state, crumb, host_box]
    go.click(search, [host_box, k_slider, trail_state], OUTPUTS)
    host_box.submit(search, [host_box, k_slider, trail_state], OUTPUTS)
    k_slider.release(search, [host_box, k_slider, trail_state], OUTPUTS)
    table.select(follow, [neighbor_state, k_slider, trail_state], OUTPUTS)
    back.click(go_back, [k_slider, trail_state], OUTPUTS)


if __name__ == "__main__":
    # Gradio 6 takes ``theme``/``css`` on launch(), not on the Blocks constructor (it warns and
    # drops them there). The Space runs this file as its app_file, so this is the one entry point.
    demo.launch(theme=THEME, css=CSS)