File size: 6,791 Bytes
6e55d2d b598cb3 6e55d2d b598cb3 6e55d2d b598cb3 6e55d2d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | #!/usr/bin/env python3
"""Check a deployed knowledge Space: the gate fails closed, the surface matches the descriptor.
C9 of the Phase 1 plan, run against the live Space rather than the code. ``stage: RUNNING`` only
says the container came up; this signs in the way the orchestrator does and asks real questions.
Two tokens, both the caller's HF token, as the orchestrator sends them:
* ``token=`` — Space *access*. It gets a caller past HF's own wall on a private Space; a public
Space ignores it. The app never reads it.
* ``headers["x-orchestrator-token"]`` — the machine *identity* the allow-list checks. HF's edge
strips ``x-hf-*``, so only this custom header reaches the gate.
Checks, in order (each prints one line; payloads are never dumped):
1. no identity header → ``denied`` / ``no_token_header`` (with ``--anonymous``: no token at all,
for a public Space, which is what a stranger sends)
2. a token HF cannot resolve → ``denied`` / ``token_unresolved``
3. ``view_api()`` named endpoints and their params == ``deploy/orchestrator_registration.yaml``
4. the caller's own identity → ``ok`` on ``annotate_genes`` for ELAVL1
5. with ``--lookups``: one real lookup per endpoint → ``ok`` (keyless public sources, no model call)
Exit codes: 0 all passed, 1 a check failed, 2 could not reach the Space, 3 no token.
python scripts/check_live_surface.py --space anne-voigt/pdac-knowledge-agent-dev --lookups
python scripts/check_live_surface.py --space anne-voigt/pdac-knowledge-agent --anonymous
"""
from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
import yaml
ROOT = Path(__file__).resolve().parent.parent
DESCRIPTOR = ROOT / "deploy" / "orchestrator_registration.yaml"
DEV_SPACE = "anne-voigt/pdac-knowledge-agent-dev"
# Shaped like an HF token so the gate asks HF about it, and HF says no.
_UNRESOLVABLE_TOKEN = "hf_" + "0" * 34
LOOKUPS = {
"drugs_for_targets": {"genes": "KRAS"},
"trials_for_genes": {"genes": "CD40", "condition": "pancreatic cancer", "status": "RECRUITING"},
"annotate_genes": {"genes": "ELAVL1"},
"pathways_for_genes": {"genes": "ISG15"},
"targets_for_disease": {"disease": "MONDO_0005184", "size": "5"},
}
def _token() -> str | None:
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
if token:
return token
try:
from huggingface_hub import get_token
return get_token()
except Exception: # noqa: BLE001 — an old hub, or no cache; the caller reports it
return None
def _call(client, endpoint: str, **kwargs) -> dict:
return json.loads(client.predict(**kwargs, api_name=f"/{endpoint}"))
def _describe(payload: dict) -> str:
"""The cause in a few named fields, never the rows."""
if payload.get("status") == "ok":
return (
f"status='ok' rows={len(payload.get('rows', []))} "
f"unresolved={len(payload.get('unresolved', []))} "
f"source_errors={[e.get('reason') for e in payload.get('source_errors', [])]}"
)
keys = ("status", "machine_auth", "reason", "message")
return " ".join(f"{k}={payload[k]!r}" for k in keys if k in payload)[:300]
def _expect_denied(client, why: str, label: str) -> bool:
payload = _call(client, "annotate_genes", genes="ELAVL1")
passed = payload.get("status") == "denied" and payload.get("machine_auth") == why
print(f"{'PASS' if passed else 'FAIL'} {label}: {_describe(payload)}")
return passed
def _surface_matches(client) -> bool:
descriptor = yaml.safe_load(DESCRIPTOR.read_text())
expected = {t["api_name"]: list(t["params"]) for t in descriptor["tools"]}
named = client.view_api(return_format="dict", print_info=False)["named_endpoints"]
# Gradio's own endpoints start with "/_" (hf_oauth adds /_check_login_status); listed, not compared.
internal = sorted(name for name in named if name.startswith("/_"))
live = {
name: [p.get("parameter_name") or p.get("label") for p in spec.get("parameters", [])]
for name, spec in named.items()
if name not in internal
}
passed = live == expected
print(
f"{'PASS' if passed else 'FAIL'} view_api == descriptor: {sorted(live)}"
f" (gradio internals: {internal})"
)
if not passed:
for name in sorted(set(live) | set(expected)):
if live.get(name) != expected.get(name):
print(f" {name}: live={live.get(name)} descriptor={expected.get(name)}")
return passed
def main() -> int:
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
ap.add_argument("--space", default=DEV_SPACE, help=f"Space id (default {DEV_SPACE})")
ap.add_argument(
"--anonymous", action="store_true", help="check 1 sends no token at all (public Space)"
)
ap.add_argument("--lookups", action="store_true", help="one real lookup per endpoint")
args = ap.parse_args()
token = _token()
if not token:
print("check: no HF token (set HF_TOKEN or `hf auth login`)", file=sys.stderr)
return 3
from gradio_client import Client
try:
no_identity = Client(args.space, token=None if args.anonymous else token, verbose=False)
unresolvable = Client(
args.space,
token=token,
headers={"x-orchestrator-token": _UNRESOLVABLE_TOKEN},
verbose=False,
)
caller = Client(
args.space, token=token, headers={"x-orchestrator-token": token}, verbose=False
)
except Exception as exc: # noqa: BLE001 — any failure here means "could not reach it"
print(f"check: could not open {args.space}: {type(exc).__name__}: {exc}", file=sys.stderr)
return 2
print(f"check: {args.space}")
results = [
_expect_denied(
no_identity,
"no_token_header",
"no token at all" if args.anonymous else "no identity header",
),
_expect_denied(unresolvable, "token_unresolved", "unresolvable token"),
_surface_matches(caller),
]
authorized = _call(caller, "annotate_genes", genes="ELAVL1")
results.append(authorized.get("status") == "ok")
print(f"{'PASS' if results[-1] else 'FAIL'} own identity: {_describe(authorized)}")
if args.lookups:
for endpoint, kwargs in LOOKUPS.items():
payload = _call(caller, endpoint, **kwargs)
results.append(payload.get("status") == "ok" and payload.get("kind") == "annotation")
print(f"{'PASS' if results[-1] else 'FAIL'} {endpoint}: {_describe(payload)}")
return 0 if all(results) else 1
if __name__ == "__main__":
sys.exit(main())
|