Lab 10: A2A Protocol¶

"""Lab 10: serve the triage agent over the A2A protocol."""

import asyncio

import uvicorn
from a2a.helpers import (
    get_message_text,
    new_data_artifact_update_event,
    new_task_from_user_message,
    new_text_status_update_event,
)
from a2a.server.agent_execution import AgentExecutor
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.routes import create_agent_card_routes, create_jsonrpc_routes
from a2a.server.tasks import InMemoryTaskStore
from a2a.types import AgentCapabilities, AgentCard, AgentInterface, AgentSkill, TaskState
from starlette.applications import Starlette

from triage import TASK, triage

HOST, PORT = "127.0.0.1", 9999

AGENT_CARD = AgentCard(
    name="Checkout triage agent",
    description="Investigates checkout incidents from logs, config, deploys, and metrics. Read-only.",
    version="1.0.0",
    supported_interfaces=[
        AgentInterface(protocol_binding="JSONRPC", protocol_version="1.0", url=f"http://{HOST}:{PORT}/a2a/jsonrpc/")
    ],
    capabilities=AgentCapabilities(streaming=True),
    default_input_modes=["text/plain"],
    default_output_modes=["application/json"],
    skills=[
        AgentSkill(
            id="triage",
            name="Incident triage",
            description="Finds the root cause of the current checkout alert and returns a triage report.",
            tags=["incident", "sre", "checkout"],
            examples=["Why did the checkout alert fire?"],
        )
    ],
)


class TriageExecutor(AgentExecutor):
    async def execute(self, context, event_queue):
        task = context.current_task or new_task_from_user_message(context.message)
        await event_queue.enqueue_event(task)  # a Task must be the first event

        def status(state, text):
            return new_text_status_update_event(
                task_id=task.id, context_id=task.context_id, state=state, text=text)

        await event_queue.enqueue_event(status(TaskState.TASK_STATE_WORKING, "Investigating"))
        # The caller's text is extra context, never a replacement for our own task and rules.
        request = get_message_text(context.message)
        result = await asyncio.to_thread(
            triage, task=f"{TASK}\n\nRequest from the calling agent:\n{request}", log=lambda *_: None)

        if result["report"] is None:
            await event_queue.enqueue_event(status(TaskState.TASK_STATE_FAILED, f"No report: {result['status']}"))
            return
        await event_queue.enqueue_event(new_data_artifact_update_event(
            task_id=task.id, context_id=task.context_id, name="triage-report", data=result["report"]))
        await event_queue.enqueue_event(
            status(TaskState.TASK_STATE_COMPLETED, f"Done in {result['steps']} steps"))

    async def cancel(self, context, event_queue):
        raise NotImplementedError("this agent does not support cancellation")


handler = DefaultRequestHandler(
    agent_executor=TriageExecutor(), task_store=InMemoryTaskStore(), agent_card=AGENT_CARD)
routes = create_agent_card_routes(AGENT_CARD) + create_jsonrpc_routes(handler, rpc_url="/a2a/jsonrpc/")
app = Starlette(routes=routes)

if __name__ == "__main__":
    uvicorn.run(app, host=HOST, port=PORT)

"""Lab 10: call a remote agent over A2A, and wrap it as a tool for your own agent."""

import asyncio
import json
import sys

from a2a.client import create_client
from a2a.helpers import get_data_parts, get_message_text, new_text_message
from a2a.types import Role, SendMessageRequest, TaskState

AGENT_URL = "http://127.0.0.1:9999/"


async def ask_remote_agent(question: str, url: str = AGENT_URL, verbose: bool = True) -> dict:
    client = await create_client(url)  # fetches /.well-known/agent-card.json first
    request = SendMessageRequest(message=new_text_message(question, role=Role.ROLE_USER))
    report, state = None, None
    async for event in client.send_message(request):
        if event.HasField("task"):
            if verbose:
                print(f"task {event.task.id}")
        elif event.HasField("status_update"):
            status = event.status_update.status
            state = TaskState.Name(status.state)
            if verbose:
                print(f"status {state}: {get_message_text(status.message)}")
        elif event.HasField("artifact_update"):
            report = get_data_parts(event.artifact_update.artifact.parts)[0]
        elif event.HasField("message"):
            print(get_message_text(event.message))
    return {"state": state, "report": report}


# The same call as a tool for the agent loop from Lab 2: run_agent(..., handlers={"ask_triage_agent": ...})
ASK_TRIAGE_TOOL = {
    "name": "ask_triage_agent",
    "description": "Ask the checkout team's triage agent to investigate the checkout service. "
    "Returns its report as JSON. Takes a minute or two.",
    "strict": True,
    "input_schema": {
        "type": "object",
        "properties": {"question": {"type": "string"}},
        "required": ["question"],
        "additionalProperties": False,
    },
}


def ask_triage_agent(question: str) -> str:
    return json.dumps(asyncio.run(ask_remote_agent(question, verbose=False)))


if __name__ == "__main__":
    question = " ".join(sys.argv[1:]) or "Why did the checkout alert fire?"
    result = asyncio.run(ask_remote_agent(question))
    print(json.dumps(result["report"], indent=2))

from a2a_client import ASK_TRIAGE_TOOL, ask_triage_agent
from agent import run_agent

result = run_agent(
    "Customers report failed payments. Is checkout involved? Ask the checkout team's agent.",
    tools=[ASK_TRIAGE_TOOL],
    handlers={"ask_triage_agent": ask_triage_agent},
)
print(result["answer"])