atlas-iac/services/hermes/scripts/hermes_model_routing.py

#!/usr/bin/env python3
"""Keep Hermes agent provider catalogs and managed profiles current."""

from __future__ import annotations

import copy
import json
import os
import re
import shutil
import subprocess
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable

import yaml


CODEX_BASELINE = "gpt-5.6-terra"
CLAUDE_BASELINE = "claude-opus-5"
CLAUDE_SUBSCRIPTION_MODELS = (
    "claude-haiku-4-5-20251001",
    "claude-fable-5",
    "claude-sonnet-5",
    "claude-opus-5",
)
EFFORTS = ("low", "medium", "high", "xhigh")
ATLAS_FALLBACK = {
    "provider": "custom",
    "model": "qwen2.5:14b-instruct-q4_0",
    "base_url": "http://hermes-model-gate.hermes.svc.cluster.local:11434/v1",
    "api_key": "ollama",
}
# Backwards-compatible name used by the focused unit tests and status tooling.
LOCAL_FALLBACK = ATLAS_FALLBACK
SWITCHYARD_PROVIDER = "atlas-switchyard"
SWITCHYARD_API = "http://hermes-switchyard.hermes.svc.cluster.local:9005/v1"
SWITCHYARD_AUTO_ROUTE = "atlas/auto/maximum"
ROUTING_CATALOG_PATH = os.environ.get("HERMES_ROUTING_CATALOG_PATH", "").strip()
MANAGED_ENV_KEYS = {
    "GIT_ASKPASS",
    "GIT_TERMINAL_PROMPT",
}
RUNTIME_SECRET_ENV_KEYS = {
    "ANTHROPIC_API_KEY",
    "API_SERVER_KEY",
    "CLAUDE_API_KEY",
    "CLAUDE_CODE_OAUTH_TOKEN",
    "GITEA_TOKEN",
    "GITEA_USERNAME",
    "HERMES_IMAGE_BROKER_KEY",
    "OPENAI_API_KEY",
}


@dataclass(frozen=True)
class Catalog:
    """Non-secret provider discovery result."""

    provider: str
    models: list[str]
    live: bool
    connected: bool
    state: str


def _unique_models(models: Iterable[str]) -> list[str]:
    """Return normalized model IDs without changing provider order."""
    seen: set[str] = set()
    result: list[str] = []
    for model in models:
        value = str(model or "").strip()
        if value and value.lower() not in seen:
            result.append(value)
            seen.add(value.lower())
    return result


def model_version(model: str) -> tuple[int, ...]:
    """Extract a sortable model version while ignoring dated aliases."""
    value = re.sub(r"-\d{8}$", "", model.lower())
    groups = re.findall(r"\d+(?:\.\d+)*", value)
    parts: list[int] = []
    for group in groups:
        parts.extend(int(piece) for piece in group.split("."))
    return tuple(parts) or (0,)


def _codex_quality(model: str) -> int:
    """Rank known Codex capability tiers for delegated implementation work."""
    value = model.lower()
    if "sol" in value:
        return 60
    if "terra" in value:
        return 50
    if "codex" in value and "spark" not in value:
        return 45
    if re.fullmatch(r"gpt-\d+(?:\.\d+)*", value):
        return 40
    if "luna" in value:
        return 20
    if "mini" in value or "nano" in value:
        return 10
    if "spark" in value:
        return 5
    return 0


def choose_codex_model(
    models: Iterable[str], current: str = CODEX_BASELINE, *, balanced: bool = False
) -> str:
    """Choose a current full Codex model, retaining current when unknown."""
    candidates = [m for m in _unique_models(models) if m.lower().startswith("gpt-")]
    full = [m for m in candidates if _codex_quality(m) >= 40]
    pool = full or candidates
    if not pool:
        return current
    newest = max(model_version(model) for model in pool)
    latest = [model for model in pool if model_version(model) == newest]
    if balanced:
        terra = [model for model in latest if "terra" in model.lower()]
        if terra:
            return max(terra, key=lambda model: (_codex_quality(model), model))
    return max(latest, key=lambda model: (_codex_quality(model), model))


def _choose_by_hints(
    models: Iterable[str], hints: tuple[str, ...], current: str, prefix: str
) -> str:
    """Pick the newest model in the first available capability class."""
    candidates = [
        model
        for model in _unique_models(models)
        if model.lower().startswith(prefix)
    ]
    if not candidates:
        return current
    for hint in hints:
        tier = [model for model in candidates if hint in model.lower()]
        if tier:
            return max(tier, key=lambda model: (model_version(model), model))
    return max(candidates, key=lambda model: (model_version(model), model))


