AgapHost/ai/provision_litellm_keys.py

#!/usr/bin/env python3
"""provision_litellm_keys — kb#147 (A2A-15): turn agent-registry.yaml grants
into real LiteLLM virtual keys.

This is the ONE place a capability grant (model allow-list + budget) crosses
from data (agent-registry.yaml, version-controlled) into a live LiteLLM
key (via the proxy's /key/generate or /key/update admin API, master-key
authenticated). It deliberately does nothing destructive: --dry-run (the
default) only computes and prints the payload each agent WOULD get, making
zero network calls. --apply is required to actually create/update a key,
and needs LITELLM_MASTER_KEY in the environment (never hardcoded here, never
committed) — this is a privileged write against a live production service,
so it is not something this task runs unattended; --apply is the kb#147
handover step for a human/approved run.

Usage:
    # Safe, run-anytime: print what each agent's key WOULD look like.
    ./provision_litellm_keys.py --dry-run
    ./provision_litellm_keys.py --dry-run --id torgash

    # Privileged, requires explicit opt-in + master key (kb#147 handover):
    LITELLM_MASTER_KEY=sk-... ./provision_litellm_keys.py --apply --id adolf
"""

import argparse
import json
import os
import sys
import urllib.error
import urllib.request

import agent_registry as ar
import model_registry as mr

LITELLM_BASE_URL = os.environ.get("LITELLM_BASE_URL", "http://localhost:4000")


def agent_ids_with_grants(registry):
    return [a["id"] for a in registry["agents"] if a.get("capability_grant")]


def _http_post(path, payload, master_key):
    req = urllib.request.Request(
        f"{LITELLM_BASE_URL}{path}",
        data=json.dumps(payload).encode(),
        headers={
            "Authorization": f"Bearer {master_key}",
            "Content-Type": "application/json",
        },
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=10) as resp:
        return json.loads(resp.read())


def apply_key(spec, master_key):
    """Create (or update, if key_alias already exists) a LiteLLM virtual key
    matching `spec` (the dict returned by agent_registry.litellm_key_spec).
    Raises on any HTTP error rather than swallowing it — a failed grant
    should never look like a successful one."""
    payload = {
        "key_alias": spec["key_alias"],
        "models": spec["models"],
        "max_budget": spec["max_budget"],
        "budget_duration": spec["budget_duration"],
        "metadata": {"agent_id": spec["agent_id"], "trust_class": spec["trust_class"], "source": "kb#147 agent-registry.yaml"},
    }
    try:
        return _http_post("/key/generate", payload, master_key)
    except urllib.error.HTTPError as e:
        body = e.read().decode(errors="replace")
        raise SystemExit(f"LiteLLM /key/generate failed for {spec['key_alias']}: {e.code} {body}")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--registry", default=None)
    ap.add_argument("--model-registry", default=None)
    ap.add_argument("--id", default=None, help="only this agent id (default: every agent with a capability_grant)")
    mode = ap.add_mutually_exclusive_group()
    mode.add_argument("--dry-run", action="store_true", default=True, help="default: compute + print only, no network call")
    mode.add_argument("--apply", action="store_true", help="actually call LiteLLM /key/generate (needs LITELLM_MASTER_KEY) -- privileged, kb#147 handover step")
    args = ap.parse_args()

    reg = ar.load_registry(args.registry)
    model_reg = mr.load_registry(args.model_registry)

    ids = [args.id] if args.id else agent_ids_with_grants(reg)
    if not ids:
        print("no agents with a capability_grant in the registry", file=sys.stderr)
        sys.exit(1)

    master_key = os.environ.get("LITELLM_MASTER_KEY")
    if args.apply and not master_key:
        print("error: --apply requires LITELLM_MASTER_KEY in the environment", file=sys.stderr)
        sys.exit(2)

    for agent_id in ids:
        spec = ar.litellm_key_spec(reg, agent_id, model_reg)
        if args.apply:
            result = apply_key(spec, master_key)
            print(json.dumps({"agent_id": agent_id, "key_alias": spec["key_alias"], "applied": True, "litellm_response_keys": list(result.keys())}))
        else:
            print(json.dumps({"mode": "dry-run", **spec}, indent=2))


if __name__ == "__main__":
    main()