automaton/scripts/detect_models.py

#!/usr/bin/env python3
"""Probe opencode.json and localhost endpoints to produce a candidate models.json.

Usage:
    python3 scripts/detect_models.py [--json] [--write]

Without --json, prints a human-readable report.
With --json, emits the candidate models.json to stdout as the last JSON line.
With --write, writes the candidate to ~/.automaton/models.json (idempotent,
never overwrites an existing file unless --force is also given).

Probing strategy (stdlib only):
  1. Parse opencode.json (or opencode.jsonc) for configured provider+model pairs.
  2. Probe localhost endpoints to find locally-running LLM servers:
     - http://localhost:8080/v1/models (llama.cpp / generic OpenAI-compatible)
     - http://localhost:11434/api/tags (Ollama)
     - http://localhost:1234/v1/models (LM Studio)
     - http://localhost:8000/v1/models (vLLM)
  3. Merge results into a candidate models.json.
"""

from __future__ import annotations

import json
import os
import re
import sys
from pathlib import Path
from typing import Optional

AUTOMATON_DIR = Path.home() / ".automaton"

# ---------------------------------------------------------------------------
# opencode.json parsing
# ---------------------------------------------------------------------------


def _find_opencode_json() -> Optional[Path]:
    """Locate the opencode config file (opencode.json or opencode.jsonc)."""
    candidates = [
        Path.cwd() / "opencode.json",
        Path.cwd() / "opencode.jsonc",
        AUTOMATON_DIR / "opencode.json",
        AUTOMATON_DIR / "opencode.jsonc",
        Path.home() / ".opencode.json",
        Path.home() / ".config" / "opencode" / "opencode.json",
        Path.home() / ".config" / "opencode" / "opencode.jsonc",
    ]
    for p in candidates:
        if p.exists():
            return p
    return None


def _parse_opencode_models(config_path: Path) -> list[dict]:
    """Extract model entries from an opencode.json config.

    Expected structure (common patterns):
      {
        "providers": {
          "opencode": { "model": "glm-4.6", ... },
          ...
        }
      }
    or a flatter:
      {
        "model": "glm-4.6",
        ...
      }
    """
    try:
        content = config_path.read_text(encoding="utf-8")
    except OSError:
        return []
    # Strip JSONC comments (// line comments only, sufficient for our use)
    content = re.sub(r"//.*", "", content)
    try:
        data = json.loads(content)
    except json.JSONDecodeError:
        return []
    if not isinstance(data, dict):
        return []
    models: list[dict] = []
    seen: set[str] = set()

    # Check top-level "model" field (single-model config)
    single = data.get("model")
    if isinstance(single, str) and single not in seen:
        seen.add(single)
        models.append({"name": single, "provider": "opencode", "context_window": None, "location": "remote"})

    # Check providers dict
    providers = data.get("providers") or {}
    for prov_name, prov_cfg in providers.items():
        if isinstance(prov_cfg, dict):
            model_name = prov_cfg.get("model")
            if isinstance(model_name, str) and model_name not in seen:
                seen.add(model_name)
                models.append({"name": model_name, "provider": prov_name, "context_window": None, "location": "remote"})

    # Check "models" list (explicit model roster)
    model_list = data.get("models")
    if isinstance(model_list, list):
        for entry in model_list:
            if isinstance(entry, dict):
                name = entry.get("name") or entry.get("model")
                if isinstance(name, str) and name not in seen:
                    seen.add(name)
                    models.append({
                        "name": name,
                        "provider": entry.get("provider", "opencode"),
                        "context_window": entry.get("context_window"),
                        "location": entry.get("location", "remote"),
                    })
    return models


# ---------------------------------------------------------------------------
# Localhost probing
# ---------------------------------------------------------------------------


def _fetch_json(url: str, timeout: int = 5) -> Optional[dict]:
    """Fetch a JSON response from a URL using urllib (stdlib)."""
    import urllib.request
    import urllib.error
    try:
        req = urllib.request.Request(url, method="GET")
        with urllib.request.urlopen(req, timeout=timeout) as resp:
            body = resp.read().decode("utf-8")
            return json.loads(body)
    except (OSError, urllib.error.URLError, json.JSONDecodeError, ValueError):
        return None


def _probe_ollama() -> list[dict]:
    """Probe Ollama: GET http://localhost:11434/api/tags → models[].name"""
    data = _fetch_json("http://localhost:11434/api/tags")
    if not data:
        return []
    models_list = data.get("models") or []
    return [
        {"name": m.get("name"), "provider": "ollama", "context_window": None, "location": "http://localhost:11434"}
        for m in models_list
        if isinstance(m, dict) and isinstance(m.get("name"), str)
    ]


def _probe_openai_compatible(url: str, provider: str) -> list[dict]:
    """Probe an OpenAI-compatible /v1/models endpoint."""
    data = _fetch_json(url)
    if not data:
        return []
    model_list = data.get("data") or []
    return [
        {"name": m.get("id"), "provider": provider, "context_window": None, "location": url}
        for m in model_list
        if isinstance(m, dict) and isinstance(m.get("id"), str)
    ]