def choose_codex_for_effort(
    models: Iterable[str], effort: str, current: str = CODEX_BASELINE
) -> str:
    """Choose the account-visible Codex tier for an allowed effort level."""
    hints = {
        "low": ("luna", "mini", "spark", "terra", "sol", "codex"),
        "medium": ("terra", "codex", "sol", "luna", "mini", "spark"),
        "high": ("sol", "codex", "terra", "luna", "mini", "spark"),
        "xhigh": ("sol", "codex", "terra", "luna", "mini", "spark"),
    }
    if effort not in EFFORTS:
        raise ValueError(f"unsupported effort: {effort}")
    return _choose_by_hints(models, hints[effort], current, "gpt-")


def _claude_quality(model: str) -> int:
    """Rank known Claude capability tiers for architecture and review work."""
    value = model.lower()
    if "opus" in value:
        return 40
    if "sonnet" in value:
        return 30
    if "fable" in value:
        return 20
    if "haiku" in value:
        return 10
    return 0


def choose_claude_model(models: Iterable[str], current: str = CLAUDE_BASELINE) -> str:
    """Choose the newest full Claude reasoning model visible to the account."""
    candidates = [m for m in _unique_models(models) if m.lower().startswith("claude-")]
    full = [m for m in candidates if _claude_quality(m) >= 30]
    pool = full or candidates
    if not pool:
        return current
    return max(
        pool,
        key=lambda model: (model_version(model), _claude_quality(model), model),
    )


def choose_claude_for_effort(
    models: Iterable[str], effort: str, current: str = CLAUDE_BASELINE
) -> str:
    """Choose the account-visible Claude tier for an allowed effort level."""
    hints = {
        "low": ("haiku", "fable", "sonnet", "opus"),
        "medium": ("sonnet", "fable", "opus", "haiku"),
        "high": ("opus", "sonnet", "fable", "haiku"),
        "xhigh": ("opus", "sonnet", "fable", "haiku"),
    }
    if effort not in EFFORTS:
        raise ValueError(f"unsupported effort: {effort}")
    return _choose_by_hints(models, hints[effort], current, "claude-")


def _read_yaml(path: Path) -> dict[str, Any]:
    """Read a mapping from YAML, returning an empty mapping when unavailable."""
    if not path.is_file():
        return {}
    try:
        value = yaml.safe_load(path.read_text(encoding="utf-8"))
    except (OSError, yaml.YAMLError):
        return {}
    return value if isinstance(value, dict) else {}


def _atomic_write(path: Path, content: str, mode: int | None = None) -> bool:
    """Replace a file only when its content changes."""
    path.parent.mkdir(parents=True, exist_ok=True)
    try:
        if path.read_text(encoding="utf-8") == content:
            if mode is not None:
                path.chmod(mode)
            return False
    except OSError:
        pass
    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    temporary.write_text(content, encoding="utf-8")
    if mode is not None:
        temporary.chmod(mode)
    os.replace(temporary, path)
    return True


def _write_yaml(path: Path, value: dict[str, Any], mode: int | None = None) -> bool:
    """Serialize a mapping and atomically update the target YAML file."""
    return _atomic_write(path, yaml.safe_dump(value, sort_keys=False), mode)


def _read_json(path: Path) -> dict[str, Any]:
    """Read a JSON mapping without treating a partial write as valid state."""
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, ValueError, json.JSONDecodeError):
        return {}
    return value if isinstance(value, dict) else {}


def _previous_provider_models(
    previous: dict[str, Any], provider: str
) -> tuple[dict[str, str], dict[str, str], list[str]]:
    """Return last-known-good effort, tier, and catalog values for a provider."""
    providers = previous.get("providers", {})
    record = providers.get(provider, {}) if isinstance(providers, dict) else {}
    if not isinstance(record, dict):
        return {}, {}, []
    resolved = record.get("resolved", {})
    tiers = record.get("tiers", {})
    models = record.get("models", [])
    return (
        dict(resolved) if isinstance(resolved, dict) else {},
        dict(tiers) if isinstance(tiers, dict) else {},
        _unique_models(models if isinstance(models, list) else []),
    )


