Spaces:
Paused
Paused
| """`flock` CLI: a thin client over the controller API. | |
| Config: FLOCK_URL and FLOCK_TOKEN. Inside an agent the worker sets FLOCK_AGENT_TOKEN (that | |
| agent's scoped token, which wins over FLOCK_TOKEN), FLOCK_AGENT and FLOCK_PARENT; see | |
| docs/spawn.md for what an agent may call. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import random | |
| import time | |
| from email.utils import parsedate_to_datetime | |
| from pathlib import Path | |
| from typing import Annotated, Any | |
| import httpx | |
| import typer | |
| import yaml | |
| from pydantic import ValidationError | |
| from flock.models import AGENT_NAME_RE, SwarmSpec | |
| app = typer.Typer(help="Host coding-harness agent swarms on HF Jobs.", no_args_is_help=True) | |
| def client() -> httpx.Client: | |
| """Overridden in tests to point at an in-process app.""" | |
| url = os.environ.get("FLOCK_URL", "http://localhost:7860") | |
| token = os.environ.get("FLOCK_AGENT_TOKEN") or os.environ.get("FLOCK_TOKEN", "") | |
| return httpx.Client(base_url=url, headers={"Authorization": f"Bearer {token}"}, timeout=60) | |
| RETRY_ATTEMPTS = 5 # requests in all, per call | |
| RETRY_BASE = 1.0 # s, first backoff before jitter | |
| RETRY_CAP = 30.0 # s, longest single wait, Retry-After included | |
| _sleep = time.sleep # tests replace it | |
| def _retry_after(r: httpx.Response) -> float | None: | |
| raw = r.headers.get("retry-after") | |
| if not raw: | |
| return None | |
| try: | |
| return max(0.0, float(raw)) | |
| except ValueError: | |
| pass | |
| try: | |
| when = parsedate_to_datetime(raw) | |
| except (TypeError, ValueError): | |
| return None | |
| return max(0.0, when.timestamp() - time.time()) if when.tzinfo else None | |
| def _retryable(method: str, r: httpx.Response | None) -> bool: | |
| """429 (the Space edge throttling 1,000 agents) and 502/503/504 never reached the app or | |
| were not processed: retry any method. A 500 or a dropped connection may have been | |
| processed, so only reads retry on those.""" | |
| if r is None: | |
| return method == "GET" | |
| return r.status_code in (429, 502, 503, 504) or (r.status_code == 500 and method == "GET") | |
| def _request(c: httpx.Client, method: str, path: str, **kw) -> httpx.Response: | |
| """One controller request, retried up to RETRY_ATTEMPTS times with exponential backoff and | |
| full jitter, never sooner than Retry-After.""" | |
| for attempt in range(RETRY_ATTEMPTS): | |
| last = attempt == RETRY_ATTEMPTS - 1 | |
| try: | |
| r = c.request(method, path, **kw) | |
| except httpx.TransportError: | |
| if last or not _retryable(method, None): | |
| raise | |
| r = None | |
| if r is not None and (last or not _retryable(method, r)): | |
| return r | |
| backoff = random.uniform(0, min(RETRY_CAP, RETRY_BASE * 2 ** attempt)) | |
| wait = max(backoff, (_retry_after(r) or 0.0) if r is not None else 0.0) | |
| _sleep(min(RETRY_CAP, wait)) | |
| raise AssertionError("unreachable") | |
| def _call(method: str, path: str, *, timeout: float | None = None, **kw) -> Any: | |
| if timeout is not None: # per request, without the TestClient's `timeout` deprecation | |
| kw["extensions"] = {"timeout": httpx.Timeout(timeout).as_dict()} | |
| with client() as c: | |
| r = _request(c, method, path, **kw) | |
| if r.status_code >= 400: | |
| typer.echo(f"error {r.status_code}: {r.text}", err=True) | |
| raise typer.Exit(1) | |
| return r.json() | |
| def _as_agent() -> bool: | |
| """Running inside an agent with its scoped token: names are agents, never swarms.""" | |
| return bool(os.environ.get("FLOCK_AGENT_TOKEN")) | |
| def _is_swarm(name: str) -> bool: | |
| """Only a name shaped like an agent's (`...-NNNN`) costs a lookup.""" | |
| if not AGENT_NAME_RE.match(name): | |
| return True | |
| with client() as c: | |
| return _request(c, "GET", f"/swarms/{name}").status_code == 200 | |
| def _table(rows: list[list[Any]], header: list[str]) -> None: | |
| cells = [header] + [[("" if v is None else str(v)) for v in row] for row in rows] | |
| widths = [max(len(r[i]) for r in cells) for i in range(len(header))] | |
| for r in cells: | |
| typer.echo(" ".join(v.ljust(w) for v, w in zip(r, widths, strict=True)).rstrip()) | |
| def _states(d: dict[str, int]) -> str: | |
| return " ".join(f"{k}:{v}" for k, v in sorted(d.items())) or "-" | |
| def _print_status(st: dict) -> None: | |
| typer.echo( | |
| f"{st['swarm']} agents={st['agents_total']} workers={st['workers_total']} " | |
| f"pending={st['pending_prompts']} cost=${st['cost_usd']:.2f}" | |
| ) | |
| rows = [ | |
| [ | |
| t["name"], | |
| t["health"], | |
| t["replicas"], | |
| _states(t["agents_by_state"]), | |
| _states(t["workers_by_state"]), | |
| f"${t['cost_usd']:.2f}", | |
| ] | |
| for t in st["types"].values() | |
| ] | |
| _table(rows, ["type", "health", "replicas", "agents", "workers", "cost"]) | |
| for t in st["types"].values(): | |
| if t.get("launch_error"): | |
| what = ("launch refused, not retrying until the next apply" | |
| if t["health"] == "launch_refused" else "last launch error, retrying") | |
| typer.echo(f"warning: type {t['name']}: {what}: {t['launch_error']}") | |
| for pv in (st.get("providers") or {}).values(): | |
| cap = f"/{pv['max_working']}" if pv.get("max_working") else "" | |
| kind = pv["type"] + (" self-hosted" if pv["self_hosted"] else "") | |
| typer.echo(f"provider {pv['name']} ({kind}): working {pv['working']}{cap}, " | |
| f"types {','.join(pv['types']) or '-'}") | |
| if pv.get("zero_usage_agents"): | |
| typer.echo(f"warning: provider {pv['name']}: {pv['zero_usage_agents']} agents finished " | |
| "turns with zero reported tokens; their cost is unknown (shown as $0)") | |
| def load_spec(spec_file: str) -> SwarmSpec: | |
| """Read a spec and inline each type's system_prompt_file (relative to the spec file).""" | |
| path = Path(spec_file) | |
| try: | |
| spec = SwarmSpec.model_validate(yaml.safe_load(path.read_text(encoding="utf-8"))) | |
| except ValidationError as e: | |
| for err in e.errors(): | |
| where = ".".join(str(x) for x in err["loc"]) | |
| typer.echo(f"{spec_file}: {where + ': ' if where else ''}{err['msg']}", err=True) | |
| raise typer.Exit(1) from None | |
| for name, ts in spec.types.items(): | |
| if ts.system_prompt_file and ts.system_prompt is None: | |
| src = Path(ts.system_prompt_file) | |
| src = src if src.is_absolute() else path.parent / src | |
| if not src.is_file(): | |
| typer.echo(f"type {name}: system_prompt_file {src} not found", err=True) | |
| raise typer.Exit(1) | |
| ts.system_prompt = src.read_text(encoding="utf-8") | |
| return spec | |
| def _provider_backed(spec: SwarmSpec) -> bool: | |
| return any(ts.resolved_model is not None for ts in spec.types.values()) | |
| # Harnesses that cannot add request-body fields, so a model's `extra` is kept but not sent. | |
| # The `llm` harness (and any harness not listed) sends it. | |
| EXTRA_NOT_SENT = frozenset({"claude", "codex", "pi"}) | |
| def _warn_extra(spec: SwarmSpec) -> None: | |
| for name, ts in spec.types.items(): | |
| if ts.harness.value in EXTRA_NOT_SENT and ts.resolved_model is not None and ts.resolved_model.extra: | |
| typer.echo(f"warning: type {name}: `extra` is kept in the spec but not sent by the " | |
| f"{ts.harness.value} harness (docs/open-models.md)", err=True) | |
| def _warn_unverified(models: list[str], why: str) -> None: | |
| for m in models: | |
| typer.echo(f"warning: {m}: unverified ({why})", err=True) | |
| def _unprobed(spec: SwarmSpec) -> list[str]: | |
| return sorted({ts.resolved_model.qualified for ts in spec.types.values() | |
| if ts.resolved_model is not None}) | |
| def _preflight(spec: SwarmSpec, strict: bool) -> bool: | |
| """POST /swarms/{swarm}/preflight with the spec; print one row per model. True when no probe | |
| failed and, with `strict`, every model was actually probed. Never silent about a model that | |
| was not probed.""" | |
| rep = _call("POST", f"/swarms/{spec.swarm}/preflight", json=spec.model_dump(mode="json"), | |
| timeout=300) | |
| rows = [ | |
| [r["model"], ",".join(r["types"]), r["api"], | |
| "unverified" if r.get("unverified") else "ok" if r["ok"] else "FAIL", | |
| "yes" if r["tool_call_seen"] else "no", "yes" if r.get("follow_up_ok") else "no", | |
| "" if r["latency_ms"] is None else r["latency_ms"], r["error"] or ""] | |
| for r in rep["results"] | |
| ] | |
| if rows: | |
| _table(rows, ["model", "types", "api", "result", "tool_call", "follow_up", "latency_ms", | |
| "error"]) | |
| for r in rep["results"]: | |
| if r.get("warning"): | |
| typer.echo(f"warning: {r['model']}: {r['warning']}", err=True) | |
| else: | |
| typer.echo("preflight: no provider-backed models in this spec") | |
| unverified = int(rep.get("unverified") or 0) | |
| if unverified: | |
| typer.echo(f"warning: {unverified} model(s) unverified" | |
| + (" (--strict: counted as failure)" if strict else ""), err=True) | |
| return bool(rep["ok"]) and not (strict and unverified) | |
| def apply( | |
| spec_file: str, | |
| preflight: bool = typer.Option( | |
| None, "--preflight/--skip-preflight", | |
| help="Probe provider-backed models before applying (default: on when the spec has any)."), | |
| strict: bool = typer.Option( | |
| False, "--strict", help="Treat models that were not probed (unverified) as failures."), | |
| ) -> None: | |
| """Apply a swarm spec (YAML). Refuses when a provider-backed model fails preflight; a model | |
| that could not be probed is reported as unverified and applied anyway unless --strict.""" | |
| spec = load_spec(spec_file) | |
| _warn_extra(spec) | |
| if preflight is None: | |
| preflight = _provider_backed(spec) | |
| if preflight: | |
| if not _preflight(spec, strict): | |
| typer.echo("preflight failed: not applying (fix the provider or pass --skip-preflight)", | |
| err=True) | |
| raise typer.Exit(1) | |
| elif unprobed := _unprobed(spec): | |
| _warn_unverified(unprobed, "preflight skipped") | |
| if strict: | |
| typer.echo("--strict: not applying unverified models", err=True) | |
| raise typer.Exit(1) | |
| st = _call("PUT", f"/swarms/{spec.swarm}", json=spec.model_dump(mode="json")) | |
| _print_status(st) | |
| def preflight_cmd( | |
| spec_file: str, | |
| strict: bool = typer.Option(False, "--strict", help="Exit 1 on unverified models too."), | |
| ) -> None: | |
| """Send one small tool-calling request to each provider-backed model in a spec (through | |
| the controller, which holds the keys). Exit 1 if any model fails; models that could not be | |
| probed are reported as unverified (exit 1 only with --strict).""" | |
| spec = load_spec(spec_file) | |
| _warn_extra(spec) | |
| if not _preflight(spec, strict): | |
| raise typer.Exit(1) | |
| def status(swarm: str = typer.Argument(None)) -> None: | |
| """Swarm status: counts by type and state. Given an agent name (or inside an agent), that | |
| agent's status as JSON.""" | |
| if not swarm and _as_agent(): | |
| swarm = os.environ.get("FLOCK_AGENT") | |
| if swarm and (_as_agent() or not _is_swarm(swarm)): | |
| typer.echo(json.dumps(_call("GET", f"/agents/{swarm}"), indent=2)) | |
| return | |
| if swarm: | |
| _print_status(_call("GET", f"/swarms/{swarm}")["status"]) | |
| return | |
| swarms = _call("GET", "/swarms") | |
| if not swarms: | |
| typer.echo("no swarms") | |
| for st in swarms: | |
| _print_status(st) | |
| typer.echo("") | |
| def agents( | |
| swarm: str, | |
| type: str = typer.Option(None, "--type"), | |
| state: str = typer.Option(None, "--state"), | |
| parent: str = typer.Option(None, "--parent", help="only the children of this agent"), | |
| ) -> None: | |
| """List a swarm's agents.""" | |
| params = {k: v for k, v in (("type", type), ("state", state), ("parent", parent)) if v} | |
| rows, cursor = [], None | |
| while True: | |
| page = _call( | |
| "GET", | |
| f"/swarms/{swarm}/agents", | |
| params={**params, "cursor": cursor} if cursor else params, | |
| ) | |
| for a in page["agents"]: | |
| rows.append( | |
| [ | |
| a["name"], | |
| a["state"], | |
| a.get("provider") or "-", | |
| a["worker"], | |
| a["turns"], | |
| a["pending"], | |
| f"${a['cost_usd']:.4f}", # a short turn costs a fraction of a cent | |
| a["last_heartbeat"], | |
| ] | |
| ) | |
| cursor = page["cursor"] | |
| if not cursor: | |
| break | |
| _table(rows, ["name", "state", "provider", "worker", "turns", "pending", "cost", | |
| "last_heartbeat"]) | |
| def agent(name: str) -> None: | |
| """One agent's status as JSON.""" | |
| typer.echo(json.dumps(_call("GET", f"/agents/{name}"), indent=2)) | |
| def prompt( | |
| args: Annotated[list[str], typer.Argument(metavar="[NAME] TEXT")], | |
| timeout: str = typer.Option(None, "--timeout"), | |
| reply_to: str = typer.Option(None, "--reply-to"), | |
| swarm: str = typer.Option(None, "--swarm", help="with --type: the swarm"), | |
| type: str = typer.Option(None, "--type", help="prompt every agent of this type (one request)"), | |
| state: str = typer.Option(None, "--state", help="with --type: only agents in this state, e.g. idle"), | |
| skip_pending: bool = typer.Option( | |
| None, "--skip-pending/--stack", | |
| help="with --type/--children: skip agents that already have an undelivered prompt " | |
| "(default: on with --state idle, off otherwise)"), | |
| children: bool = typer.Option( | |
| False, "--children", help="prompt the direct children of --parent (one request)"), | |
| parent: str = typer.Option(None, "--parent", help="with --children; default: $FLOCK_AGENT"), | |
| ) -> None: | |
| """Send a prompt to an agent (`prompt NAME TEXT`), to every agent of a type | |
| (`prompt --swarm S --type T [--state idle] TEXT`) or to your children | |
| (`prompt --children [--type T] [--state idle] TEXT`); prints the ack (counts for the bulk | |
| forms). With `--state idle` an agent that already has a pending prompt is skipped, so | |
| repeating the call never stacks prompts (`skipped_pending` in the counts).""" | |
| body = {"timeout": timeout, "reply_to": reply_to} | |
| if children: | |
| parent = parent or os.environ.get("FLOCK_AGENT") | |
| if not parent or swarm is not None or len(args) != 1: | |
| typer.echo("usage: flock prompt --children [--parent P] [--type T] [--state STATE] " | |
| "TEXT (parent defaults to $FLOCK_AGENT)", err=True) | |
| raise typer.Exit(2) | |
| body.update(type=type, state=state, text=args[0]) | |
| if skip_pending is not None: | |
| body["skip_pending"] = skip_pending | |
| typer.echo(json.dumps(_call("POST", f"/agents/{parent}/children/prompts", json=body))) | |
| return | |
| if type is not None: | |
| if swarm is None or len(args) != 1: | |
| typer.echo("usage: flock prompt --swarm S --type T [--state STATE] TEXT", err=True) | |
| raise typer.Exit(2) | |
| body.update(type=type, state=state, text=args[0]) | |
| if skip_pending is not None: | |
| body["skip_pending"] = skip_pending | |
| typer.echo(json.dumps(_call("POST", f"/swarms/{swarm}/prompts", json=body))) | |
| return | |
| if swarm is not None or state is not None or len(args) != 2: | |
| typer.echo("usage: flock prompt NAME TEXT (or --swarm S --type T TEXT)", err=True) | |
| raise typer.Exit(2) | |
| body["text"] = args[1] | |
| typer.echo(json.dumps(_call("POST", f"/agents/{args[0]}/prompt", json=body))) | |
| def _short(text: Any, n: int = 160) -> str: | |
| text = " ".join(str(text).split()) | |
| return text if len(text) <= n else text[: n - 1] + "…" | |
| def _blocks_text(content: Any) -> str: | |
| if isinstance(content, str): | |
| return content | |
| parts = [] | |
| for b in content or []: | |
| if b.get("type") == "text": | |
| parts.append(b.get("text", "")) | |
| elif b.get("type") == "toolCall": | |
| parts.append(f"[{b.get('name')} {_short(json.dumps(b.get('arguments')), 80)}]") | |
| return " ".join(parts) | |
| def pretty_event(ev: dict) -> str: | |
| """One line per event: time, turn, kind, then the text or a summary (`events --pretty`).""" | |
| ts = str(ev.get("ts", ""))[11:19] | |
| turn = ev.get("turn_id") or "-" | |
| data = ev.get("data") or {} | |
| kind = ev.get("type", "?") | |
| if kind == "message": | |
| m = data.get("message") or {} | |
| role = m.get("role", "?") | |
| if role == "toolResult": | |
| return f"{ts} {turn:>5} tool {m.get('toolName', '')}: {_short(_blocks_text(m.get('content')), 120)}" | |
| return f"{ts} {turn:>5} {role:<8} {_short(_blocks_text(m.get('content')))}" | |
| detail = data.get("data", data) if isinstance(data.get("data"), dict) else data | |
| if kind == "turn_ended": | |
| toks = detail.get("tokens") or {} | |
| used = " ".join(f"{model} in={u.get('input', 0)} out={u.get('output', 0)}" | |
| for model, u in toks.items()) | |
| return f"{ts} {turn:>5} {kind:<8} {detail.get('status')} {used}".rstrip() | |
| keep = {k: v for k, v in detail.items() if v is not None} | |
| summary = " ".join(f"{k}={_short(v, 60)}" for k, v in keep.items()) | |
| return f"{ts} {turn:>5} {kind:<8} {summary}".rstrip() | |
| def events( | |
| name: str, | |
| after: int = typer.Option(0, "--after"), | |
| follow: bool = typer.Option(False, "--follow", "-f"), | |
| wait: bool = typer.Option(False, "--wait", help="block until an event after --after"), | |
| timeout: float = typer.Option(600, "--timeout", help="seconds --wait blocks at most"), | |
| pretty: bool = typer.Option(False, "--pretty", help="one readable line per event"), | |
| ) -> None: | |
| """Print an agent's event stream (JSONL); --follow keeps long-polling; --wait returns once | |
| something new arrived (the last line on stderr is the cursor for the next --after); | |
| --pretty prints time, turn, kind and the text or a summary instead of JSON.""" | |
| deadline = time.monotonic() + timeout | |
| while True: | |
| poll = follow or wait | |
| page = _call( | |
| "GET", f"/agents/{name}/events", params={"after": after, "wait": 25 if poll else 0} | |
| ) | |
| for ev in page["events"]: | |
| typer.echo(pretty_event(ev) if pretty else json.dumps(ev)) | |
| after = page["next"] | |
| if wait and (page["events"] or time.monotonic() >= deadline): | |
| typer.echo(f"next: {after}", err=True) | |
| break | |
| if not poll: | |
| break | |
| def scale(swarm: str, type: str, replicas: int) -> None: | |
| """Set a type's replicas (edits the stored spec and re-applies it).""" | |
| spec = SwarmSpec.model_validate(_call("GET", f"/swarms/{swarm}")["spec"]) | |
| if type not in spec.types: | |
| typer.echo(f"no type {type!r} in swarm {swarm!r}", err=True) | |
| raise typer.Exit(1) | |
| spec.types[type].replicas = replicas | |
| _print_status(_call("PUT", f"/swarms/{swarm}", json=spec.model_dump(mode="json"))) | |
| def remove(swarm: str) -> None: | |
| """Drain and remove a swarm. Given a spawned agent's name (or inside an agent), remove that | |
| agent and everything it spawned.""" | |
| if _as_agent() or not _is_swarm(swarm): | |
| typer.echo(json.dumps(_call("DELETE", f"/agents/{swarm}"))) | |
| return | |
| typer.echo(json.dumps(_call("DELETE", f"/swarms/{swarm}"))) | |
| def spawn( | |
| type: str, | |
| prompt: str = typer.Option(None, "--prompt", help="first prompt for the child"), | |
| key: str = typer.Option(None, "--key", help="idempotency key: a repeat returns the same child"), | |
| parent: str = typer.Option(None, "--parent", help="default: $FLOCK_AGENT (yourself)"), | |
| as_json: bool = typer.Option(False, "--json", help="print the whole ack"), | |
| ) -> None: | |
| """Spawn a child agent of TYPE; prints its name.""" | |
| parent = parent or os.environ.get("FLOCK_AGENT") | |
| if not parent: | |
| typer.echo("no parent: pass --parent or set FLOCK_AGENT", err=True) | |
| raise typer.Exit(1) | |
| body = {"type": type, "prompt": prompt, "key": key} | |
| ack = _call("POST", f"/agents/{parent}/spawn", json=body) | |
| typer.echo(json.dumps(ack) if as_json else ack["name"]) | |
| def whoami() -> None: | |
| """Who the configured token acts as: the controller, or an agent with its lineage.""" | |
| typer.echo(json.dumps(_call("GET", "/whoami"), indent=2)) | |
| def workers(swarm: str) -> None: | |
| """List a swarm's workers.""" | |
| rows = [ | |
| [ | |
| w["id"], | |
| w["type"], | |
| w["state"], | |
| w["job_id"], | |
| len(w["agents"]), | |
| w["crashes"], | |
| w["deadline"], | |
| w["last_heartbeat"], | |
| w.get("last_error") or "", | |
| ] | |
| for w in _call("GET", "/workers", params={"swarm": swarm}) | |
| ] | |
| _table(rows, ["id", "type", "state", "job", "agents", "crashes", "deadline", "last_heartbeat", | |
| "last_error"]) | |
| typer.echo("full job logs: hf jobs logs <job>", err=True) | |
| if __name__ == "__main__": | |
| app() | |