|
| 1 | +"""Orchestrator + sub-agent swarm with cryptographically signed tool calls. |
| 2 | +
|
| 3 | +The orchestrator is a PydanticAI agent. Each of its tools delegates to a |
| 4 | +specialized sub-agent, which signs the action with its own key and embeds its |
| 5 | +delegation token. Every action in the audit log is traceable to the human |
| 6 | +who bootstrapped the swarm. |
| 7 | +""" |
| 8 | + |
| 9 | +from __future__ import annotations |
| 10 | + |
| 11 | +import sys |
| 12 | +from dataclasses import dataclass |
| 13 | + |
| 14 | +from pydantic_ai import Agent, RunContext |
| 15 | +from rich.console import Console |
| 16 | +from rich.text import Text |
| 17 | + |
| 18 | +from agent_swarm import tools |
| 19 | +from agent_swarm.audit import append_envelope, clear, read_all, save_swarm_keys |
| 20 | +from agent_swarm.identities import SwarmIdentity, make_swarm |
| 21 | +from agent_swarm.signing import CapabilityError, sign_tool_call |
| 22 | + |
| 23 | +console = Console() |
| 24 | + |
| 25 | + |
| 26 | +@dataclass |
| 27 | +class OrchestratorDeps: |
| 28 | + data_agent: SwarmIdentity |
| 29 | + analysis_agent: SwarmIdentity |
| 30 | + notify_agent: SwarmIdentity |
| 31 | + |
| 32 | + |
| 33 | +_orchestrator = Agent( |
| 34 | + "openai:gpt-4o-mini", |
| 35 | + deps_type=OrchestratorDeps, |
| 36 | + system_prompt=( |
| 37 | + "You are a data analyst orchestrator. " |
| 38 | + "Coordinate your specialized sub-agents to fulfill the user's request. " |
| 39 | + "Use read_data to fetch data, analyze_data to summarize it, " |
| 40 | + "and send_notification to deliver results. " |
| 41 | + "Be concise — one short paragraph max." |
| 42 | + ), |
| 43 | +) |
| 44 | + |
| 45 | + |
| 46 | +def _record(agent: SwarmIdentity, tool_name: str, args: dict, cap: str) -> None: |
| 47 | + """Sign the tool call and append it to the audit log; print status line.""" |
| 48 | + envelope = sign_tool_call(agent, tool_name, args, cap) |
| 49 | + append_envelope(envelope) |
| 50 | + |
| 51 | + display_args = ", ".join( |
| 52 | + f"{k}={repr(v[:40] + '...' if isinstance(v, str) and len(v) > 40 else v)}" |
| 53 | + for k, v in args.items() |
| 54 | + ) |
| 55 | + line = Text() |
| 56 | + line.append(f" [{agent.name}] ", style="bold cyan") |
| 57 | + line.append(f"{tool_name}({display_args})", style="dim") |
| 58 | + line.append(" ✓ signed", style="green bold") |
| 59 | + console.print(line) |
| 60 | + |
| 61 | + |
| 62 | +@_orchestrator.tool |
| 63 | +def read_data(ctx: RunContext[OrchestratorDeps], path: str) -> str: |
| 64 | + """Read data from a CSV file (delegated to DataAgent).""" |
| 65 | + _record(ctx.deps.data_agent, "read_csv", {"path": path}, "read_data") |
| 66 | + return tools.read_csv(path) |
| 67 | + |
| 68 | + |
| 69 | +@_orchestrator.tool |
| 70 | +def analyze_data(ctx: RunContext[OrchestratorDeps], data: str) -> str: |
| 71 | + """Summarize the provided data (delegated to AnalysisAgent).""" |
| 72 | + _record(ctx.deps.analysis_agent, "summarize", {"data": data}, "analyze") |
| 73 | + return tools.summarize(data) |
| 74 | + |
| 75 | + |
| 76 | +@_orchestrator.tool |
| 77 | +def send_notification(ctx: RunContext[OrchestratorDeps], channel: str, message: str) -> bool: |
| 78 | + """Send a notification to a channel (delegated to NotifyAgent).""" |
| 79 | + args = {"channel": channel, "message": message} |
| 80 | + _record(ctx.deps.notify_agent, "send_notification", args, "notify") |
| 81 | + return tools.send_notification(channel, message) |
| 82 | + |
| 83 | + |
| 84 | +def _print_identity_tree( |
| 85 | + human: SwarmIdentity, |
| 86 | + orchestrator: SwarmIdentity, |
| 87 | + sub_agents: list[SwarmIdentity], |
| 88 | +) -> None: |
| 89 | + console.print("\n[bold]Swarm identity chain:[/bold]") |
| 90 | + console.print(f" [yellow]{human.name}[/yellow] [dim]{human.did[:40]}...[/dim]") |
| 91 | + caps = ", ".join(orchestrator.capabilities) |
| 92 | + did_abbrev = orchestrator.did[:40] |
| 93 | + console.print(f" └─ [cyan]{orchestrator.name}[/cyan] [dim]{did_abbrev}...[/dim] [{caps}]") |
| 94 | + for i, agent in enumerate(sub_agents): |
| 95 | + prefix = "└─" if i == len(sub_agents) - 1 else "├─" |
| 96 | + caps = ", ".join(agent.capabilities) |
| 97 | + console.print( |
| 98 | + f" {prefix} [green]{agent.name}[/green]" |
| 99 | + f" [dim]{agent.did[:40]}...[/dim] [{caps}]" |
| 100 | + ) |
| 101 | + console.print() |
| 102 | + |
| 103 | + |
| 104 | +def _demo_scope_violation(sub_agents: list[SwarmIdentity]) -> None: |
| 105 | + """Show that a sub-agent cannot exceed its granted capabilities.""" |
| 106 | + data_agent = sub_agents[0] |
| 107 | + console.print("[bold]Scope enforcement demo:[/bold]") |
| 108 | + console.print( |
| 109 | + f" Attempting to use [green]{data_agent.name}[/green] for a" |
| 110 | + " [red]'notify'[/red] action it was never granted..." |
| 111 | + ) |
| 112 | + try: |
| 113 | + sign_tool_call(data_agent, "send_notification", {"channel": "team", "message": "hi"}, "notify") # noqa: E501 |
| 114 | + console.print(" [red]ERROR: scope check did not fire[/red]") |
| 115 | + except CapabilityError as e: |
| 116 | + console.print(f" [green]✓ Blocked:[/green] [dim]{e}[/dim]\n") |
| 117 | + |
| 118 | + |
| 119 | +def main() -> None: |
| 120 | + prompt = ( |
| 121 | + " ".join(sys.argv[1:]) |
| 122 | + or "Read data/sales.csv, analyze it, and send a summary notification to the team channel." |
| 123 | + ) |
| 124 | + |
| 125 | + clear() |
| 126 | + human, orchestrator, sub_agents = make_swarm() |
| 127 | + data_agent, analysis_agent, notify_agent = sub_agents |
| 128 | + |
| 129 | + save_swarm_keys(human, orchestrator, sub_agents) |
| 130 | + _print_identity_tree(human, orchestrator, sub_agents) |
| 131 | + _demo_scope_violation(sub_agents) |
| 132 | + |
| 133 | + console.print(f"[bold]Task:[/bold] {prompt}\n") |
| 134 | + |
| 135 | + result = _orchestrator.run_sync( |
| 136 | + prompt, |
| 137 | + deps=OrchestratorDeps( |
| 138 | + data_agent=data_agent, |
| 139 | + analysis_agent=analysis_agent, |
| 140 | + notify_agent=notify_agent, |
| 141 | + ), |
| 142 | + ) |
| 143 | + |
| 144 | + console.print(f"\n[bold]Result:[/bold] {result.output}\n") |
| 145 | + |
| 146 | + entries = read_all() |
| 147 | + signed_dids = {e.get("identity") for e in entries} |
| 148 | + agent_count = len(signed_dids) |
| 149 | + |
| 150 | + console.print(f"[green]✓[/green] {len(entries)} action(s) across {agent_count} agent(s)") |
| 151 | + console.print("[dim]Run [bold]verify-swarm[/bold] to verify the full delegation chain.[/dim]\n") |
0 commit comments