"""ThoxRoute Mesh Controller — FastAPI app for the always-on HF Space. This is the front-door for the ThoxRoute Colab mesh. It runs on a free CPU HF Space (always-on), provides: - /v1/chat/completions (OpenAI-compatible, routes to mesh) - /v1/models (list active mesh models) - /mesh/register (Colab notebooks register their ngrok URLs) - /mesh/deregister (clean shutdown) - /mesh/status (health dashboard) - /mesh/health (JSON health for monitoring) - / (web UI dashboard) The controller loads ThoxRoute-1.5B as a CPU sidecar (via Ollama) for stage-1 intent classification, then forwards the request to the best available mesh endpoint (Colab notebook or HF Space model server). """ from __future__ import annotations import json import os import time import logging import subprocess import sys import urllib.request import urllib.error from typing import Optional from fastapi import FastAPI, Request, HTTPException from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse from pydantic import BaseModel from mesh import MeshEndpoint, get_registry, open_endpoint logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(name)s] %(message)s") logger = logging.getLogger("thoxroute.controller") app = FastAPI(title="ThoxRoute Mesh Controller", version="0.1.0") # ── Route → preferred models (same as thox_routes.yaml) ─────────────────── ROUTE_MODEL_MAP: dict[str, list[str]] = { "thox.chat.general": ["thoxmythos-9b", "thoxintel-27b", "thoxroute-1.5b"], "thox.reasoning.deep": ["thoxintel-27b", "thoxmythos-9b"], "thox.reasoning.frontier": ["thoxintel-27b"], "thox.code.generation": ["thoxmythos-9b", "thoxintel-27b"], "thox.summarize.longcontext": ["thoxintel-27b", "thoxmythos-9b"], "thox.local.private": [], # never route to Colab "thox.agentic.task": ["thoxmythos-9b", "thoxintel-27b"], } ROUTE_DESCRIPTIONS: dict[str, str] = { "thox.chat.general": "General conversation, open-ended questions, everyday assistant chat", "thox.reasoning.deep": "Complex problem solving, hard reasoning, math, logic, proofs", "thox.reasoning.frontier": "Hardest problems, frontier-level reasoning, research-grade analysis", "thox.code.generation": "Writing, editing, explaining, refactoring, debugging source code", "thox.summarize.longcontext": "Summarizing or extracting from long documents, transcripts", "thox.local.private": "Private, confidential, on-device requests", "thox.agentic.task": "Multi-step autonomous tasks, build, implement, refactor", } # ── Classifier sidecar (ThoxRoute-1.5B via Ollama) ─────────────────────── CLASSIFIER_BASE_URL = os.environ.get("THOXROUTE_CLASSIFIER_BASE_URL", "http://127.0.0.1:11434/v1") CLASSIFIER_MODEL = os.environ.get("THOXROUTE_CLASSIFIER_MODEL", "thoxroute-1.5b") CLASSIFIER_FALLBACK_MODEL = os.environ.get("THOXROUTE_CLASSIFIER_FALLBACK_MODEL", "gemma3:4b") # ── Request/Response models ───────────────────────────────────────────── class ChatMessage(BaseModel): role: str content: str class ChatRequest(BaseModel): model: str = "" messages: list[ChatMessage] temperature: float = 0.7 max_tokens: int = 2048 stream: bool = False class RegisterRequest(BaseModel): model_id: str base_url: str source: str = "colab" gpu_type: str = "T4" space_id: str = "" space_url: str = "" # Node-tier fields sent by thoxmesh_node. Declared explicitly because # pydantic drops undeclared fields silently -- which previously meant the # node's lease and telemetry vanished on arrival and consumers could not # tell a live node from one whose Colab VM had been reclaimed. alias: str = "" protocol: str = "openai" context_window: int = 0 status: str = "unknown" capabilities: dict = {} # ── Classifier ──────────────────────────────────────────────────────────── def classify_intent(user_message: str) -> str: """Use ThoxRoute-1.5B sidecar to classify user intent.""" routes_text = "\n".join(f"- {name}: {desc}" for name, desc in ROUTE_DESCRIPTIONS.items()) prompt = ( f"You are ThoxRoute, an intent classifier. Given the user message, " f"classify it into exactly ONE of these routes:\n\n{routes_text}\n\n" f"User message: {user_message}\n\n" f"Respond with ONLY the route name, nothing else." ) payload = { "model": CLASSIFIER_MODEL, "messages": [ {"role": "system", "content": "You are ThoxRoute intent classifier."}, {"role": "user", "content": prompt}, ], "temperature": 0, "max_tokens": 64, } try: req = urllib.request.Request( f"{CLASSIFIER_BASE_URL}/chat/completions", data=json.dumps(payload).encode(), headers={"Content-Type": "application/json"}, ) resp = json.loads(urllib.request.urlopen(req, timeout=20).read()) text = resp["choices"][0]["message"]["content"].strip() for route_name in ROUTE_DESCRIPTIONS: if route_name in text: return route_name except Exception as e: logger.warning(f"Classifier failed: {e}, using fallback") return "thox.chat.general" # ── Mesh proxy ──────────────────────────────────────────────────────────── def forward_to_mesh( messages: list[dict], temperature: float, max_tokens: int, route: str, model_id: str, base_url: str, ) -> dict: """Forward a request to a mesh endpoint and return the response.""" payload = { "model": model_id, "messages": messages, "temperature": temperature, "max_tokens": max_tokens, } # Endpoints register their base_url two different ways, and both are valid: # the original Colab notebooks register the OpenAI base (".../v1"), while # thoxmesh_node registers the node's bare origin. Appending # "/chat/completions" blindly 404s the latter, so normalise here rather than # forcing one convention on nodes that are already deployed. root = base_url.rstrip("/") if not root.endswith("/v1"): root = f"{root}/v1" target = f"{root}/chat/completions" # Unauthenticated first, retried with the HF token only on an auth-shaped status. # A public Space rejects a supplied credential on the Space-to-Space path; a private # one 404s without it. See open_endpoint(). resp = json.loads( open_endpoint(target, data=json.dumps(payload).encode(), timeout=120).read() ) return resp def route_and_complete(messages: list[dict], temperature: float, max_tokens: int) -> dict: """Classify intent, select mesh endpoint, forward request.""" user_msg = next((m["content"] for m in reversed(messages) if m["role"] == "user"), "") route = classify_intent(user_msg) registry = get_registry() preferred = ROUTE_MODEL_MAP.get(route, []) ep = registry.select_for_route(preferred) if not ep: return { "route": route, "model": None, "error": f"No available mesh endpoint for route '{route}'. Start a Colab notebook or HF Space model server.", "response": None, } try: result = forward_to_mesh(messages, temperature, max_tokens, route, ep.model_id, ep.base_url) return { "route": route, "model": ep.model_id, "endpoint": ep.base_url, "source": ep.source, "response": result["choices"][0]["message"]["content"], } except Exception as e: # Mark endpoint as potentially down with registry._lock: if ep.model_id in registry._endpoints: registry._endpoints[ep.model_id].healthy = False return { "route": route, "model": ep.model_id, "error": f"Endpoint {ep.base_url} failed: {e}", "response": None, } # ── API endpoints ───────────────────────────────────────────────────────── @app.get("/") async def dashboard(): """Mesh status dashboard.""" registry = get_registry() endpoints = registry.list_all() healthy = [e for e in endpoints if e["healthy"]] unhealthy = [e for e in endpoints if not e["healthy"]] rows = "" for ep in endpoints: status = "🟢" if ep["healthy"] else "🔴" age = int(time.time() - ep.get("last_seen", ep.get("registered_at", 0))) rows += f""" {status} {ep['model_id']} {ep['source']} {ep['gpu_type']} {ep['base_url'][:60]}... {age}s ago """ return HTMLResponse(f""" ThoxRoute Mesh Controller

