app.py

"""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"""
        <tr>
          <td>{status} {ep['model_id']}</td>
          <td>{ep['source']}</td>
          <td>{ep['gpu_type']}</td>
          <td><code>{ep['base_url'][:60]}...</code></td>
          <td>{age}s ago</td>
        </tr>"""

    return HTMLResponse(f"""<!DOCTYPE html>
<html><head><title>ThoxRoute Mesh Controller</title>
<style>
  body {{ font-family: -apple-system, sans-serif; max-width: 900px; margin: 40px auto; padding: 20px; }}
  h1 {{ color: #00D9FF; }}
  table {{ border-collapse: collapse; width: 100%; }}
  th, td {{ border: 1px solid #333; padding: 8px 12px; text-align: left; }}
  th {{ background: #1a1a2e; color: #00D9FF; }}
  .stat {{ display: inline-block; padding: 10px 20px; margin: 5px; background: #1a1a2e; border-radius: 8px; }}
  .stat strong {{ color: #00D9FF; font-size: 24px; }}
  code {{ font-size: 12px; }}
</style></head>
<body>
<h1>🧭 ThoxRoute Mesh Controller</h1>
<div>
  <div class="stat"><strong>{len(healthy)}</strong> healthy</div>
  <div class="stat"><strong>{len(unhealthy)}</strong> unhealthy</div>
  <div class="stat"><strong>{len(endpoints)}</strong> total</div>
</div>
<h2>Mesh Endpoints</h2>
<table>
  <tr><th>Model</th><th>Source</th><th>GPU</th><th>Endpoint</th><th>Last Seen</th></tr>
  {rows if rows else '<tr><td colspan="5">No endpoints registered. Start a Colab notebook to join the mesh.</td></tr>'}
</table>
<h2>API</h2>
<ul>
  <li><code>POST /v1/chat/completions</code> — OpenAI-compatible (routes through mesh)</li>
  <li><code>GET /v1/models</code> — List active mesh models</li>
  <li><code>POST /mesh/register</code> — Register a model endpoint</li>
  <li><code>POST /mesh/deregister</code> — Deregister a model endpoint</li>
  <li><code>GET /mesh/status</code> — This dashboard (JSON)</li>
</ul>
</body></html>""")


@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)