scripts/admin.py

"""Admin CLI for managing the live llm-service deployment.

Usage:
    python -m scripts.admin status
    python -m scripts.admin update-key <provider> <new-key>
    python -m scripts.admin list-models
    python -m scripts.admin add-models --provider <name> --key <api-key>
    python -m scripts.admin remove-model <model-id>
    python -m scripts.admin generate-consumer-key --alias <name>
    python -m scripts.admin revoke-consumer-key <key>
"""

from __future__ import annotations

import json
import sys
import tomllib
from pathlib import Path

import click
import httpx
import tomlkit
import yaml


SECRETS_PATH = Path(__file__).resolve().parent.parent / "secrets.toml"
PROVIDERS_PATH = Path(__file__).resolve().parent.parent / "free_llm_providers.yaml"


def load_service_config() -> tuple[str, str]:
    """Return (proxy_url, master_key) from secrets.toml."""
    doc = tomllib.loads(SECRETS_PATH.read_text())
    master_key = doc.get("service", {}).get("master_key", "")
    proxy_url = doc.get("service", {}).get("proxy_url", "https://ronheichman-llm-service.hf.space")
    return proxy_url, master_key


def api_get(path: str, proxy_url: str, master_key: str) -> dict:
    with httpx.Client(timeout=120) as client:
        resp = client.get(f"{proxy_url}{path}", headers={"Authorization": f"Bearer {master_key}"})
        resp.raise_for_status()
        return resp.json()


def api_post(path: str, payload: dict, proxy_url: str, master_key: str) -> dict:
    with httpx.Client(timeout=120) as client:
        resp = client.post(
            f"{proxy_url}{path}",
            json=payload,
            headers={"Authorization": f"Bearer {master_key}"},
        )
        resp.raise_for_status()
        return resp.json()


@click.group()
def cli() -> None:
    """Manage the live llm-service deployment."""


@cli.command()
def status() -> None:
    """Show deployment health and model counts."""
    proxy_url, master_key = load_service_config()
    health = api_get("/health", proxy_url, master_key)
    info = api_get("/model/info", proxy_url, master_key)

    models = info.get("data", [])
    providers: dict[str, int] = {}
    db_count = 0
    yaml_count = 0
    for m in models:
        mid = m.get("model_info", {}).get("id", "?")
        provider = mid.split("::")[0]
        providers[provider] = providers.get(provider, 0) + 1
        if m.get("model_info", {}).get("db_model", False):
            db_count += 1
        else:
            yaml_count += 1

    click.echo(f"Health: {json.dumps(health, indent=2)[:200]}")
    click.echo(f"\n{len(models)} deployments ({yaml_count} from config, {db_count} from DB)")
    for p in sorted(providers):
        click.echo(f"  {p}: {providers[p]}")


@cli.command("list-models")
def list_models() -> None:
    """List all model deployments with provider, model, and source."""
    proxy_url, master_key = load_service_config()
    info = api_get("/model/info", proxy_url, master_key)

    for m in info.get("data", []):
        mid = m.get("model_info", {}).get("id", "?")
        db = m.get("model_info", {}).get("db_model", False)
        source = "db" if db else "yaml"
        click.echo(f"  [{source}] {mid}")