🧭 ThoxRoute Mesh Controller

{len(healthy)} healthy
{len(unhealthy)} unhealthy
{len(endpoints)} total

Mesh Endpoints

{rows if rows else ''}
ModelSourceGPUEndpointLast Seen
No endpoints registered. Start a Colab notebook to join the mesh.

API

""") @app.get("/v1/models") async def list_models(): """OpenAI-compatible models list — returns active mesh endpoints.""" registry = get_registry() models = [] for ep in registry.list_healthy(): models.append({"id": ep.model_id, "object": "model", "owned_by": "thox"}) return {"object": "list", "data": models} @app.post("/v1/chat/completions") async def chat_completions(req: ChatRequest): """OpenAI-compatible chat completions — routes through the mesh.""" messages = [m.model_dump() for m in req.messages] result = route_and_complete(messages, req.temperature, req.max_tokens) if result.get("error"): raise HTTPException(status_code=503, detail=result["error"]) return { "id": f"thoxroute-{int(time.time())}", "object": "chat.completion", "created": int(time.time()), "model": result["model"], "choices": [{ "index": 0, "message": {"role": "assistant", "content": result["response"]}, "finish_reason": "stop", }], "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, "thoxroute_meta": { "route": result["route"], "model": result["model"], "source": result.get("source", "unknown"), "endpoint": result.get("endpoint", ""), }, } @app.post("/mesh/register") async def register_endpoint(req: RegisterRequest): """Register a model endpoint (called by Colab notebooks and HF Spaces).""" registry = get_registry() ep = MeshEndpoint( model_id=req.model_id, base_url=req.base_url, source=req.source, gpu_type=req.gpu_type, registered_at=time.time(), space_id=req.space_id, space_url=req.space_url, alias=req.alias or req.model_id, protocol=req.protocol, context_window=req.context_window, status=req.status, capabilities=req.capabilities or {}, ) result = registry.register(ep) return result @app.post("/mesh/deregister") async def deregister_endpoint(request: Request): """Deregister a model endpoint.""" body = await request.json() model_id = body.get("model_id", "") if not model_id: raise HTTPException(status_code=400, detail="model_id required") registry = get_registry() return registry.deregister(model_id) @app.get("/mesh/status") async def mesh_status(): """JSON status for monitoring.""" registry = get_registry() endpoints = registry.list_all() return { "total": len(endpoints), "healthy": sum(1 for e in endpoints if e["healthy"]), "unhealthy": sum(1 for e in endpoints if not e["healthy"]), "endpoints": endpoints, } @app.get("/mesh/health") async def mesh_health(): """Simple health check.""" return {"status": "ok", "timestamp": time.time()} # ── Startup ─────────────────────────────────────────────────────────────── @app.on_event("startup") async def startup(): """Start the mesh registry and health monitor.""" registry = get_registry() logger.info("ThoxRoute Mesh Controller started") logger.info(f"Classifier: {CLASSIFIER_BASE_URL} (model={CLASSIFIER_MODEL})") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=7860)