llama-swap-sync.sh

#!/usr/bin/env bash
# llama-swap-sync.sh
# Syncs models from llama-swap config.yaml → opencode.json
#
# Usage:
#   ./llama-swap-sync.sh               # run sync
#   ./llama-swap-sync.sh --reconfigure # reset saved settings
#   ./llama-swap-sync.sh --help
#
# Dependencies: python3 (stdlib only — no jq needed)

set -euo pipefail

# ── Config ────────────────────────────────────────────────────────────────────
STATE_FILE="${XDG_CONFIG_HOME:-$HOME/.config}/llama-swap-sync/settings.conf"

# ── Colours ───────────────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'

info()    { echo -e "${CYAN}[INFO]${NC}  $*"; }
success() { echo -e "${GREEN}[OK]${NC}    $*"; }
warn()    { echo -e "${YELLOW}[WARN]${NC}  $*"; }
error()   { echo -e "${RED}[ERROR]${NC} $*" >&2; }
die()     { error "$*"; exit 1; }

# ── Dependency check ──────────────────────────────────────────────────────────
command -v python3 &>/dev/null || die "Required command not found: python3"

# ── First-run setup / reconfigure ────────────────────────────────────────────
# Pass "reconfigure" as $1 to pre-populate defaults from existing settings file
first_run_setup() {
  local mode="${1:-setup}"

  echo ""
  echo -e "${BOLD}╔══════════════════════════════════════════════╗${NC}"
  echo -e "${BOLD}║      llama-swap → opencode.json sync         ║${NC}"

  if [[ "$mode" == "reconfigure" ]]; then
    echo -e "${BOLD}║              Reconfiguring                   ║${NC}"
    echo -e "${BOLD}╚══════════════════════════════════════════════╝${NC}"
    echo ""
    warn "Settings file found: $STATE_FILE"
    info "Current values are shown as defaults — press Enter to keep them."
    echo ""
    # Load existing values as defaults
    # shellcheck source=/dev/null
    source "$STATE_FILE"
  else
    echo -e "${BOLD}║            First-time setup                  ║${NC}"
    echo -e "${BOLD}╚══════════════════════════════════════════════╝${NC}"
    echo ""
    info "No saved settings found. Let's get you configured."
    echo ""
  fi

  # Use existing values (if reconfiguring) or hard-coded defaults (if fresh setup)
  local default_opencode="${OPENCODE_JSON:-$HOME/.config/opencode/opencode.json}"
  read -rp "$(echo -e "Path to ${BOLD}opencode.json${NC} [${default_opencode}]: ")" input
  OPENCODE_JSON="${input:-$default_opencode}"

  local default_yaml="${LLAMA_SWAP_YAML:-$HOME/llama-swap/config.yaml}"
  read -rp "$(echo -e "Path to llama-swap ${BOLD}config.yaml${NC} [${default_yaml}]: ")" input
  LLAMA_SWAP_YAML="${input:-$default_yaml}"

  local default_url="${BASE_URL:-http://127.0.0.1:41234/v1}"
  read -rp "$(echo -e "llama-swap ${BOLD}baseURL${NC} [${default_url}]: ")" input
  BASE_URL="${input:-$default_url}"

  local default_pid="${PROVIDER_ID:-llama-swap}"
  read -rp "$(echo -e "Provider ${BOLD}ID${NC} in opencode.json [${default_pid}]: ")" input
  PROVIDER_ID="${input:-$default_pid}"

  local default_pname="${PROVIDER_NAME:-llama-swap}"
  read -rp "$(echo -e "Provider ${BOLD}display name${NC} [${default_pname}]: ")" input
  PROVIDER_NAME="${input:-$default_pname}"

  local default_output="${DEFAULT_OUTPUT:-8192}"
  read -rp "$(echo -e "Default ${BOLD}max output tokens${NC} [${default_output}]: ")" input
  DEFAULT_OUTPUT="${input:-$default_output}"

  echo ""
  mkdir -p "$(dirname "$STATE_FILE")"
  cat > "$STATE_FILE" <<EOF
OPENCODE_JSON=${OPENCODE_JSON}
LLAMA_SWAP_YAML=${LLAMA_SWAP_YAML}
BASE_URL=${BASE_URL}
PROVIDER_ID=${PROVIDER_ID}
PROVIDER_NAME=${PROVIDER_NAME}
DEFAULT_OUTPUT=${DEFAULT_OUTPUT}
EOF
  success "Settings saved to $STATE_FILE"
  echo ""
}

