UmbraAudio/scripts/lmstudio_control.py

#!/usr/bin/env python3
"""LM Studio model control script for UmbraAudio.

Usage:
    python scripts/lmstudio_control.py status       # Show LM Studio status
    python scripts/lmstudio_control.py start        # Start server
    python scripts/lmstudio_control.py stop         # Stop server
    python scripts/lmstudio_control.py load [MODEL] # Load a model
    python scripts/lmstudio_control.py unload       # Unload current model
"""
from __future__ import annotations

import sys

from umbraaudio.voices.lmstudio_lifecycle import (
    get_loaded_model,
    is_running,
    load_model,
    server_status,
    start_server,
    stop_server,
    unload_all,
)

PREFERRED_MODELS = [
    "gemma-4-12b-it-qat",
    "google/gemma-4-e4b",
    "qwen3.5-9b",
]


def cmd_start():
    if is_running():
        print("LM Studio is already running")
        return
    start_server()
    print("LM Studio server started")


def cmd_stop():
    if not is_running():
        print("LM Studio is not running")
        return
    stop_server()
    print("LM Studio server stopped")


def cmd_unload():
    if not is_running():
        print("LM Studio is not running")
        return
    loaded = get_loaded_model()
    if not loaded:
        print("No model loaded")
        return
    print(f"Unloading {loaded}...")
    if unload_all():
        print("Unloaded successfully")
    else:
        print("Unload failed")


def cmd_load(model: str | None = None):
    if not is_running():
        print("LM Studio is not running. Use 'start' first.")
        return

    current = get_loaded_model()
    if current:
        if current == model or (not model and current in PREFERRED_MODELS):
            print(f"Model already loaded: {current}")
            return
        print(f"Unloading {current}...")
        unload_all()

    if not model:
        import json
        import subprocess

        result = subprocess.run(
            ["lms", "ls", "--json"],
            capture_output=True, text=True, timeout=30,
        )
        if result.returncode == 0:
            available = [
                m["modelKey"] for m in json.loads(result.stdout)
                if m.get("type") == "llm"
            ]
            for preferred in PREFERRED_MODELS:
                if preferred in available:
                    model = preferred
                    break
        if not model:
            print("No preferred model found and no model specified")
            return

    print(f"Loading {model}...")
    if load_model(model):
        print(f"Loaded successfully: {model}")
    else:
        print(f"Load failed: {model}")


def cmd_status():
    status = server_status()
    if not status["api_responding"]:
        print("LM Studio is NOT running")
        print("\nStart it with:")
        print("  python scripts/lmstudio_control.py start")
        return

    print("LM Studio is running")
    print(f"  Server: {'running' if status['server_running'] else 'stopped'}")

    models = status["loaded_models"]
    if models:
        print(f"  Loaded: {models[0]}")
    else:
        print("  No model loaded")

    import json
    import subprocess

    result = subprocess.run(
        ["lms", "ls", "--json"],
        capture_output=True, text=True, timeout=30,
    )
    if result.returncode == 0:
        available = [
            m["modelKey"] for m in json.loads(result.stdout)
            if m.get("type") == "llm"
        ]
        if available:
            print(f"\nAvailable models ({len(available)}):")
            for m in available:
                marker = " *" if models and m == models[0] else ""
                preferred = " (recommended)" if m in PREFERRED_MODELS else ""
                print(f"  {m}{marker}{preferred}")


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        cmd_status()
        return

    cmd = sys.argv[1].lower()

    if cmd == "status":
        cmd_status()
    elif cmd == "start":
        cmd_start()
    elif cmd == "stop":
        cmd_stop()
    elif cmd == "unload":
        cmd_unload()
    elif cmd == "load":
        model = sys.argv[2] if len(sys.argv) > 2 else None
        cmd_load(model)
    else:
        print(f"Unknown command: {cmd}")
        print(__doc__)


if __name__ == "__main__":
    main()