opencode-patches/discover_models.py

#!/usr/bin/env python3
"""
Auto-discover models from 9Router and update opencode.json.

Usage:
  python3 discover_models.py [--config PATH] [--output PATH] [--dry-run] [--push]

Scans /v1/models on the 9router endpoint, filters to text-capable models,
and generates the provider.models section for opencode.json.

Environment variables:
  NINEROUTER_BASE_URL  - 9Router base URL (default: https://ollama.wilddragon.net/v1)
  NINEROUTER_API_KEY   - 9Router API key
  FORGEJO_TOKEN        - Forgejo API token (for --push)
"""

import json
import sys
import os
import argparse
import urllib.request
import urllib.error
import re

# 9Router endpoint config
BASE_URL = os.environ.get("NINEROUTER_BASE_URL", "https://ollama.wilddragon.net/v1")
API_KEY = os.environ.get("NINEROUTER_API_KEY", "")
CONFIG_PATH = os.path.expanduser("~/.config/opencode/opencode.json")

# Forgejo
FORGEJO_URL = "https://forge.wilddragon.net"
FORGEJO_OWNER = "zgaetano"
FORGEJO_REPO = "opencode-patches"
FORGEJO_TOKEN = os.environ.get("FORGEJO_TOKEN", "")

# Model IDs to skip (non-text, ASR, image, etc.)
SKIP_PATTERNS = [
    "asr", "tts", "image", "embed", "vision-preview-only",
    "parakeet", "whisper", "speech", "audio",
    "-review",
]

# Provider prefixes in 9router model IDs -> display names
PROVIDER_NAMES = {
    "cc": "CC",
    "openrouter": "OpenRouter",
    "nvidia": "Nvidia",
    "ollama": "Ollama Cloud",
    "ollama-local": "Ollama Local",
    "gemini": "Google Gemini",
    "ds": "DeepSeek",
    "cx": "Codex",
    "mistral": "Mistral",
    "jinx": "Jinx",
    "ag": "Antigravity",
    "cl": "Cloudliner",
    "kr": "KR Cloud",
    "cf": "Cloudflare",
    "g0i": "G0I Cloud",
    "oc": "OC Cloud",
}

