AgapHost/openai/agent_registry.py

#!/usr/bin/env python3
"""agent_registry — reads agent-registry.yaml (kb#134, A2A-2).

Sibling of model_registry.py (kb#133): same load/get/CLI shape, same
philosophy ("registry is data, not logic"). This module's one piece of
real logic is effective_card() — the mechanism that makes "switching a
backbone is one field" true rather than aspirational (see the header
comment in agent-registry.yaml): an agent's static fields (persona, trust
class, tool scope, memory banks, preferred tier) never duplicate what the
CURRENT backbone provides (tier, cost_class, availability a(t), context
window). Those are resolved at read time by dereferencing `backbone` into
model-registry.yaml (via model_registry.get_model) or, for runtimes
model-registry.yaml deliberately excludes (kb#133: Claude Code), into this
file's own `runtimes:` section.

Usage (library):
    from agent_registry import load_registry, get_agent, effective_card, trust_rank
    reg = load_registry()
    adolf = get_agent(reg, "adolf")
    card = effective_card(reg, "adolf")   # -> persona/trust/tools/memory + resolved tier/cost/a(t)
    trust_rank(reg, "torgash") < trust_rank(reg, "adolf")   # sandboxed < trusted

Usage (CLI, for manual verification):
    ./agent_registry.py list
    ./agent_registry.py get --id adolf
    ./agent_registry.py effective-card --id adolf
    ./agent_registry.py trust-rank --id torgash
    ./agent_registry.py can-reach-vault --id torgash   # kb#147 sanity check
"""

import argparse
import json
import os
import sys

import yaml

import model_registry as mr

HERE = os.path.dirname(os.path.abspath(__file__))
DEFAULT_REGISTRY_PATH = os.path.join(HERE, "agent-registry.yaml")


class AgentRegistryError(Exception):
    pass


def load_registry(path=None):
    """Load and lightly validate agent-registry.yaml."""
    path = path or DEFAULT_REGISTRY_PATH
    with open(path) as f:
        reg = yaml.safe_load(f)
    if not reg or "agents" not in reg:
        raise AgentRegistryError(f"{path}: missing top-level 'agents' list")
    ids = [a["id"] for a in reg["agents"]]
    if len(ids) != len(set(ids)):
        raise AgentRegistryError(f"{path}: duplicate agent ids in {ids}")
    runtime_ids = [r["id"] for r in reg.get("runtimes", [])]
    if len(runtime_ids) != len(set(runtime_ids)):
        raise AgentRegistryError(f"{path}: duplicate runtime ids in {runtime_ids}")
    return reg


def get_agent(registry, agent_id):
    for a in registry["agents"]:
        if a["id"] == agent_id:
            return a
    raise AgentRegistryError(f"unknown agent id: {agent_id!r}")


def get_runtime(registry, runtime_id):
    for r in registry.get("runtimes", []):
        if r["id"] == runtime_id:
            return r
    raise AgentRegistryError(f"unknown runtime id: {runtime_id!r}")


# ---------------------------------------------------------------------------
# trust_rank — §5 "human > trusted > sandboxed > untrusted" as a comparable
# integer, for kb#140 (routing) / kb#147 (grant enforcement) to use directly
# instead of re-deriving an ordering from the string.
# ---------------------------------------------------------------------------

def trust_rank(registry, agent_id):
    a = get_agent(registry, agent_id)
    classes = registry.get("trust_classes", {})
    cls = a["trust_class"]
    if cls not in classes:
        raise AgentRegistryError(f"{agent_id}: unknown trust_class {cls!r} (have: {sorted(classes)})")
    return classes[cls]["rank"]


def can_reach_vault(registry, agent_id):
    """kb#147 sanity check: vault access = trusted-or-human only (DECIDED).
    Cross-checks the agent's declared tool_allowlist.vault_access flag
    against its trust rank, so a registry typo (vault_access: true on a
    sandboxed agent) is a raised error, not a silent enforcement gap."""
    a = get_agent(registry, agent_id)
    classes = registry.get("trust_classes", {})
    trusted_rank = classes.get("trusted", {}).get("rank")
    declared = ((a.get("tool_allowlist") or {}).get("vault_access")) if a.get("tool_allowlist") else False
    rank = trust_rank(registry, agent_id)
    if declared and rank < trusted_rank:
        raise AgentRegistryError(
            f"{agent_id}: tool_allowlist.vault_access=true but trust_class="
            f"{a['trust_class']!r} (rank {rank}) < trusted (rank {trusted_rank}) — "
            "registry inconsistency, fix before this is load-bearing for kb#147"
        )
    return bool(declared) and rank >= trusted_rank