def build_routing_catalog(
    codex: Catalog, claude: Catalog, previous: dict[str, Any] | None = None
) -> dict[str, Any]:
    """Build a current catalog while retaining working routes during outages."""
    previous = previous or {}
    providers: dict[str, Any] = {}
    specifications = (
        (
            "codex",
            codex,
            choose_codex_for_effort,
            {"luna": "low", "terra": "medium", "sol": "high"},
            {
                "low": "gpt-5.6-luna",
                "medium": "gpt-5.6-terra",
                "high": "gpt-5.6-sol",
                "xhigh": "gpt-5.6-sol",
            },
        ),
        (
            "claude",
            claude,
            choose_claude_for_effort,
            {
                "haiku": "low",
                "fable": "medium",
                "sonnet": "high",
                "opus": "xhigh",
            },
            {
                "low": "claude-haiku-4-5-20251001",
                "medium": "claude-sonnet-5",
                "high": "claude-opus-5",
                "xhigh": "claude-opus-5",
            },
        ),
    )
    for name, discovered, chooser, tier_efforts, defaults in specifications:
        old_resolved, old_tiers, old_models = _previous_provider_models(previous, name)
        source_models = discovered.models if discovered.live else old_models
        resolved: dict[str, str] = {}
        for effort in EFFORTS:
            current = str(old_resolved.get(effort) or defaults[effort])
            resolved[effort] = (
                chooser(source_models, effort, current)
                if source_models
                else current
            )
        tiers: dict[str, str] = {}
        for tier, effort in tier_efforts.items():
            current = str(old_tiers.get(tier) or defaults[effort])
            tier_matches = [model for model in source_models if tier in model.lower()]
            tiers[tier] = (
                max(tier_matches, key=lambda model: (model_version(model), model))
                if tier_matches
                else (resolved[effort] if discovered.live else current)
            )
        providers[name] = {
            "state": discovered.state,
            "connected": discovered.connected,
            "live": discovered.live,
            "models": _unique_models(
                discovered.models
                if discovered.live
                else (old_models or discovered.models)
            ),
            "resolved": resolved,
            "tiers": tiers,
        }
    return {
        "schema_version": 1,
        "updated_at": int(time.time()),
        "providers": providers,
    }


def write_routing_catalog(
    path: Path, codex: Catalog, claude: Catalog
) -> dict[str, Any]:
    """Atomically publish the catalog consumed by hosted and worker brokers."""
    catalog = build_routing_catalog(codex, claude, _read_json(path))
    _atomic_write(path, json.dumps(catalog, indent=2, sort_keys=True) + "\n", 0o644)
    return catalog


def _read_env(path: Path) -> dict[str, str]:
    """Read the small dotenv subset used by Hermes provider credentials."""
    values: dict[str, str] = {}
    try:
        lines = path.read_text(encoding="utf-8").splitlines()
    except OSError:
        return values
    for line in lines:
        value = line.strip()
        if not value or value.startswith("#") or "=" not in value:
            continue
        key, raw = value.removeprefix("export ").split("=", 1)
        raw = raw.strip()
        if len(raw) >= 2 and raw[0] == raw[-1] and raw[0] in "\"'":
            raw = raw[1:-1]
        values[key.strip()] = raw
    return values


def _update_profile_env(path: Path, source: dict[str, str]) -> None:
    """Refresh non-secret managed settings and remove stale credential copies."""
    try:
        old_lines = path.read_text(encoding="utf-8").splitlines()
    except OSError:
        old_lines = []
    kept = [
        line
        for line in old_lines
        if not any(
            line.lstrip().startswith(f"{key}=")
            for key in MANAGED_ENV_KEYS | RUNTIME_SECRET_ENV_KEYS
        )
    ]
    kept.extend(
        f"{key}={source[key]}" for key in sorted(MANAGED_ENV_KEYS) if source.get(key)
    )
    _atomic_write(path, "\n".join(kept).rstrip() + "\n", 0o600)


def _existing_model(config: dict[str, Any], provider: str, default: str) -> str:
    """Return the last configured model for a provider or a safe baseline."""
    model_cfg = config.get("model", {})
    if isinstance(model_cfg, dict) and model_cfg.get("provider") == provider:
        return str(model_cfg.get("model") or model_cfg.get("default") or default)
    fallbacks = config.get("fallback_providers", [])
    if isinstance(fallbacks, list):
        for fallback in fallbacks:
            if isinstance(fallback, dict) and fallback.get("provider") == provider:
                return str(fallback.get("model") or default)
    return default


def codex_cli_authenticated() -> bool:
    """Return whether the installed Codex CLI has a usable local login."""
    codex = shutil.which("codex")
    if codex:
        try:
            status = subprocess.run(
                [codex, "login", "status"],
                capture_output=True,
                check=False,
                text=True,
                timeout=10,
            )
            detail = f"{status.stdout}\n{status.stderr}".lower()
            if status.returncode == 0 and (
                "logged in" in detail or "authenticated" in detail
            ):
                return True
        except (OSError, subprocess.SubprocessError):
            pass
    return False


