Spaces:
Sleeping
Sleeping
File size: 6,617 Bytes
97ce220 | 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 | """A2A protocol layer (plan §7.2), built on the official a2a-sdk (0.3.x, pydantic).
- Each agent publishes an AgentCard at /.well-known/agent-card.json.
- Messaging is JSON-RPC 2.0 `message/send`; payloads travel as DataPart.
- `serve_agent` runs an agent as a Starlette app on its own port (background thread).
- `send_payload` is the client: resolves the peer's card, sends, returns the reply dict.
"""
from __future__ import annotations
import asyncio
import threading
import uuid
from collections.abc import Callable
import httpx
import uvicorn
from a2a.client import A2ACardResolver, ClientConfig, ClientFactory
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.apps import A2AStarletteApplication
from a2a.server.events import EventQueue
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.tasks import InMemoryTaskStore
from a2a.types import (
AgentCapabilities,
AgentCard,
AgentSkill,
DataPart,
Message,
Part,
Role,
)
def make_card(name: str, description: str, port: int, skills: list[dict],
streaming: bool = False) -> AgentCard:
return AgentCard(
protocol_version="0.3.0",
name=name,
description=description,
version="1.0.0",
url=f"http://127.0.0.1:{port}",
capabilities=AgentCapabilities(streaming=streaming),
default_input_modes=["application/json"],
default_output_modes=["application/json"],
skills=[AgentSkill(id=s["id"], name=s.get("name", s["id"]),
description=s["description"], tags=s.get("tags", []))
for s in skills],
)
def extract_payload(message: Message) -> dict | None:
"""First DataPart payload in a message, else None."""
for part in message.parts or []:
inner = part.root if isinstance(part, Part) else part
if isinstance(inner, DataPart):
return inner.data
return None
def payload_message(payload: dict, role: Role = Role.agent) -> Message:
return Message(
message_id=uuid.uuid4().hex,
role=role,
parts=[Part(root=DataPart(data=payload))],
)
class PayloadExecutor(AgentExecutor):
"""Adapts a plain `handler(payload: dict) -> dict` into an A2A AgentExecutor."""
def __init__(self, handler: Callable[[dict], dict]):
self.handler = handler
async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
payload = extract_payload(context.message) if context.message else None
try:
if asyncio.iscoroutinefunction(self.handler):
reply = await self.handler(payload or {})
else:
reply = await asyncio.to_thread(self.handler, payload or {})
except Exception as e: # noqa: BLE001 — surface agent errors as structured replies
reply = {"status": "error", "detail": str(e)}
await event_queue.enqueue_event(payload_message(reply))
async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
await event_queue.enqueue_event(payload_message({"status": "cancelled"}))
class AgentServer:
"""Runs one agent as an A2A Starlette app on a background uvicorn thread."""
def __init__(self, card: AgentCard, handler: Callable[[dict], dict], port: int):
self.card = card
self.port = port
handler_ = DefaultRequestHandler(
agent_executor=PayloadExecutor(handler),
task_store=InMemoryTaskStore(),
)
self.app = A2AStarletteApplication(agent_card=card, http_handler=handler_).build()
self._server: uvicorn.Server | None = None
self._thread: threading.Thread | None = None
def start(self, wait_ready_s: float = 10.0) -> None:
config = uvicorn.Config(self.app, host="127.0.0.1", port=self.port,
log_level="warning")
self._server = uvicorn.Server(config)
self._thread = threading.Thread(target=self._server.run, daemon=True)
self._thread.start()
# block until the card endpoint answers
deadline = wait_ready_s
import time
while deadline > 0:
try:
r = httpx.get(f"http://127.0.0.1:{self.port}/.well-known/agent-card.json",
timeout=1.0)
if r.status_code == 200:
return
except httpx.HTTPError:
pass
time.sleep(0.2)
deadline -= 0.2
raise RuntimeError(f"A2A agent on port {self.port} did not become ready")
def stop(self) -> None:
if self._server:
self._server.should_exit = True
if self._thread:
self._thread.join(timeout=5)
async def send_payload(base_url: str, payload: dict, timeout_s: float = 120.0) -> dict:
"""A2A client: resolve the peer's AgentCard, send one message/send, return reply."""
async with httpx.AsyncClient(timeout=timeout_s) as http:
card = await A2ACardResolver(http, base_url).get_agent_card()
client = ClientFactory(ClientConfig(httpx_client=http, streaming=False)).create(card)
message = payload_message(payload, role=Role.user)
reply_payload: dict | None = None
async for event in client.send_message(message):
if isinstance(event, Message):
reply_payload = extract_payload(event)
elif isinstance(event, tuple): # (Task, UpdateEvent)
task = event[0]
if task and task.artifacts:
for artifact in task.artifacts:
for part in artifact.parts:
inner = part.root if isinstance(part, Part) else part
if isinstance(inner, DataPart):
reply_payload = inner.data
if task and task.status and task.status.message:
got = extract_payload(task.status.message)
if got is not None:
reply_payload = got
if reply_payload is None:
raise RuntimeError(f"no DataPart reply from {base_url}")
return reply_payload
def export_card(card: AgentCard, path) -> None:
"""Write the AgentCard JSON (spec camelCase) to protocols/a2a_cards/."""
import json
from pathlib import Path
Path(path).parent.mkdir(parents=True, exist_ok=True)
Path(path).write_text(
json.dumps(card.model_dump(by_alias=True, exclude_none=True), indent=2),
encoding="utf-8",
)
|