# Known model name mappings for clean display
MODEL_NAME_MAP = {
    "claude-opus-4-7": "Claude Opus 4.7",
    "claude-opus-4-6": "Claude Opus 4.6",
    "claude-opus-4-5-20251101": "Claude Opus 4.5",
    "claude-opus-4-8": "Claude Opus 4.8",
    "claude-sonnet-4-6": "Claude Sonnet 4.6",
    "claude-sonnet-4-5-20250929": "Claude Sonnet 4.5",
    "claude-sonnet-4-5": "Claude Sonnet 4.5",
    "claude-sonnet-4": "Claude Sonnet 4",
    "claude-haiku-4-5-20251001": "Claude Haiku 4.5",
    "claude-haiku-4-5": "Claude Haiku 4.5",
    "claude-opus-4-6-thinking": "Claude Opus 4.6 Thinking",
    "claude-opus-4.7": "Claude Opus 4.7",
    "claude-sonnet-4.6": "Claude Sonnet 4.6",
    "claude-opus-4.6": "Claude Opus 4.6",
    "gpt-5.5": "GPT-5.5",
    "gpt-5.4": "GPT-5.4",
    "gpt-5.3-codex": "GPT-5.3 Codex",
    "gpt-oss-120b": "GPT OSS 120B",
    "gpt-oss-120b-medium": "GPT OSS 120B Medium",
    "gpt-oss:120b": "GPT OSS 120B",
    "gpt-image-2": "GPT Image 2",
    "gemini-3.1-pro-preview": "Gemini 3.1 Pro Preview",
    "gemini-3.1-flash-lite-preview": "Gemini 3.1 Flash Lite Preview",
    "gemini-3-flash-preview": "Gemini 3 Flash Preview",
    "gemini-2.0-flash-lite": "Gemini 2.0 Flash Lite",
    "gemini-2.5-pro": "Gemini 2.5 Pro",
    "gemini-2.5-flash": "Gemini 2.5 Flash",
    "gemini-2.5-flash-lite": "Gemini 2.5 Flash Lite",
    "gemini-3.1-pro-high": "Gemini 3.1 Pro High",
    "gemini-3.1-pro-low": "Gemini 3.1 Pro Low",
    "gemini-3-flash": "Gemini 3 Flash",
    "deepseek-v4-pro": "DeepSeek V4 Pro",
    "deepseek-v4-pro-max": "DeepSeek V4 Pro Max",
    "deepseek-v4-pro-none": "DeepSeek V4 Pro None",
    "deepseek-v4-flash": "DeepSeek V4 Flash",
    "deepseek-chat": "DeepSeek Chat",
    "deepseek-reasoner": "DeepSeek Reasoner",
    "deepseek-3.2": "DeepSeek 3.2",
    "deepseek-r1-14b": "DeepSeek R1 14B",
    "deepseek-r1-distill-qwen-32b": "DeepSeek R1 Distill Qwen 32B",
    "qwen3.5": "Qwen 3.5",
    "qwen3-coder-80b": "Qwen 3 Coder 80B",
    "qwen3-coder-next": "Qwen 3 Coder Next",
    "qwen3-coder-next-thinking": "Qwen 3 Coder Next (Thinking)",
    "qwen3-coder-next-agentic": "Qwen 3 Coder Next (Agentic)",
    "qwen3-coder-next-thinking-agentic": "Qwen 3 Coder Next (Thinking+Agentic)",
    "qwen2.5-coder-32b-instruct": "Qwen 2.5 Coder 32B",
    "qwen2.5-coder:14b": "Qwen 2.5 Coder 14B",
    "qwq-32b": "QwQ 32B",
    "glm-5": "GLM 5",
    "glm-5.1": "GLM 5.1",
    "glm-4.7-flash": "GLM 4.7 Flash",
    "kimi-k2.5": "Kimi K2.5",
    "kimi-k2.6": "Kimi K2.6",
    "mistral-large-latest": "Mistral Large",
    "mistral-medium-latest": "Mistral Medium",
    "mistral-small-3.1-24b-instruct": "Mistral Small 3.1 24B",
    "codestral-latest": "Codestral",
    "llama-3.2-1b-instruct": "Llama 3.2 1B",
    "llama-3.2-3b-instruct": "Llama 3.2 3B",
    "llama-3.1-8b-instruct-fp8-fast": "Llama 3.1 8B (FP8 Fast)",
    "llama-3.1-8b-instruct-awq": "Llama 3.1 8B (AWQ)",
    "llama-3.1-70b-instruct-fp8-fast": "Llama 3.1 70B (FP8 Fast)",
    "llama-3.3-70b-instruct-fp8-fast": "Llama 3.3 70B (FP8 Fast)",
    "nemotron-3-nano-30b-a3b:free": "Nemotron 3 Nano 30B (Free)",
    "nemotron-3-nano-omni-30b-a3b-reasoning:free": "Nemotron 3 Nano Omni Reasoning (Free)",
    "nemotron-3-super-120b-a12b:free": "Nemotron 3 Super 120B (Free)",
    "trinity-large-thinking:free": "Trinity Large Thinking (Free)",
    "gemma-4-26b-a4b-it:free": "Gemma 4 26B IT (Free)",
    "gemma-4-31b-it": "Gemma 4 31B IT",
    "kiro-claude-sonnet-4-5": "Kiro Claude Sonnet 4.5",
    "kiro-claude-sonnet-4-5-agentic": "Kiro Claude Sonnet 4.5 (Agentic)",
    "kiro-claude-sonnet-4": "Kiro Claude Sonnet 4",
    "kiro-claude-haiku-4-5": "Kiro Claude Haiku 4.5",
    "kiro-claude-haiku-4-5-agentic": "Kiro Claude Haiku 4.5 (Agentic)",
    "kiro-auto": "Kiro Auto",
    "kiro-glm-5": "Kiro GLM 5",
    "kiro-glm-5-agentic": "Kiro GLM 5 (Agentic)",
    "kiro-minimax-m2-1": "Kiro MiniMax M2.1",
    "kiro-minimax-m2-1-agentic": "Kiro MiniMax M2.1 (Agentic)",
    "kiro-minimax-m2-5": "Kiro MiniMax M2.5",
    "kiro-minimax-m2-5-agentic": "Kiro MiniMax M2.5 (Agentic)",
    "kiro-deepseek-3-2": "Kiro DeepSeek 3.2",
    "kiro-deepseek-3-2-agentic": "Kiro DeepSeek 3.2 (Agentic)",
    "kiro-qwen3-coder-next-agentic": "Kiro Qwen 3 Coder Next (Agentic)",
    "minimax-m2.5": "MiniMax M2.5",
    "minimax-m2.1": "MiniMax M2.1",
    "minimax-m2.5-thinking": "MiniMax M2.5 (Thinking)",
    "minimax-m2.5-agentic": "MiniMax M2.5 (Agentic)",
    "minimax-m2.5-thinking-agentic": "MiniMax M2.5 (Thinking+Agentic)",
    "minimax-m2.1-thinking": "MiniMax M2.1 (Thinking)",
    "minimax-m2.1-agentic": "MiniMax M2.1 (Agentic)",
    "minimax-m2.1-thinking-agentic": "MiniMax M2.1 (Thinking+Agentic)",
    "claudecode": "ClaudeCode",
    "kat-coder-pro": "Kat Coder Pro",
    "owl-alpha": "Owl Alpha",
}


def friendly_name(model_id: str) -> str:
    """Turn 'cc/claude-sonnet-4-5-20250929' into 'Claude Sonnet 4.5'."""
    raw = model_id.rsplit("/", 1)[-1]
    clean = re.sub(r"-\d{8}$", "", raw)
    if clean in MODEL_NAME_MAP:
        return MODEL_NAME_MAP[clean]
    if raw in MODEL_NAME_MAP:
        return MODEL_NAME_MAP[raw]
    name = clean.replace("-", " ").replace("_", " ")
    return " ".join(w.capitalize() for w in name.split())


