server.py

"""FastAPI backend for the embedding tool selector (Docker Space, no Gradio).

Builds the tool index once at startup, then serves:
  POST /api/search  -> {query, ms, total, results:[{rank, name, description, domain, score, params}]}
  GET  /api/health  -> readiness (tool + domain counts)
  GET  /            -> static/index.html

The model repo may be private/gated, so HF_TOKEN (a Space secret) is needed to pull it.

    uvicorn server:app --host 0.0.0.0 --port 7860
"""

import os
import time
import traceback

from fastapi import FastAPI
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

import search as S

STATIC = os.path.join(os.path.dirname(os.path.abspath(__file__)), "static")

INDEX, INDEX_ERROR = None, None
try:
    print(f"[server] building tool index with {S.MODEL_ID} ...", flush=True)
    INDEX = S.load_or_build_index()
    print(f"[server] ready: {len(INDEX['tools'])} tools across {len(INDEX['domains'])} domains", flush=True)
except Exception:
    INDEX_ERROR = traceback.format_exc()
    print("INDEX BUILD FAILED:\n" + INDEX_ERROR, flush=True)

app = FastAPI(title="Embedding tool selection")


class SearchRequest(BaseModel):
    q: str
    k: int = 5


@app.get("/api/health")
def health():
    return {
        "status": "ok" if INDEX is not None else "error",
        "model": S.MODEL_ID,
        "tools": (len(INDEX["tools"]) if INDEX is not None else 0),
        "domains": (INDEX["domains"] if INDEX is not None else []),
    }


@app.post("/api/search")
def do_search(req: SearchRequest):
    if INDEX is None:
        return JSONResponse({"error": INDEX_ERROR or "index not ready"}, status_code=503)
    q = (req.q or "").strip()
    total = len(INDEX["tools"])
    if not q:
        return {"query": "", "ms": 0, "total": total, "results": []}
    t0 = time.time()
    results = S.search(INDEX, q, k=max(1, min(int(req.k), 12)))
    return {"query": q, "ms": int((time.time() - t0) * 1000), "total": total, "results": results}


@app.get("/")
def index():
    return FileResponse(os.path.join(STATIC, "index.html"))


app.mount("/", StaticFiles(directory=STATIC), name="static")