gen_key.py

#!/usr/bin/env python3
"""Generate an API key for the LLM serving nginx auth and update .env."""

import secrets
import string
import sys
import os

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))


def generate_key(prefix: str = "sk-lm", part1_len: int = 10, part2_len: int = 20) -> str:
    """Generate key like: sk-lm-XXXXXXXXXX:YYYYYYYYYYYYYYYYYYYY"""
    chars = string.ascii_letters + string.digits
    part1 = "".join(secrets.choice(chars) for _ in range(part1_len))
    part2 = "".join(secrets.choice(chars) for _ in range(part2_len))
    return f"{prefix}-{part1}:{part2}"


def update_env(key: str) -> None:
    """Update DPI_FACTORY_API_KEY in .env file."""
    env_path = os.path.join(SCRIPT_DIR, ".env")

    if os.path.exists(env_path):
        with open(env_path) as f:
            lines = f.readlines()
        with open(env_path, "w") as f:
            for line in lines:
                if line.startswith("DPI_FACTORY_API_KEY="):
                    f.write(f"DPI_FACTORY_API_KEY={key}\n")
                else:
                    f.write(line)
    else:
        with open(env_path, "w") as f:
            f.write(f"DPI_FACTORY_API_KEY={key}\n")
            f.write("DPI_FACTORY_API_BASE=http://125.188.35.185:19101/v1\n")


def main() -> None:
    key = generate_key()
    print(f"Generated API key:\n  {key}")

    if "--no-save" in sys.argv:
        print("\nDry run — .env not updated.")
        return

    update_env(key)
    print(f"\n✓ Saved to .env")
    print(f"  Run ./reload_nginx.sh to apply the new key.")


if __name__ == "__main__":
    main()