# ── Load saved settings ───────────────────────────────────────────────────────
load_settings() {
  if [[ ! -f "$STATE_FILE" ]]; then
    first_run_setup setup
  else
    # shellcheck source=/dev/null
    source "$STATE_FILE"
  fi
  [[ -f "$LLAMA_SWAP_YAML" ]] || die "llama-swap config not found: $LLAMA_SWAP_YAML"
  [[ -f "$OPENCODE_JSON"   ]] || die "opencode.json not found: $OPENCODE_JSON"
}

# ── Parse config.yaml → JSON array of {id, name, context} ────────────────────
parse_models() {
  python3 - "$LLAMA_SWAP_YAML" <<'PYEOF'
import sys, re, json

yaml_file = sys.argv[1]
with open(yaml_file) as f:
    content = f.read()

models = {}
in_models = False
current_model = None
current_cmd_lines = []
in_cmd_block = False

for raw_line in content.splitlines():
    line = raw_line.rstrip()
    stripped = line.strip()
    indent = len(line) - len(line.lstrip())

    if re.match(r'^models\s*:', line):
        in_models = True
        current_model = None
        continue

    if not in_models:
        continue

    if indent == 0 and stripped and not re.match(r'^models\s*:', line):
        in_models = False
        continue

    if indent == 2 and stripped and not stripped.startswith('#'):
        m = re.match(r'^  ["\']?(.+?)["\']?\s*:\s*$', line)
        if m:
            if current_model and current_cmd_lines:
                models[current_model] = ' '.join(current_cmd_lines)
            current_model = m.group(1).strip()
            current_cmd_lines = []
            in_cmd_block = False
            continue

    if current_model is None:
        continue

    if re.match(r'^\s{4}cmd\s*:\s*\|\s*$', line):
        current_cmd_lines = []
        in_cmd_block = True
        continue

    if re.match(r'^\s{4}cmd\s*:\s*\S', line) and not in_cmd_block:
        cmd_val = re.sub(r'^\s+cmd\s*:\s*', '', line).strip()
        current_cmd_lines = [cmd_val]
        continue

    if in_cmd_block:
        if indent >= 6:
            current_cmd_lines.append(stripped.rstrip('\\').strip())
        elif stripped and not stripped.startswith('#'):
            in_cmd_block = False

if current_model and current_cmd_lines:
    models[current_model] = ' '.join(current_cmd_lines)

results = []
for model_id, cmd in models.items():
    ctx_match = re.search(r'(?:-c|--ctx-size)\s+(\d+)', cmd)
    context = int(ctx_match.group(1)) if ctx_match else None
    results.append({"id": model_id, "name": model_id, "context": context})

print(json.dumps(results))
PYEOF
}

# ── Backup opencode.json ──────────────────────────────────────────────────────
backup_opencode() {
  local backup="${OPENCODE_JSON}.bak.$(date +%Y%m%d_%H%M%S)"
  cp "$OPENCODE_JSON" "$backup"
  info "Backup: $backup"
}

