api/routing.py

"""RosDiff's intelligence & routing layer: which model, which simulator, and what it will cost.

Every generation goes through `plan_scene` / `plan_code`, which look at the request (and what retrieval found
for it), rate its complexity, and choose the cheapest model tier that should handle it reliably. When a draft
fails validation, `model_for_attempt` escalates to the next tier for the repair, so hard tasks end up on the
strongest model without anyone picking it. Users never see a model picker; the plan's tier and reason are shown
as information.

There are two registries, one for scenes and one for code, of capable low-cost Chinese open-weight models on
OpenRouter chosen by benchmark and price rather than by release date (see README for sources). Either can be
replaced with SCENE_MODELS / CODE_MODELS (JSON) without code changes.

Debugging: the first repair after a failed check goes to a free model through OpenCode Zen (DEBUG_MODEL, default
opencode/deepseek-v4-flash-free), never through OpenRouter. If that call fails (rate limit, model withdrawn,
OpenCode missing) the same attempt falls back to the next paid tier, so a free model can never cost a repair.

Simulators: RosDiff scenes are MuJoCo (MJCF) scenes, so the routing chooses between what can run them today:
MuJoCo (one run with video), MuJoCo Warp (many randomized copies on the GPU) and Isaac Sim (a photoreal RTX
render of the same motion). Other simulators are listed as not available yet rather than pretended.
"""

from __future__ import annotations

import json
import os
import re
from dataclasses import asdict, dataclass, field

TIERS = ("low", "mid", "high")


@dataclass(frozen=True)
class ModelSpec:
    id: str  # OpenRouter model id
    tier: str  # low | mid | high
    provider: str
    input_usd_per_m: float
    output_usd_per_m: float
    note: str = ""
    free: bool = False  # served free through OpenCode Zen: recorded, never billed


# OpenRouter list prices (USD per million tokens, Sept 2026) and the benchmark each was chosen on.
CODE_REGISTRY = (
    ModelSpec("deepseek/deepseek-v4-flash", "low", "DeepSeek", 0.089, 0.177, "79.0% SWE-bench Verified; very cheap"),
    ModelSpec("deepseek/deepseek-v4-pro", "mid", "DeepSeek", 0.435, 0.87, "80.6% SWE-bench Verified"),
    ModelSpec(
        "z-ai/glm-5.2", "high", "Zhipu AI (Z.ai)", 0.49, 1.54, "62.1% SWE-bench Pro, best open model on hard code"
    ),
)
SCENE_REGISTRY = (
    ModelSpec(
        "deepseek/deepseek-v4-flash", "low", "DeepSeek", 0.089, 0.177, "listed for completeness; scenes start at mid"
    ),
    ModelSpec("z-ai/glm-5.2", "mid", "Zhipu AI (Z.ai)", 0.49, 1.54, "within a point of Claude Opus 4.8 on FrontierSWE"),
    ModelSpec(
        "moonshotai/kimi-k2.6", "high", "Moonshot AI", 0.74, 3.49, "80.2% SWE-bench Verified; top open-weight AA index"
    ),
)
DEFAULT_REGISTRY = CODE_REGISTRY  # kept for callers that don't say which task


def load_registry(task: str = "code") -> tuple[ModelSpec, ...]:
    """SCENE_MODELS / CODE_MODELS (or MODEL_REGISTRY for both) = '[{"id": ..., "tier": ..., ...}, ...]'."""
    raw = os.environ.get("SCENE_MODELS" if task == "scene" else "CODE_MODELS", "").strip()
    raw = raw or os.environ.get("MODEL_REGISTRY", "").strip()
    if not raw:
        return SCENE_REGISTRY if task == "scene" else CODE_REGISTRY
    specs = tuple(ModelSpec(**m) for m in json.loads(raw))
    if {s.tier for s in specs} != set(TIERS):
        raise ValueError("a model registry needs one model for each tier: low, mid, high")
    return specs


def all_models() -> dict[str, ModelSpec]:
    return {m.id: m for task in ("code", "scene") for m in load_registry(task)}


def model_for_tier(tier: str, registry: tuple[ModelSpec, ...] | None = None, task: str = "code") -> ModelSpec:
    registry = registry or load_registry(task)
    return next(m for m in registry if m.tier == tier)


def free_debugger(settings) -> ModelSpec | None:
    if not getattr(settings, "free_debug", False) or not settings.debug_model:
        return None
    return ModelSpec(settings.debug_model, "free", "OpenCode Zen", 0.0, 0.0, "free debugging model", free=True)