# ---------------------------------------------------------------------------
# effective_card — the full A2A Agent Card: static agent fields merged with
# the CURRENT backbone's resolved tier/cost_class/a(t)/context_tokens.
# ---------------------------------------------------------------------------

def _resolve_backbone(registry, backbone_id, model_registry=None):
    """backbone can point into either this file's runtimes: (flat-subscription
    runtimes model-registry.yaml deliberately excludes, kb#133) or into
    model-registry.yaml's models: (everything else). Try runtimes first —
    it's the smaller, local list."""
    try:
        return "runtime", get_runtime(registry, backbone_id)
    except AgentRegistryError:
        pass
    model_registry = model_registry if model_registry is not None else mr.load_registry()
    return "model", mr.get_model(model_registry, backbone_id)


def effective_card(registry, agent_id, model_registry=None):
    """Return the full A2A Agent Card for `agent_id`: its own persona/trust/
    tools/memory fields plus tier/cost_class/context_tokens/lifecycle
    resolved from whatever `backbone` currently names. Humans and trivial
    endpoints with backbone=None get resolved_backbone=None — there's
    nothing to dereference."""
    a = get_agent(registry, agent_id)
    card = dict(a)  # shallow copy; don't mutate the loaded registry
    backbone_id = a.get("backbone")
    if backbone_id is None:
        card["resolved_backbone"] = None
        return card

    source, backbone = _resolve_backbone(registry, backbone_id, model_registry)
    card["resolved_backbone"] = {
        "source": source,               # "runtime" | "model" — which file it came from
        "id": backbone_id,
        "tier": backbone["tier"],
        "cost_class": backbone["cost_class"],
        "lifecycle": backbone["lifecycle"],
        "context_tokens": backbone.get("context_tokens"),
        "metered": backbone.get("metered", False),
    }
    return card


# ---------------------------------------------------------------------------
# litellm_key_spec — kb#147 (A2A-15): turn an agent's static registry fields
# into the LiteLLM virtual-key grant provision_litellm_keys.py provisions.
# "Grants live in the agent registry, not scattered configs" (kb#147 accept-
# ance bar) means the model allow-list and budget are COMPUTED here from
# preferred_tier + trust_class, never hand-typed per agent.
# ---------------------------------------------------------------------------

# Ascending order matching model-registry.yaml's routing.tiers keys. An
# agent may use its preferred tier and anything below it (a "large"-
# preferring agent degrades to "small" gracefully; a "small"-only agent
# never gets "large" — that asymmetry IS the sandboxed/trusted split this
# key spec exists to enforce).
_TIER_ORDER = ["small", "large"]


def _reachable_tiers(preferred_tier):
    if preferred_tier not in _TIER_ORDER:
        return []
    return _TIER_ORDER[: _TIER_ORDER.index(preferred_tier) + 1]


