app.py

#!/usr/bin/env python3
"""Switchyard Lab — interactive multi-model routing playground.

Offline-first: routes scenarios with educational algorithms and shows
tokenomics. Optional live chat when OPENROUTER_API_KEY / NVIDIA_API_KEY set.
"""

from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Any

import httpx
import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

from switchyard_ops.metrics import summarize_run  # noqa: E402
from switchyard_ops.models import Step, StepType  # noqa: E402
from switchyard_ops.routers import (  # noqa: E402
    EscalationRouter,
    LlmClassifierRouter,
    RandomRouter,
    SimpleSwitchyard,
    StageRouter,
)
from switchyard_ops.loaders import load_scenarios, scenario_steps  # noqa: E402

load_dotenv(ROOT / ".env")

HOST = os.getenv("HOST", "127.0.0.1")
PORT = int(os.getenv("PORT", "7871"))
OFFLINE = os.getenv("OFFLINE", "").strip().lower() in {"1", "true", "yes"}

STATIC_DIR = ROOT / "static"
HEADER_SRC = ROOT / "assets" / "header.jpg"
HEADER_DST = STATIC_DIR / "header.jpg"

app = FastAPI(title="Switchyard Lab", version="0.1.0")


def _ensure_header() -> None:
    STATIC_DIR.mkdir(parents=True, exist_ok=True)
    if HEADER_SRC.exists() and (
        not HEADER_DST.exists() or HEADER_SRC.stat().st_mtime > HEADER_DST.stat().st_mtime
    ):
        HEADER_DST.write_bytes(HEADER_SRC.read_bytes())


_ensure_header()
app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")


STRATEGIES = {
    "policy": "Plan/execute policy (SimpleSwitchyard)",
    "stage": "Stage router (tool/error signals)",
    "classifier": "LLM classifier (content score)",
    "escalation": "Escalation (weak-first + judge)",
    "random": "Random 30/70 strong/weak",
}


def _make_router(name: str):
    if name == "policy":
        return SimpleSwitchyard()
    if name == "stage":
        return StageRouter()
    if name == "classifier":
        return LlmClassifierRouter(base_threshold=0.5)
    if name == "escalation":
        return EscalationRouter()
    if name == "random":
        return RandomRouter(weights={"strong": 0.3, "weak": 0.7})
    raise HTTPException(400, f"unknown strategy {name}")


@app.get("/")
def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/api/meta")
def meta() -> dict[str, Any]:
    return {
        "name": "Switchyard Lab",
        "version": (ROOT / "VERSION").read_text(encoding="utf-8").strip(),
        "strategies": STRATEGIES,
        "offline": OFFLINE or not _api_key(),
        "upstream": "https://github.com/NVIDIA-NeMo/Switchyard",
    }


@app.get("/api/scenarios")
def scenarios() -> list[dict[str, Any]]:
    return [
        {"id": s["id"], "title": s["title"], "description": s["description"], "steps": len(s["steps"])}
        for s in load_scenarios()
    ]


class SimulateBody(BaseModel):
    scenario_id: str = "refund-investigation"
    strategy: str = "stage"


@app.post("/api/simulate")
def simulate(body: SimulateBody) -> dict[str, Any]:
    scenarios_list = load_scenarios()
    sc = next((s for s in scenarios_list if s["id"] == body.scenario_id), None)
    if not sc:
        raise HTTPException(404, "scenario not found")
    steps = scenario_steps(sc)
    router = _make_router(body.strategy)
    decisions = [router.route(s) for s in steps]
    report = summarize_run(steps, decisions)
    return {
        "scenario": {"id": sc["id"], "title": sc["title"], "description": sc["description"]},
        "strategy": body.strategy,
        "steps": [
            {
                "type": s.type.value,
                "content": s.content,
                "signals": s.signals,
                "risk": s.risk,
                "failures": s.failures,
                "tokens_in": s.tokens_in,
                "tokens_out": s.tokens_out,
                "target": d.target,
                "model_id": d.model_id,
                "reason": d.reason,
                "score": d.score,
            }
            for s, d in zip(steps, decisions)
        ],
        "report": report,
    }


class CompareBody(BaseModel):
    scenario_id: str = "coding-pr"


@app.post("/api/compare")
def compare(body: CompareBody) -> dict[str, Any]:
    scenarios_list = load_scenarios()
    sc = next((s for s in scenarios_list if s["id"] == body.scenario_id), None)
    if not sc:
        raise HTTPException(404, "scenario not found")
    steps = scenario_steps(sc)
    rows = []
    for name in STRATEGIES:
        router = _make_router(name)
        decisions = [router.route(s) for s in steps]
        r = summarize_run(steps, decisions)
        rows.append(
            {
                "strategy": name,
                "label": STRATEGIES[name],
                "routed_cost_usd": r["routed_cost_usd"],
                "frontier_only_usd": r["frontier_only_usd"],
                "savings_pct": r["savings_pct"],
                "tier_counts": r["tier_counts"],
            }
        )
    return {"scenario_id": body.scenario_id, "rows": rows}


def _api_key() -> str:
    return (
        os.getenv("OPENROUTER_API_KEY", "").strip()
        or os.getenv("NVIDIA_API_KEY", "").strip()
        or os.getenv("OPENAI_API_KEY", "").strip()
    )


def _base_url() -> str:
    raw = (
        os.getenv("OPENAI_BASE_URL", "").strip()
        or os.getenv("NVIDIA_BASE_URL", "").strip()
        or "https://openrouter.ai/api/v1"
    )
    base = raw.rstrip("/")
    if base and not base.endswith("/v1"):
        base = f"{base}/v1"
    return base


class ChatBody(BaseModel):
    message: str = Field(..., min_length=1)
    model: str = "openai/gpt-4o-mini"


@app.post("/api/chat")
def chat(body: ChatBody) -> dict[str, Any]:
    if OFFLINE or not _api_key():
        # Educational offline reply
        router = LlmClassifierRouter()
        step = Step(StepType.EXECUTE, body.message, complexity="high" if len(body.message) > 80 else "low")
        d = router.route(step)
        return {
            "offline": True,
            "model": d.model_id,
            "target": d.target,
            "reason": d.reason,
            "content": (
                f"[offline simulation] Classifier would route this prompt to **{d.target}** "
                f"({d.model_id}). Reason: {d.reason}. Set OPENROUTER_API_KEY for a live completion."
            ),
        }
    try:
        with httpx.Client(timeout=60.0) as client:
            r = client.post(
                f"{_base_url()}/chat/completions",
                headers={
                    "Authorization": f"Bearer {_api_key()}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": body.model,
                    "messages": [
                        {
                            "role": "system",
                            "content": "You are a concise assistant explaining multi-model LLM routing.",
                        },
                        {"role": "user", "content": body.message},
                    ],
                },
            )
            r.raise_for_status()
            data = r.json()
            content = data["choices"][0]["message"]["content"]
            return {"offline": False, "model": body.model, "content": content}
    except Exception as e:
        raise HTTPException(502, f"upstream error: {e}") from e


def main() -> None:
    print(f"Switchyard Lab → http://{HOST}:{PORT}")
    uvicorn.run(app, host=HOST, port=PORT, log_level="info")


if __name__ == "__main__":
    main()