def is_free(model_id: str | None) -> bool:
    return bool(model_id) and model_id.startswith("opencode/")


def escalate(tier: str, steps: int = 1) -> str:
    return TIERS[min(TIERS.index(tier) + steps, len(TIERS) - 1)]


# --------------------------------------------------------------------------- complexity

_CODE_COMPLEX = (
    "nav2", "navigation stack", "moveit", "slam", "costmap", "planner", "behavior tree", "behaviour tree",
    "state machine", "lifecycle", "ros2_control", "hardware interface", "controller manager", "action server",
    "perception", "point cloud", "pointcloud", "kalman", "ekf", "localization", "localisation", "multi-robot",
    "fleet", "autonomous", "inverse kinematics", "trajectory", "mpc", "reinforcement",
)  # fmt: skip
_CODE_MEDIUM = (
    "service", "action", "parameter", "launch", "timer", "tf", "transform", "filter", "odometry", "odom",
    "lidar", "laser", "imu", "camera", "image", "twist", "subscriber", "publisher", "client", "path", "c++",
)  # fmt: skip
_SCENE_COMPLEX = (
    "maze", "warehouse", "rough terrain", "heightfield", "uneven", "hinge", "door that opens", "drawer", "joint",
    "articulated", "pendulum", "tendon", "cloth", "rope", "multiple robots", "two robots", "conveyor", "stairs with",
    "spiral", "ramp and stairs",
)  # fmt: skip
_SCENE_MEDIUM = (
    "stairs", "steps", "ramp", "slope", "table", "shelf", "wall", "doorway", "obstacle", "boxes", "cubes", "ball",
    "friction", "terrain", "tray",
)  # fmt: skip


def _count(text: str, words: tuple[str, ...]) -> list[str]:
    t = text.lower()
    return [w for w in words if w in t]


@dataclass
class Plan:
    task: str  # scene | code
    complexity: str  # simple | medium | complex
    tier: str
    model: str
    reason: str
    category: str  # simulation_ai | coding (repairs are booked as debugging)
    estimate_inr: tuple[float, float] = (0.0, 0.0)
    simulator: str | None = None
    simulator_reason: str | None = None
    signals: list[str] = field(default_factory=list)
    debugger: str | None = None  # the free model tried first on the first repair

    def to_dict(self) -> dict:
        d = asdict(self)
        d["estimate_inr"] = list(self.estimate_inr)
        return d


def _level(score: int) -> str:
    return "complex" if score >= 4 else "medium" if score >= 2 else "simple"


def _tier_for(level: str) -> str:
    return {"simple": "low", "medium": "mid", "complex": "high"}[level]


def plan_scene(description: str, robot: str | None, settings, n_numbers: int | None = None) -> Plan:
    complex_hits = _count(description, _SCENE_COMPLEX)
    medium_hits = _count(description, _SCENE_MEDIUM)
    objects = n_numbers if n_numbers is not None else len(re.findall(r"\d+(?:\.\d+)?", description))
    # everyday parts add up to 'medium' at most; only genuinely hard ones push a task to the top tier
    score = 3 * len(complex_hits) + min(len(medium_hits), 3) + (objects >= 6) + (len(description) > 400)
    level = _level(score)
    # Scenes must be spatially exact; the cheapest tier is only used for very small ones.
    tier = max(_tier_for(level), "mid", key=TIERS.index)
    signals = complex_hits + medium_hits
    reason = {
        "simple": "a small scene",
        "medium": "a scene with several objects to place exactly",
        "complex": "a scene with articulated or intricate parts",
    }[level]
    sim, sim_reason = choose_simulator(description)
    plan = Plan("scene", level, tier, model_for_tier(tier, task="scene").id, reason, "simulation_ai", signals=signals)
    plan.simulator, plan.simulator_reason = sim, sim_reason
    plan.debugger = (d := free_debugger(settings)) and d.id
    plan.estimate_inr = estimate_inr(plan, settings)
    return plan