def provider_label(model_id: str) -> str:
    """Get display label for the provider prefix."""
    prefix = model_id.split("/")[0] if "/" in model_id else ""
    return PROVIDER_NAMES.get(prefix, prefix.upper())


def should_skip(model_id: str) -> bool:
    """Skip non-coding models."""
    lower = model_id.lower()
    return any(p in lower for p in SKIP_PATTERNS)


def fetch_models(base_url: str, api_key: str) -> list:
    """Fetch model list from 9Router."""
    url = f"{base_url}/models"
    req = urllib.request.Request(url, headers={"Authorization": f"Bearer {api_key}"})
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            data = json.loads(resp.read().decode())
            return data.get("data", [])
    except urllib.error.URLError as e:
        print(f"ERROR: Failed to fetch models: {e}", file=sys.stderr)
        sys.exit(1)


def build_model_entry(model_id: str) -> dict:
    """Build an opencode model entry."""
    label = friendly_name(model_id)
    provider = provider_label(model_id)
    return {"name": f"{provider} - {label}"}


def discover(base_url: str, api_key: str) -> dict:
    """Main discovery logic. Returns {model_id: entry} dict."""
    raw_models = fetch_models(base_url, api_key)
    result = {}
    for m in raw_models:
        mid = m.get("id", "")
        if not mid or mid == "Main" or should_skip(mid):
            continue
        result[mid] = build_model_entry(mid)
    return result


def load_config(path: str) -> dict:
    try:
        with open(path) as f:
            return json.load(f)
    except FileNotFoundError:
        return {"$schema": "https://opencode.ai/config.json", "provider": {}}
    except json.JSONDecodeError as e:
        print(f"ERROR: Invalid JSON in {path}: {e}", file=sys.stderr)
        sys.exit(1)


def save_config(path: str, config: dict):
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w") as f:
        json.dump(config, f, indent=2)
        f.write("\n")


def push_to_forgejo(token: str, models: dict):
    """Push discovered models to Forgejo as JSON."""
    if not token:
        print("WARN: No Forgejo token, skipping push", file=sys.stderr)
        return
    data = json.dumps(models, indent=2).encode()
    import base64
    content = base64.b64encode(data).decode()
    url = f"{FORGEJO_URL}/api/v1/repos/{FORGEJO_OWNER}/{FORGEJO_REPO}/contents/discovered_models.json"
    # Check if file exists first
    req = urllib.request.Request(url, headers={"Authorization": f"token {token}"})
    sha = None
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            existing = json.loads(resp.read().decode())
            sha = existing.get("sha")
    except urllib.error.HTTPError as e:
        if e.code != 404:
            print(f"WARN: Forgejo check failed: {e}", file=sys.stderr)
            return
    payload = json.dumps({
        "message": f"Auto-discover: {len(models)} models",
        "content": content,
        **({"sha": sha} if sha else {}),
    }).encode()
    req = urllib.request.Request(url, data=payload, headers={
        "Authorization": f"token {token}",
        "Content-Type": "application/json",
    })
    req.get_method = lambda: "PUT"
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            print(f"Pushed to Forgejo: {resp.status}", file=sys.stderr)
    except urllib.error.URLError as e:
        print(f"WARN: Forgejo push failed: {e}", file=sys.stderr)


def main():
    parser = argparse.ArgumentParser(description="Auto-discover 9Router models for opencode")
    parser.add_argument("--config", default=CONFIG_PATH, help="Path to opencode.json")
    parser.add_argument("--output", help="Output path (default: same as --config)")
    parser.add_argument("--dry-run", action="store_true", help="Print result, don't write")
    parser.add_argument("--push", action="store_true", help="Push discovered models to Forgejo")
    parser.add_argument("--api-key", default=API_KEY, help="9Router API key")
    parser.add_argument("--base-url", default=BASE_URL, help="9Router base URL")
    args = parser.parse_args()

    if not args.api_key:
        print("ERROR: No API key. Set NINEROUTER_API_KEY or use --api-key", file=sys.stderr)
        sys.exit(1)

    discovered = discover(args.base_url, args.api_key)
    print(f"Discovered {len(discovered)} models", file=sys.stderr)

    if args.dry_run:
        for mid, entry in discovered.items():
            print(f"  {mid} -> {entry['name']}")
        return

    output_path = args.output or args.config
    config = load_config(args.config)
    if "provider" not in config:
        config["provider"] = {}
    if "9router" not in config["provider"]:
        config["provider"]["9router"] = {
            "name": "9router",
            "npm": "@ai-sdk/openai-compatible",
            "options": {
                "baseURL": args.base_url,
                "apiKey": args.api_key,
            },
            "models": {},
        }

    config["provider"]["9router"]["models"] = discovered
    save_config(output_path, config)
    print(f"Wrote {len(discovered)} models to {output_path}", file=sys.stderr)

    if args.push:
        push_to_forgejo(FORGEJO_TOKEN, discovered)


if __name__ == "__main__":
    main()