def discover_codex_models() -> Catalog:
    """Use the authenticated Codex endpoint; catalogs are status-only fallback."""
    token = ""
    try:
        from hermes_cli.auth import resolve_codex_runtime_credentials

        credentials = resolve_codex_runtime_credentials(refresh_if_expiring=True) or {}
        token = str(credentials.get("api_key") or "").strip()
    except Exception:
        token = ""
    live: list[str] = []
    if token:
        try:
            from hermes_cli.codex_models import _fetch_models_from_api

            live = _unique_models(_fetch_models_from_api(token))
        except Exception:
            live = []
    if live:
        return Catalog("openai-codex", live, True, True, "connected")
    try:
        from hermes_cli.models import provider_model_ids

        known = _unique_models(provider_model_ids("openai-codex", force_refresh=True))
    except Exception:
        known = []
    # The owner agent deliberately uses the local Codex app-server and the
    # Codex CLI's ChatGPT login. That credential is not exported into Hermes'
    # bearer-token store, so API discovery above may be empty even though the
    # runtime is fully authenticated. Treat a successful CLI status check as a
    # connected, degraded catalog: routing may use Codex while retaining the
    # last-known working model names until app-server discovery is available.
    if codex_cli_authenticated():
        return Catalog(
            "openai-codex",
            known,
            False,
            True,
            "connected-cli",
        )
    state = "degraded" if token else "not-configured"
    return Catalog("openai-codex", known, False, bool(token), state)


def discover_claude_models() -> Catalog:
    """Use the native subscription login before any API-key catalog fallback."""
    claude = shutil.which("claude") or os.environ.get("HERMES_CLAUDE_BIN", "")
    if claude:
        try:
            environment = os.environ.copy()
            environment.pop("ANTHROPIC_API_KEY", None)
            environment.pop("CLAUDE_API_KEY", None)
            status = subprocess.run(
                [claude, "auth", "status"],
                capture_output=True,
                check=False,
                text=True,
                timeout=15,
                env=environment,
            )
            detail = json.loads(status.stdout) if status.stdout.strip() else {}
            if (
                status.returncode == 0
                and isinstance(detail, dict)
                and detail.get("loggedIn") is True
                and detail.get("apiProvider") == "firstParty"
            ):
                try:
                    from hermes_cli.models import provider_model_ids

                    known = _unique_models(
                        provider_model_ids("anthropic", force_refresh=True)
                    )
                except Exception:
                    known = []
                return Catalog(
                    "anthropic",
                    _unique_models((*CLAUDE_SUBSCRIPTION_MODELS, *known)),
                    True,
                    True,
                    "connected-subscription",
                )
        except (OSError, subprocess.SubprocessError, ValueError, json.JSONDecodeError):
            pass
    try:
        from hermes_cli.models import provider_model_ids

        known = _unique_models(provider_model_ids("anthropic", force_refresh=True))
    except Exception:
        known = []
    return Catalog("anthropic", known, False, False, "not-configured")


def _profile_config(
    base: dict[str, Any],
    primary_provider: str,
    primary_model: str,
    fallback: dict[str, str],
    effort: str,
) -> dict[str, Any]:
    """Derive a worker configuration with an explicit cross-provider fallback."""
    config = copy.deepcopy(base)
    config["model"] = {
        "provider": primary_provider,
        "default": primary_model,
        "model": primary_model,
        "openai_runtime": "codex_app_server",
    }
    config["fallback_providers"] = [fallback]
    # Local OpenAI-compatible runtimes top out at the equivalent of high.
    # Never silently downgrade an explicitly xhigh task after both hosted
    # provider lanes are exhausted.
    if effort != "xhigh":
        config["fallback_providers"].append(copy.deepcopy(ATLAS_FALLBACK))
    agent = config.setdefault("agent", {})
    if isinstance(agent, dict):
        agent["reasoning_effort"] = effort
    config["toolsets"] = []
    return config


def _switchyard_profile_config(
    base: dict[str, Any], route: str, effort: str
) -> dict[str, Any]:
    """Derive a profile that cannot bypass the Switchyard authority."""
    config = copy.deepcopy(base)
    providers = config.setdefault("providers", {})
    if not isinstance(providers, dict):
        providers = {}
        config["providers"] = providers
    providers[SWITCHYARD_PROVIDER] = {
        "name": "Atlas Switchyard",
        "api": SWITCHYARD_API,
        "api_key": "atlas-switchyard",
        "default_model": route,
        "transport": "chat_completions",
    }
    config["model"] = {
        "provider": SWITCHYARD_PROVIDER,
        "default": route,
        "model": route,
    }
    config["fallback_providers"] = []
    config["model_catalog"] = {"enabled": True, "ttl_hours": 1}
    agent = config.setdefault("agent", {})
    if isinstance(agent, dict):
        agent["reasoning_effort"] = effort
    config["toolsets"] = []
    return config


