scripts/hf_deploy.py

"""Deploy llm-service to Hugging Face Spaces.

Stages the Docker Space files, sets API key secrets from local secrets.toml,
and uploads everything via the huggingface_hub API.

Usage:
    python -m scripts.hf_deploy <space_id>

Example:
    python -m scripts.hf_deploy ronheichman/llm-service
"""
from __future__ import annotations

import os
import shutil
import sys
import tempfile
import tomllib
from datetime import date
from pathlib import Path

from huggingface_hub import HfApi

REPO_ROOT = Path(__file__).resolve().parent.parent

STAGE_FILES = {
    "Dockerfile.hf": "Dockerfile",
    "hf_readme.md": "README.md",
    "hf_entrypoint.sh": "hf_entrypoint.sh",
    "free_llm_providers.yaml": "free_llm_providers.yaml",
    "secrets.example.toml": "secrets.example.toml",
}

STAGE_DIRS = {
    "scripts": "scripts",
}


def load_secrets(secrets_path: Path) -> dict[str, dict]:
    """Parse secrets.toml and return the keys section."""
    return tomllib.loads(secrets_path.read_text()).get("keys", {})


def _as_date(value: object) -> date | None:
    if value is None:
        return None
    return date.fromisoformat(str(value)) if isinstance(value, str) else value


def provider_is_active(config: dict) -> bool:
    """Mirror generate_config's expires/disabled_until checks.

    hf_inject_secrets.py rebuilds secrets.toml on the Space from env vars and
    cannot carry disabled_until/expires metadata. Disabled providers must
    therefore be excluded at deploy time, and their stale Space secrets
    deleted, or dead providers resurrect on the Space (this is how scitely and
    iflow kept serving 4xx/5xx from the live proxy after being disabled
    locally).
    """
    today = date.today()
    expires = _as_date(config.get("expires"))
    if expires is not None and expires < today:
        return False
    disabled_until = _as_date(config.get("disabled_until"))
    return not (disabled_until is not None and disabled_until > today)


def secrets_to_env_vars(keys: dict[str, dict]) -> dict[str, str]:
    """Map active secrets.toml entries to HF Space secret env var names."""
    env_vars: dict[str, str] = {}
    for provider, config in keys.items():
        if not provider_is_active(config):
            continue
        upper = provider.upper()
        for i, key in enumerate(config.get("keys", [])):
            if i == 0:
                env_vars[f"LLM_KEY_{upper}"] = key
            else:
                env_vars[f"LLM_KEY_{upper}_{i + 1}"] = key

        account_id = config.get("account_id")
        if account_id:
            env_vars[f"LLM_ACCOUNT_ID_{upper}"] = str(account_id)

    return env_vars


def stale_env_var_names(keys: dict[str, dict]) -> list[str]:
    """Space secret names to delete for providers that are not active."""
    names: list[str] = []
    for provider, config in keys.items():
        if provider_is_active(config):
            continue
        upper = provider.upper()
        names.append(f"LLM_KEY_{upper}")
        names.extend(f"LLM_KEY_{upper}_{i}" for i in range(2, 2 + len(config.get("keys", []))))
        if config.get("account_id"):
            names.append(f"LLM_ACCOUNT_ID_{upper}")
    return names


def stage_files(staging_dir: Path) -> None:
    """Copy all deployment files into the staging directory."""
    for src_name, dst_name in STAGE_FILES.items():
        src = REPO_ROOT / src_name
        if not src.exists():
            print(f"WARNING: {src} not found, skipping")
            continue
        shutil.copy2(src, staging_dir / dst_name)

    for src_name, dst_name in STAGE_DIRS.items():
        src = REPO_ROOT / src_name
        if not src.exists():
            print(f"WARNING: {src}/ not found, skipping")
            continue
        shutil.copytree(src, staging_dir / dst_name, dirs_exist_ok=True)


SERVICE_SECRETS = {
    "STORE_MODEL_IN_DB": "True",
}


def deploy(
    space_id: str,
    secrets_path: Path | None,
    database_url: str | None,
    master_key: str | None,
) -> str:
    """Deploy to HF Spaces. Returns the Space URL."""
    token = os.environ.get("HF_TOKEN")
    if token is None and secrets_path and secrets_path.exists():
        hf_entry = load_secrets(secrets_path).get("huggingface", {})
        hf_keys = hf_entry.get("keys", [])
        if hf_keys:
            token = hf_keys[0]
            print("Using huggingface key from secrets.toml as deploy token")
    api = HfApi(token=token)

    api.create_repo(repo_id=space_id, repo_type="space", space_sdk="docker", exist_ok=True)

    for name, value in SERVICE_SECRETS.items():
        api.add_space_secret(repo_id=space_id, key=name, value=value)

    if database_url:
        api.add_space_secret(repo_id=space_id, key="DATABASE_URL", value=database_url)
        print("Set DATABASE_URL Space secret")

    if master_key:
        api.add_space_secret(repo_id=space_id, key="LITELLM_MASTER_KEY", value=master_key)
        print("Set LITELLM_MASTER_KEY Space secret")

    if secrets_path and secrets_path.exists():
        keys = load_secrets(secrets_path)
        env_vars = secrets_to_env_vars(keys)
        for name, value in env_vars.items():
            api.add_space_secret(repo_id=space_id, key=name, value=value)
        print(f"Set {len(env_vars)} LLM key secret(s)")

        stale = stale_env_var_names(keys)
        for name in stale:
            api.delete_space_secret(repo_id=space_id, key=name)
        if stale:
            print(f"Deleted {len(stale)} stale secret(s) for disabled/expired providers: "
                  + ", ".join(stale))
    else:
        print("No secrets.toml found, skipping LLM key injection")

    with tempfile.TemporaryDirectory() as tmp:
        staging = Path(tmp)
        stage_files(staging)
        api.upload_folder(
            repo_id=space_id,
            repo_type="space",
            folder_path=str(staging),
        )

    url = f"https://huggingface.co/spaces/{space_id}"
    print(f"Deployed to {url}")
    return url


def main() -> None:
    if len(sys.argv) < 2:
        print("Usage: python -m scripts.hf_deploy <space_id> [--database-url URL] [--master-key KEY]")
        print("Example: python -m scripts.hf_deploy ronheichman/llm-service --database-url postgresql://...")
        sys.exit(1)

    space_id = sys.argv[1]
    secrets_path = REPO_ROOT / "secrets.toml"

    db_url = None
    mk = None
    args = sys.argv[2:]
    i = 0
    while i < len(args):
        if args[i] == "--database-url" and i + 1 < len(args):
            db_url = args[i + 1]
            i += 2
        elif args[i] == "--master-key" and i + 1 < len(args):
            mk = args[i + 1]
            i += 2
        else:
            i += 1

    deploy(space_id, secrets_path, db_url, mk)


if __name__ == "__main__":
    main()