def litellm_key_spec(registry, agent_id, model_registry=None):
    """Return the LiteLLM virtual-key grant for `agent_id`: which
    litellm_model_name values it may use and its default budget, derived
    from THIS registry's data (preferred_tier, trust_class) plus
    model-registry.yaml's routing.tiers/metered_opt_in — never hand-entered
    per agent. Models with no litellm_model_name (e.g. `kimi`, called
    directly via the adolf-llm wrapper, never through LiteLLM) are outside
    LiteLLM's enforcement surface by construction and are excluded, not
    silently allowed.

    provision_litellm_keys.py consumes this dict's `models`/`max_budget`/
    `budget_duration`/`key_alias` as the body of a LiteLLM /key/generate (or
    /key/update) call. This function makes no network call itself.
    """
    a = get_agent(registry, agent_id)
    model_registry = model_registry if model_registry is not None else mr.load_registry()
    grant = a.get("capability_grant") or {}
    key_alias = grant.get("litellm_key_alias", agent_id)

    opted_in = set(model_registry.get("routing", {}).get("metered_opt_in", []) or [])
    opted_in_key = f"agent:{agent_id}"
    pools = model_registry.get("routing", {}).get("tiers", {})

    models = []
    for tier in _reachable_tiers(a.get("preferred_tier")):
        for model_id in pools.get(tier, []):
            m = mr.get_model(model_registry, model_id)
            name = m.get("litellm_model_name")
            if not name:
                continue  # not LiteLLM-routed (e.g. kimi's adolf-llm wrapper) -- nothing to grant/deny here
            if m.get("metered") and opted_in_key not in opted_in:
                continue  # §3a: no metered API by default, per-key opt-in only
            if name not in models:
                models.append(name)

    # kb#128 gap (flagged 2026-07-26, closed 2026-07-30): the raw litellm_
    # model_names above (e.g. "ollama/gemma3:4b") are the BACKING deployments
    # for openai/litellm-config.yaml's alias model_names -- tier-small/
    # tier-large (alvis's "tier" routing mode) and auto_router/
    # complexity_router (alvis's "automatic" routing mode). Without granting
    # the aliases too, a provisioned key could reach a model directly but not
    # by tier or through the router, so "all three routing modes exercisable"
    # (kb#128 acceptance) wasn't actually true per-agent. Gate exactly like
    # the raw grants above -- reachable tiers, not a separate allow-list --
    # so an agent's routing-mode access never exceeds its direct-model access:
    #   - "small" reachable  -> tier-small (mirrors the always-granted small
    #     pool; every agent with a backbone gets at least this).
    #   - "large" reachable  -> tier-large, PLUS auto_router/complexity_router.
    #     Both routers' pools include tier-large in their upper bands (COMPLEX/
    #     REASONING, or the semantic "complex reasoning" route), so granting
    #     them to a small-only (sandboxed) agent would let automatic routing
    #     escalate it past its trust class -- exactly the asymmetry
    #     _reachable_tiers()/kb#147 exists to prevent. A small-only agent gets
    #     neither router: it can still call tier-small directly.
    reachable = _reachable_tiers(a.get("preferred_tier"))
    if "small" in reachable and "tier-small" not in models:
        models.append("tier-small")
    if "large" in reachable:
        for alias in ("tier-large", "auto_router", "complexity_router"):
            if alias not in models:
                models.append(alias)

    classes = registry.get("trust_classes", {})
    cls = classes.get(a["trust_class"], {})
    return {
        "agent_id": agent_id,
        "key_alias": key_alias,
        "trust_class": a["trust_class"],
        "models": models,
        "max_budget": cls.get("default_budget_usd"),
        "budget_duration": cls.get("budget_duration"),
        "mcp_auth_token_env": grant.get("mcp_auth_token_env"),
    }


# ---------------------------------------------------------------------------
# CLI — manual verification only, not part of the library contract.
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--registry", default=None, help="path to agent-registry.yaml (default: sibling file)")
    ap.add_argument("--model-registry", default=None, help="path to model-registry.yaml (default: sibling file)")
    sub = ap.add_subparsers(dest="cmd", required=True)

    p = sub.add_parser("get")
    p.add_argument("--id", required=True)

    p = sub.add_parser("effective-card")
    p.add_argument("--id", required=True)

    p = sub.add_parser("trust-rank")
    p.add_argument("--id", required=True)

    p = sub.add_parser("can-reach-vault")
    p.add_argument("--id", required=True)

    p = sub.add_parser("litellm-key-spec")
    p.add_argument("--id", required=True)

    sub.add_parser("list")

    args = ap.parse_args()
    reg = load_registry(args.registry)
    model_reg = mr.load_registry(args.model_registry) if args.model_registry else None

    try:
        if args.cmd == "get":
            print(json.dumps(get_agent(reg, args.id), indent=2))
        elif args.cmd == "effective-card":
            print(json.dumps(effective_card(reg, args.id, model_reg), indent=2))
        elif args.cmd == "trust-rank":
            print(trust_rank(reg, args.id))
        elif args.cmd == "can-reach-vault":
            ok = can_reach_vault(reg, args.id)
            print(json.dumps({"id": args.id, "can_reach_vault": ok}))
            sys.exit(0 if ok else 1)
        elif args.cmd == "litellm-key-spec":
            print(json.dumps(litellm_key_spec(reg, args.id, model_reg), indent=2))
        elif args.cmd == "list":
            for a in reg["agents"]:
                backbone = a.get("backbone") or "-"
                print(f"{a['id']:22} trust={a['trust_class']:9} "
                      f"tier={str(a.get('preferred_tier')):6} backbone={backbone:16} "
                      f"role={a['persona']['role']}")
    except AgentRegistryError as exc:
        print(f"error: {exc}", file=sys.stderr)
        sys.exit(2)


if __name__ == "__main__":
    main()