ship grants, a2a_client, discovery, sandbox SDK + tests
This commit is contained in:
robert
253
examples/multi_agent.py
Normal file
253
examples/multi_agent.py
Normal file
@@ -0,0 +1,253 @@
|
||||
"""End-to-end demo of the platform's killer flow.
|
||||
|
||||
User uploads files → main agent discovers a community-published graph agent →
|
||||
delegates a scoped grant → graph agent runs in its own RunContext, sees only
|
||||
the granted files, returns artifacts → main agent presents results.
|
||||
|
||||
Two agents in the same process for the demo. In production, the graph agent
|
||||
is on another pod (`HttpA2AClient`) and the discovery client hits the
|
||||
control plane registry — the agent code is unchanged.
|
||||
|
||||
Run::
|
||||
|
||||
cd apps/a2a
|
||||
pip install -e '.[dev]'
|
||||
python -m examples.multi_agent
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
from a2a_pack import (
|
||||
A2AAgent,
|
||||
DiscoveredAgent,
|
||||
FileType,
|
||||
InMemoryA2AClient,
|
||||
InMemoryDiscovery,
|
||||
LocalRunContext,
|
||||
LocalWorkspaceClient,
|
||||
NoAuth,
|
||||
RunContext,
|
||||
WorkspaceAccess,
|
||||
WorkspaceMode,
|
||||
skill,
|
||||
verify_grant,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# graph-agent: a community-published agent another developer wrote.
|
||||
# Tags itself for discovery, declares a workspace policy.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class GraphAgent(A2AAgent):
|
||||
name = "graph-agent-v2"
|
||||
description = "Renders dashboards from spreadsheets"
|
||||
tools_used = ("matplotlib", "pandas")
|
||||
workspace_access = WorkspaceAccess.dynamic(
|
||||
max_files=8,
|
||||
allowed_modes=(WorkspaceMode.READ_ONLY,),
|
||||
deny_patterns=("secrets/**", "**/.env"),
|
||||
)
|
||||
|
||||
@skill(
|
||||
description="Generate a dashboard from spreadsheet files",
|
||||
tags=["visualization", "spreadsheet", "chart"],
|
||||
)
|
||||
async def generate_dashboard(
|
||||
self, ctx: RunContext[NoAuth], prompt: str
|
||||
) -> dict:
|
||||
# Receive the workspace bound by the inbound grant. We can ONLY see
|
||||
# files the caller delegated; everything else (secrets, etc.) is
|
||||
# invisible.
|
||||
ws = ctx.workspace
|
||||
spreadsheets = await ws.search(
|
||||
query="data sales revenue", types=[FileType.OTHER], limit=5
|
||||
)
|
||||
files_seen = [m.path for m in spreadsheets]
|
||||
|
||||
await ctx.emit_progress(f"reading {len(files_seen)} files for: {prompt}")
|
||||
|
||||
# Pretend we generated a chart. Stage outputs as artifacts; the
|
||||
# platform stages them as patches under the grant's outputs prefix.
|
||||
chart_bytes = (
|
||||
b"\x89PNG\r\n\x1a\n" # tiny fake png prefix
|
||||
+ json.dumps({"prompt": prompt, "files": files_seen}).encode()
|
||||
)
|
||||
ref = await ctx.write_artifact(
|
||||
"charts/dashboard.png", chart_bytes, "image/png"
|
||||
)
|
||||
await ctx.emit_artifact(ref)
|
||||
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"files_used": files_seen,
|
||||
"chart": ref.uri,
|
||||
"bucket_seen": getattr(ws, "bucket", None),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# main-agent: orchestrates the user's session. The user's full workspace is
|
||||
# bound to its RunContext; it scopes a grant down to the spreadsheets only
|
||||
# before calling the graph agent.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class MainAgent(A2AAgent):
|
||||
name = "session-orchestrator"
|
||||
description = "Routes user intents to the right specialist agent"
|
||||
|
||||
@skill(description="Make me a chart from my uploaded spreadsheets")
|
||||
async def make_chart(self, ctx: RunContext[NoAuth], prompt: str) -> dict:
|
||||
# 1. Discover: find a graph-capable agent in the registry.
|
||||
candidates = await ctx.discover.find_agents(tags=["visualization"])
|
||||
if not candidates:
|
||||
raise RuntimeError("no graph-capable agent registered")
|
||||
graph = candidates[0]
|
||||
await ctx.emit_progress(f"found {graph.name}; delegating workspace")
|
||||
|
||||
# 2. Delegate: mint a grant scoped to *.xlsx, deny secrets, write
|
||||
# to charts/ only, expires in 5 minutes.
|
||||
token = await ctx.workspace.delegate(
|
||||
audience=graph.name,
|
||||
allow_patterns=("*.xlsx", "*.csv"),
|
||||
deny_patterns=("secrets/**", "**/.env"),
|
||||
outputs_prefix="charts/",
|
||||
ttl_seconds=300,
|
||||
)
|
||||
decoded = verify_grant(token)
|
||||
await ctx.emit_event_kind(
|
||||
"delegation",
|
||||
{
|
||||
"to": graph.name,
|
||||
"grant_id": decoded.grant_id,
|
||||
"allow": list(decoded.allow_patterns),
|
||||
"deny": list(decoded.deny_patterns),
|
||||
"expires_at": decoded.expires_at,
|
||||
},
|
||||
) if hasattr(ctx, "emit_event_kind") else None
|
||||
|
||||
# 3. Call: invoke the graph agent. Runtime hands the grant in the
|
||||
# body; receiving runtime materializes a workspace from it.
|
||||
result = await ctx.call(
|
||||
graph.name,
|
||||
"generate_dashboard",
|
||||
args={"prompt": prompt},
|
||||
grant=token,
|
||||
)
|
||||
|
||||
return {
|
||||
"delegated_to": graph.name,
|
||||
"grant_id": result.grant_id,
|
||||
"graph_response": result.result,
|
||||
"events_from_callee": [e["kind"] for e in result.events],
|
||||
"artifacts_from_callee": list(result.artifacts),
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# wire-up: in-memory router + discovery (replace with HTTP + control plane
|
||||
# in production with zero agent-code changes)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_in_memory_runtime(
|
||||
user_workspace: LocalWorkspaceClient, agents: dict[str, A2AAgent]
|
||||
):
|
||||
def factory(agent: A2AAgent, grant_token: str | None):
|
||||
ws = None
|
||||
if grant_token is not None:
|
||||
grant = verify_grant(grant_token)
|
||||
visible = {
|
||||
p: b
|
||||
for p, b in user_workspace._files.items()
|
||||
if not any(
|
||||
p.startswith((d.rstrip("*").rstrip("/")))
|
||||
for d in grant.deny_patterns
|
||||
if d
|
||||
)
|
||||
}
|
||||
ws = LocalWorkspaceClient(
|
||||
files=visible,
|
||||
access=WorkspaceAccess.dynamic(
|
||||
max_files=64,
|
||||
allowed_modes=(WorkspaceMode.READ_ONLY,),
|
||||
deny_patterns=tuple(grant.deny_patterns),
|
||||
),
|
||||
bucket=grant.bucket,
|
||||
issuer=grant.audience,
|
||||
)
|
||||
return LocalRunContext(auth=NoAuth(), workspace=ws)
|
||||
|
||||
return InMemoryA2AClient(agents=agents, ctx_factory=factory)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# main
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
user_workspace = LocalWorkspaceClient(
|
||||
files={
|
||||
"sales_q1.xlsx": b"q1 data",
|
||||
"sales_q2.xlsx": b"q2 data",
|
||||
"notes.md": b"# notes",
|
||||
"secrets/.env": b"DB_PASSWORD=NEVER",
|
||||
},
|
||||
access=WorkspaceAccess.dynamic(
|
||||
max_files=10,
|
||||
allowed_modes=(
|
||||
WorkspaceMode.READ_ONLY,
|
||||
WorkspaceMode.READ_WRITE_OVERLAY,
|
||||
),
|
||||
deny_patterns=("secrets/**",),
|
||||
),
|
||||
bucket="user-42-files",
|
||||
issuer="user-42",
|
||||
)
|
||||
|
||||
main_a = MainAgent()
|
||||
graph_a = GraphAgent()
|
||||
|
||||
discovery = InMemoryDiscovery(
|
||||
{
|
||||
graph_a.name: DiscoveredAgent(
|
||||
name=graph_a.name, url=None, card=graph_a.card()
|
||||
)
|
||||
}
|
||||
)
|
||||
a2a = build_in_memory_runtime(
|
||||
user_workspace=user_workspace, agents={graph_a.name: graph_a}
|
||||
)
|
||||
|
||||
print("=== main-agent ships uploaded files ===")
|
||||
print(f" user bucket: {user_workspace.bucket}")
|
||||
print(f" files: {sorted(user_workspace._files)}")
|
||||
print()
|
||||
|
||||
out = await main_a.local_invoke(
|
||||
"make_chart",
|
||||
workspace=user_workspace,
|
||||
a2a=a2a,
|
||||
discover=discovery,
|
||||
prompt="weekly burn rate by quarter",
|
||||
)
|
||||
|
||||
print("=== main-agent result ===")
|
||||
print(json.dumps(out, indent=2))
|
||||
print()
|
||||
print("=== what graph-agent could see ===")
|
||||
print(f" bucket: {out['graph_response']['bucket_seen']}")
|
||||
print(f" files_used: {out['graph_response']['files_used']}")
|
||||
print(
|
||||
" → secrets/.env is NOT in the list. The grant denied it; the "
|
||||
"callee's runtime never made it visible."
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user