def plan_code(request: str, distro: str, language: str, settings, n_interfaces: int = 0) -> Plan:
    complex_hits = _count(request, _CODE_COMPLEX)
    medium_hits = _count(request, _CODE_MEDIUM)
    score = 3 * len(complex_hits) + min(len(medium_hits), 3) + (1 if language == "cpp" else 0)
    score += 1 if n_interfaces >= 3 else 0
    score += 1 if len(request) > 500 else 0
    level = _level(score)
    tier = _tier_for(level)
    reason = {
        "simple": "a small, self-contained ROS 2 node",
        "medium": "a ROS 2 package with several parts to connect",
        "complex": "a ROS 2 system with navigation, control or planning",
    }[level]
    if language == "cpp" and level != "complex":
        reason += " in C++"
    plan = Plan("code", level, tier, model_for_tier(tier).id, reason, "coding", signals=complex_hits + medium_hits)
    plan.debugger = (d := free_debugger(settings)) and d.id
    plan.estimate_inr = estimate_inr(plan, settings)
    return plan


def candidates_for_attempt(plan: Plan, attempt: int) -> list[ModelSpec]:
    """Models to try for one attempt, in order (the next is only used if a call fails outright).

    Attempt 0: the planned tier. Attempt 1: the free debugger, falling back to one tier up. Later repairs keep
    going up one tier each. Without a free debugger every repair simply goes one tier up."""
    registry = load_registry(plan.task)
    if attempt == 0:
        return [model_for_tier(plan.tier, registry)]
    if plan.debugger:
        paid = model_for_tier(escalate(plan.tier, attempt), registry)
        if attempt == 1:
            return [ModelSpec(plan.debugger, "free", "OpenCode Zen", 0.0, 0.0, free=True), paid]
        return [paid]
    return [model_for_tier(escalate(plan.tier, attempt), registry)]


def model_for_attempt(plan: Plan, attempt: int) -> ModelSpec:
    """The paid model an attempt ends on in the worst case."""
    return candidates_for_attempt(plan, attempt)[-1]


# --------------------------------------------------------------------------- simulators

SIMULATORS = {
    "mujoco": {"label": "MuJoCo", "what": "one run with video and 3D replay", "available": True},
    "mujoco_warp": {"label": "MuJoCo Warp", "what": "64 randomized copies at once on the GPU", "available": True},
    "isaac_sim": {"label": "Isaac Sim", "what": "photoreal RTX render of the same motion", "available": True},
    "gazebo": {"label": "Gazebo", "what": "not available in RosDiff yet", "available": False},
    "webots": {"label": "Webots", "what": "not available in RosDiff yet", "available": False},
    "pybullet": {"label": "PyBullet", "what": "not available in RosDiff yet", "available": False},
}

_WANTS_RENDER = ("photoreal", "realistic render", "camera image", "rgb", "lighting", "vision", "synthetic image",
                 "rendered image", "ray tracing", "ray-traced")  # fmt: skip
_WANTS_MANY = ("robust", "random", "variation", "many times", "stress", "domain randomization", "monte carlo",
               "reinforcement", "train a policy", "statistics", "how often")  # fmt: skip


def choose_simulator(text: str) -> tuple[str, str]:
    if hits := _count(text, _WANTS_MANY):
        return "mujoco_warp", f"you asked about {hits[0]}: many randomized copies show how often it works"
    if hits := _count(text, _WANTS_RENDER):
        return "isaac_sim", f"you mentioned {hits[0]}: MuJoCo simulates, Isaac Sim renders it photoreal"
    return "mujoco", "contact-rich rigid-body physics with the real robot model, fast and precise"


# --------------------------------------------------------------------------- cost

# Typical tokens per model call, measured from RosDiff's prompts (context + robot facts / interface definitions).
TOKENS = {
    ("scene", "direct"): (16_000, 3_500),
    ("code", "direct"): (20_000, 6_000),
    ("code", "opencode"): (60_000, 9_000),  # OpenCode's agent loop re-reads context across several steps
}


def call_cost_usd(model: ModelSpec, input_tokens: int, output_tokens: int) -> float:
    return (input_tokens * model.input_usd_per_m + output_tokens * model.output_usd_per_m) / 1e6


def estimate_inr(plan: Plan, settings) -> tuple[float, float]:
    """(likely, worst case): one call on the planned tier vs. every attempt, escalating."""
    mode = "opencode" if plan.task == "code" and settings.codegen_backend == "opencode" else "direct"
    tin, tout = TOKENS[(plan.task, mode)]
    attempts = settings.mjcf_max_attempts if plan.task == "scene" else settings.code_max_attempts
    first = call_cost_usd(model_for_attempt(plan, 0), tin, tout)
    # worst case: every attempt runs, and the free debugger fails so its paid fallback is billed
    worst = sum(call_cost_usd(model_for_attempt(plan, i), tin * (1 + i // 2), tout) for i in range(attempts))
    return round(first * settings.usd_inr, 2), round(worst * settings.usd_inr, 2)