scripts/hf_inject_secrets.py

"""Build secrets.toml from environment variables for HF Spaces deployment.

Reads secrets.example.toml to discover known providers, then writes
secrets.toml entries for each provider that has a matching LLM_KEY_{PROVIDER}
env var set.

Env var conventions:
    LLM_KEY_GOOGLE          -> [keys.google] keys = ["value"]
    LLM_KEY_GOOGLE_2        -> appended as second key
    LLM_ACCOUNT_ID_GOOGLE   -> account_id = "value"
"""
from __future__ import annotations

import os
import tomllib
from pathlib import Path

import tomlkit


def discover_providers(example_path: Path) -> list[str]:
    """Return provider names from secrets.example.toml [keys.*] sections."""
    data = tomllib.loads(example_path.read_text())
    return list(data.get("keys", {}).keys())


def build_secrets_toml(example_path: Path) -> str:
    """Read env vars and produce secrets.toml content."""
    providers = discover_providers(example_path)
    doc = tomlkit.document()
    keys_table = tomlkit.table(is_super_table=True)

    for provider in providers:
        upper = provider.upper()
        primary = os.environ.get(f"LLM_KEY_{upper}")
        if primary is None:
            continue

        key_list = [primary]
        for i in range(2, 100):
            extra = os.environ.get(f"LLM_KEY_{upper}_{i}")
            if extra is None:
                break
            key_list.append(extra)

        entry = tomlkit.table()
        entry.add("keys", key_list)

        account_id = os.environ.get(f"LLM_ACCOUNT_ID_{upper}")
        if account_id is not None:
            entry.add("account_id", account_id)

        keys_table.add(provider, entry)

    doc.add("keys", keys_table)
    return tomlkit.dumps(doc)


def main() -> None:
    example_path = Path("secrets.example.toml")
    if not example_path.exists():
        raise FileNotFoundError(f"{example_path} not found in {Path.cwd()}")

    content = build_secrets_toml(example_path)
    Path("secrets.toml").write_text(content)
    count = content.count("[keys.")
    print(f"secrets.toml written with {count} provider(s)")


if __name__ == "__main__":
    main()