_ENDPOINTS = [
    ("http://localhost:8080/v1/models", "llama.cpp"),
    ("http://localhost:11434/api/tags", "ollama"),  # handled separately above
    ("http://localhost:1234/v1/models", "lm-studio"),
    ("http://localhost:8000/v1/models", "vllm"),
]


def _probe_localhost() -> list[dict]:
    """Probe all known localhost endpoints and merge results."""
    seen_names: set[str] = set()
    models: list[dict] = []
    for url, provider in _ENDPOINTS:
        if provider == "ollama":
            result = _probe_ollama()
        else:
            result = _probe_openai_compatible(url, provider)
        for m in result:
            n = m.get("name")
            if isinstance(n, str) and n not in seen_names:
                seen_names.add(n)
                models.append(m)
    return models


# ---------------------------------------------------------------------------
# Merge & write
# ---------------------------------------------------------------------------


def build_candidate_models(probe_local: bool = True) -> dict:
    """Build a candidate models.json dict.

    1. Parse models from opencode.json
    2. Optionally probe localhost endpoints
    3. Merge: opencode config models come first; local probes fill in gaps.
    4. Build result with default, advised, models[].
    """
    opencode_path = _find_opencode_json()
    config_models: list[dict] = []
    if opencode_path:
        config_models = _parse_opencode_models(opencode_path)

    local_models: list[dict] = []
    if probe_local:
        local_models = _probe_localhost()

    # Merge: key by name, config models take priority (unordered)
    merged: dict[str, dict] = {}
    for m in config_models:
        n = m["name"]
        if n not in merged:
            merged[n] = m
    for m in local_models:
        n = m.get("name")
        if n and n not in merged:
            merged[n] = m

    models_list = list(merged.values())

    # Determine default: first config model, or first local model, or empty
    default_name: Optional[str] = None
    if config_models:
        default_name = config_models[0].get("name")
    elif local_models:
        default_name = local_models[0].get("name")

    # Determine advised: if only 0-1 models, set advised=true; else false
    advised = len(models_list) <= 1

    result: dict = {
        "schema_version": 1,
        "default": default_name,
        "advised": advised,
        "models": models_list,
    }

    return result


def write_models_file(candidate: dict, force: bool = False) -> bool:
    """Write candidate models.json to AUTOMATON_DIR.

    Never overwrites an existing file unless force=True.
    Returns True if written, False if skipped.
    """
    target = AUTOMATON_DIR / "models.json"
    if target.exists() and not force:
        return False
    target.write_text(json.dumps(candidate, indent=2) + "\n")
    return True


def format_report(candidate: dict) -> str:
    """Human-readable report of the candidate models."""
    lines = []
    lines.append("=== Model Detection Report ===")
    lines.append("")
    source = "No opencode.json found" if not _find_opencode_json() else f"Config: {_find_opencode_json()}"
    lines.append(f"Source: {source}")
    lines.append("")
    models = candidate.get("models", [])
    if not models:
        lines.append("No models detected.")
    else:
        lines.append(f"Detected {len(models)} model(s):")
        for m in models:
            loc = m.get("location", "unknown")
            prov = m.get("provider", "?")
            ctx = m.get("context_window")
            ctx_str = f", context: {ctx}" if ctx else ""
            lines.append(f"  - {m['name']} ({prov}, {loc}{ctx_str})")
    lines.append("")
    lines.append(f"Default: {candidate.get('default', 'none')}")
    lines.append(f"Advised: {candidate.get('advised', False)}")
    lines.append(f"Mode: {'multi-LLM' if len(models) >= 2 else 'single-LLM'}")
    lines.append("")
    target = AUTOMATON_DIR / "models.json"
    if target.exists():
        lines.append(f"models.json already exists at {target} (use --force to overwrite)")
    else:
        lines.append(f"Ready to write to {target} (use --write to create)")
    return "\n".join(lines)


def main() -> int:
    import argparse
    parser = argparse.ArgumentParser(description="Detect available LLM models and write models.json")
    parser.add_argument("--json", action="store_true", help="Output candidate JSON on last line")
    parser.add_argument("--write", action="store_true", help="Write candidate models.json to ~/.automaton/ (idempotent)")
    parser.add_argument("--force", action="store_true", help="Overwrite existing models.json")
    parser.add_argument("--no-probe", action="store_true", help="Skip localhost endpoint probing")
    args = parser.parse_args()

    candidate = build_candidate_models(probe_local=not args.no_probe)

    if args.write:
        written = write_models_file(candidate, force=args.force)
        if written:
            print(f"Written models.json to {AUTOMATON_DIR / 'models.json'}")
        else:
            print(f"Skipped: {AUTOMATON_DIR / 'models.json'} already exists (use --force to overwrite)")

    if args.json:
        print(json.dumps(candidate))
    else:
        print(format_report(candidate))

    return 0


if __name__ == "__main__":
    sys.exit(main())