# ── Main sync ─────────────────────────────────────────────────────────────────
sync() {
  info "Reading $LLAMA_SWAP_YAML ..."
  local models_json
  models_json="$(parse_models)"

  # All JSON read / diff / write handled by python3 — no jq needed
  local py_output
  py_output="$(python3 - "$models_json" "$OPENCODE_JSON" "$PROVIDER_ID" "$PROVIDER_NAME" "$BASE_URL" "$DEFAULT_OUTPUT" <<'PYEOF'
import sys, json

models_json    = sys.argv[1]
opencode_path  = sys.argv[2]
provider_id    = sys.argv[3]
provider_name  = sys.argv[4]
base_url       = sys.argv[5]
default_output = int(sys.argv[6])

models = json.loads(models_json)

if not models:
    print("ERROR: No models parsed from config.yaml — check the file format.", file=sys.stderr)
    sys.exit(1)

# Load existing opencode.json — supports JSONC (strips // comments and trailing commas)
def load_jsonc(path):
    import re
    with open(path) as f:
        text = f.read()
    # Remove single-line // comments (not inside strings)
    text = re.sub(r'(?<!:)//(?=(?:[^"]*"[^"]*")*[^"]*$).*', '', text)
    # Remove trailing commas before } or ]
    text = re.sub(r',\s*([}\]])', r'\1', text)
    return json.loads(text)

config = load_jsonc(opencode_path)

# Read existing models so we can preserve manually set output values
existing = config.get("provider", {}).get(provider_id, {}).get("models", {})

# Build new models dict — preserve existing output if already set, else use default
new_models = {}
for m in models:
    entry = {"name": m["name"]}
    if m["context"] is not None:
        # Preserve existing output value if the model already exists in opencode.json
        existing_output = existing.get(m["id"], {}).get("limit", {}).get("output")
        entry["limit"] = {
            "context": m["context"],
            "output":  existing_output if existing_output is not None else default_output,
        }
    new_models[m["id"]] = entry

# Calculate diff vs existing provider models
added   = sorted(set(new_models) - set(existing))
removed = sorted(set(existing)   - set(new_models))

# Emit diff markers for the shell to display
print(f"ADDED:{','.join(added)}")
print(f"REMOVED:{','.join(removed)}")

# Write updated provider block — all other keys preserved
config.setdefault("provider", {})[provider_id] = {
    "npm":     "@ai-sdk/openai-compatible",
    "name":    provider_name,
    "options": {"baseURL": base_url},
    "models":  new_models,
}

with open(opencode_path, "w") as f:
    json.dump(config, f, indent=2)
    f.write("\n")

# Emit model summary lines for the shell to display
for m in models:
    ctx = f"ctx {m['context']}" if m["context"] is not None else "no -c flag"
    print(f"MODEL:{m['id']}\t{m['name']}\t{ctx}")
PYEOF
)"

  # Parse and display diff
  local added removed
  added="$(echo   "$py_output" | grep '^ADDED:'   | cut -d: -f2)"
  removed="$(echo "$py_output" | grep '^REMOVED:' | cut -d: -f2)"

  if [[ -n "$added" ]];   then echo -e "${GREEN}  + Adding:${NC}   ${added//,/ }"; fi
  if [[ -n "$removed" ]]; then echo -e "${RED}  - Removing:${NC} ${removed//,/ }"; fi
  if [[ -z "$added" && -z "$removed" ]]; then
    info "Model list unchanged — updating context sizes if needed"
  fi

  backup_opencode
  success "opencode.json updated."
  echo ""
  echo -e "${BOLD}  Active models:${NC}"
  echo "$py_output" | grep '^MODEL:' | sed 's/^MODEL://' | \
    awk -F'\t' '{ printf "  • %-30s ← %s  [%s]\n", $1, $2, $3 }'
  echo ""
}

# ── Entry point ───────────────────────────────────────────────────────────────
case "${1:-}" in
  --reconfigure|-r)
    first_run_setup reconfigure
    sync
    ;;
  --help|-h)
    echo "Usage: $(basename "$0") [--reconfigure | --help]"
    echo ""
    echo "  (no args)        Sync llama-swap config.yaml → opencode.json"
    echo "  --reconfigure    Reset saved settings and re-run setup"
    echo "  --help           Show this message"
    echo ""
    echo "Settings are stored in: $STATE_FILE"
    exit 0
    ;;
  "")
    load_settings
    sync
    ;;
  *)
    die "Unknown argument: $1 (try --help)"
    ;;
esac