def _write_profile(
    root: Path,
    name: str,
    description: str,
    soul: str,
    config: dict[str, Any],
    env_values: dict[str, str],
) -> None:
    """Create or refresh a managed Hermes worker profile."""
    profile = root / "profiles" / name
    for directory in ("logs", "sessions", "skills", "workspace", "home"):
        (profile / directory).mkdir(parents=True, exist_ok=True)
    _write_yaml(profile / "config.yaml", config)
    _write_yaml(
        profile / "profile.yaml",
        {"description": description, "description_auto": False},
    )
    _atomic_write(profile / "SOUL.md", soul.rstrip() + "\n")
    _update_profile_env(profile / ".env", env_values)


def configure_routes(
    root: Path,
    codex: Catalog,
    claude: Catalog,
    catalog_path: Path | None = None,
) -> dict[str, Any]:
    """Refresh catalogs while keeping every Hermes profile on Switchyard."""
    config_path = root / "config.yaml"
    base = _read_yaml(config_path)
    resolved_catalog_path = catalog_path or (
        Path(ROUTING_CATALOG_PATH)
        if ROUTING_CATALOG_PATH
        else root / "routing-catalog.json"
    )
    catalog = write_routing_catalog(resolved_catalog_path, codex, claude)
    providers = catalog["providers"]
    codex_models = providers["codex"]["resolved"]
    claude_models = providers["claude"]["resolved"]

    coordinator_toolsets = copy.deepcopy(base.get("toolsets"))
    base = _switchyard_profile_config(base, SWITCHYARD_AUTO_ROUTE, "high")
    # The coordinator uses its configured toolsets. Worker profiles below are
    # deliberately toolset-empty so Hermes resolves their native defaults.
    if coordinator_toolsets is None:
        base.pop("toolsets", None)
    else:
        base["toolsets"] = coordinator_toolsets
    _write_yaml(config_path, base)

    env_values = _read_env(root / ".env")
    routes: dict[str, list[str]] = {}
    for effort in EFFORTS:
        codex_name = f"codex-{effort}"
        claude_name = f"claude-{effort}"
        codex_route = {
            "low": "atlas/manual/codex/luna",
            "medium": "atlas/manual/codex/terra",
            "high": "atlas/manual/codex/sol",
            "xhigh": "atlas/manual/codex/sol",
        }[effort]
        claude_route = {
            "low": "atlas/manual/claude/haiku",
            "medium": "atlas/manual/claude/sonnet",
            "high": "atlas/manual/claude/sonnet",
            "xhigh": "atlas/manual/claude/opus",
        }[effort]
        _write_profile(
            root,
            codex_name,
            f"Codex implementation preference at {effort} effort, enforced by Switchyard.",
            "You are an implementation worker. Make focused, tested changes for the assigned task, preserve unrelated work, and report evidence and blockers to the coordinator.",
            _switchyard_profile_config(base, codex_route, effort),
            env_values,
        )
        _write_profile(
            root,
            claude_name,
            f"Claude analysis preference at {effort} effort, enforced by Switchyard.",
            "You are an architecture and review worker. Analyze the assigned task deeply, change files only when asked, and return concise conclusions, evidence, and risks to the coordinator.",
            _switchyard_profile_config(base, claude_route, effort),
            env_values,
        )
        routes[codex_name] = [codex_route]
        routes[claude_name] = [claude_route]

    _write_profile(
        root,
        "synthesis-xhigh",
        "Cross-provider synthesis and critical review, capped at xhigh effort.",
        "Synthesize the worker evidence into one answer. Resolve disagreements explicitly, verify high-risk claims, and never claim completion without cited validation.",
        _switchyard_profile_config(base, SWITCHYARD_AUTO_ROUTE, "xhigh"),
        env_values,
    )
    routes["synthesis-xhigh"] = [SWITCHYARD_AUTO_ROUTE]
    _write_yaml(
        root / "profile.yaml",
        {
            "description": "Owner-only project coordinator using Switchyard to route Hermes, Codex, Claude, and local model boundaries.",
            "description_auto": False,
        },
    )
    routes["coordinator"] = [SWITCHYARD_AUTO_ROUTE]
    routes["catalog"] = [
        *(f"openai-codex/{model}" for model in codex_models.values()),
        *(f"anthropic/{model}" for model in claude_models.values()),
    ]
    return routes