@cli.command("update-key")
@click.argument("provider")
@click.argument("new_key")
def update_key(provider: str, new_key: str) -> None:
    """Update a provider's API key on the live service.

    1. Updates secrets.toml locally
    2. Regenerates litellm config to get the deployment list
    3. Adds all models for that provider via /model/new
    """
    proxy_url, master_key = load_service_config()

    doc = tomlkit.loads(SECRETS_PATH.read_text())
    keys_section = doc.get("keys", {})
    if provider not in keys_section:
        click.echo(f"Provider '{provider}' not found in secrets.toml", err=True)
        sys.exit(1)

    keys_section[provider]["keys"] = [new_key]
    if "disabled_until" in keys_section[provider]:
        del keys_section[provider]["disabled_until"]
    SECRETS_PATH.write_text(tomlkit.dumps(doc))
    click.echo(f"Updated secrets.toml for {provider}")

    from scripts.generate_config import generate_config

    tmp_config = Path("/tmp/admin_litellm_config.yaml")
    generate_config(PROVIDERS_PATH, SECRETS_PATH, tmp_config)
    config = yaml.safe_load(tmp_config.read_text())

    provider_models = [
        m for m in config["model_list"]
        if m["model_info"]["id"].startswith(f"{provider}::")
    ]

    if not provider_models:
        click.echo(f"No models generated for {provider} (filtered out?)", err=True)
        sys.exit(1)

    added = 0
    for m in provider_models:
        m["litellm_params"]["api_key"] = new_key
        payload = {
            "model_name": m["model_name"],
            "litellm_params": m["litellm_params"],
            "model_info": m["model_info"],
        }
        resp = api_post("/model/new", payload, proxy_url, master_key)
        click.echo(f"  Added {m['model_info']['id']}")
        added += 1

    click.echo(f"\n{added} deployments added to live service for {provider}")
    click.echo("Note: old YAML-loaded deployments for this provider may still exist until next restart.")


@cli.command("add-models")
@click.option("--provider", required=True)
@click.option("--key", required=True, help="API key for the provider")
def add_models(provider: str, key: str) -> None:
    """Add all models for a provider to the live service."""
    proxy_url, master_key = load_service_config()

    from scripts.generate_config import generate_config

    tmp_config = Path("/tmp/admin_litellm_config.yaml")

    doc = tomlkit.loads(SECRETS_PATH.read_text())
    keys_section = doc.get("keys", {})
    if provider not in keys_section:
        click.echo(f"Provider '{provider}' not in secrets.toml. Adding temporarily for config generation.")
        keys_section[provider] = tomlkit.table()
        keys_section[provider].add("keys", [key])
        tmp_secrets = Path("/tmp/admin_secrets.toml")
        tmp_secrets.write_text(tomlkit.dumps(doc))
        generate_config(PROVIDERS_PATH, tmp_secrets, tmp_config)
    else:
        generate_config(PROVIDERS_PATH, SECRETS_PATH, tmp_config)

    config = yaml.safe_load(tmp_config.read_text())
    provider_models = [
        m for m in config["model_list"]
        if m["model_info"]["id"].startswith(f"{provider}::")
    ]

    for m in provider_models:
        m["litellm_params"]["api_key"] = key
        payload = {
            "model_name": m["model_name"],
            "litellm_params": m["litellm_params"],
            "model_info": m["model_info"],
        }
        api_post("/model/new", payload, proxy_url, master_key)
        click.echo(f"  Added {m['model_info']['id']}")

    click.echo(f"\n{len(provider_models)} deployments added")


@cli.command("remove-model")
@click.argument("model_id")
def remove_model(model_id: str) -> None:
    """Remove a model deployment by its ID (only works for DB-backed models)."""
    proxy_url, master_key = load_service_config()
    resp = api_post("/model/delete", {"id": model_id}, proxy_url, master_key)
    click.echo(f"Removed: {model_id}")


@cli.command("generate-consumer-key")
@click.option("--alias", required=True, help="Human-readable name for this key")
@click.option("--max-budget", type=float, help="Spending cap (dollars)")
@click.option("--duration", help="Key lifetime, e.g. '30d', '24h'")
def generate_consumer_key(alias: str, max_budget: float | None, duration: str | None) -> None:
    """Generate a virtual API key for a consumer."""
    proxy_url, master_key = load_service_config()
    payload: dict = {
        "key_alias": alias,
    }
    if max_budget is not None:
        payload["max_budget"] = max_budget
    if duration is not None:
        payload["duration"] = duration

    resp = api_post("/key/generate", payload, proxy_url, master_key)
    key = resp.get("key", "?")
    click.echo(f"Consumer key for '{alias}': {key}")
    click.echo("Give this key to the consumer. They use it instead of the master key.")


@cli.command("revoke-consumer-key")
@click.argument("key")
def revoke_consumer_key(key: str) -> None:
    """Revoke a consumer's virtual API key."""
    proxy_url, master_key = load_service_config()
    resp = api_post("/key/delete", {"keys": [key]}, proxy_url, master_key)
    click.echo(f"Revoked: {key}")


if __name__ == "__main__":
    cli()