build_layout.py

"""Precompute the 2D UMAP layout of the tool embeddings for the background constellation.

Run once locally (needs the model + HF auth + umap-learn):

    python build_layout.py

Embeds every tool's routing_text with LFM2.5-Embedding-350M (the same text and prompt
the live search uses), projects to 2D with UMAP (cosine metric), normalises to [0,1]
with a small inset, and writes static/layout.json — a list of
{key: "domain|name", name, domain, x, y}. The frontend renders one dot per tool and
lights up the ones a query retrieves. The Space itself does not need umap-learn.
"""

from __future__ import annotations

import json
import pathlib

import numpy as np

import search as S
import toolset as T

OUT = pathlib.Path(__file__).resolve().parent / "static" / "layout.json"


def _normalize(xy: np.ndarray) -> np.ndarray:
    lo, hi = xy.min(axis=0), xy.max(axis=0)
    span = np.where(hi - lo == 0, 1.0, hi - lo)
    return (xy - lo) / span                         # -> [0, 1]; the frontend adds border padding


def main() -> None:
    import umap  # heavy (numba); only needed for this offline build

    tools = T.load_catalog()
    print(f"embedding {len(tools)} tools with {S.MODEL_ID} ...", flush=True)
    emb = S._encode([t.routing_text for t in tools], prompt_name="document")

    n_neighbors = min(15, max(2, len(tools) - 1))
    reducer = umap.UMAP(
        n_components=2, metric="cosine", n_neighbors=n_neighbors,
        min_dist=0.12, random_state=42,
    )
    xy = _normalize(np.asarray(reducer.fit_transform(emb), dtype=np.float64))

    pts = [
        {"key": f"{t.domain}|{t.name}", "name": t.name, "domain": t.domain,
         "x": round(float(xy[i, 0]), 4), "y": round(float(xy[i, 1]), 4)}
        for i, t in enumerate(tools)
    ]
    OUT.write_text(json.dumps(pts), encoding="utf-8")
    print(f"wrote {len(pts)} points to {OUT}", flush=True)


if __name__ == "__main__":
    main()