a2a_example.py

"""
A2A Example: Three-Agent Collaboration via Google Agent2Agent Protocol (A2A)

Demonstrates capability discovery, task management, collaboration, and message exchange
between a client (Orchestrator) and three specialized agents:
- SearchAgent: finds candidate profiles
- FilterAgent: filters candidates by attribute
- ScheduleAgent: schedules interviews

Protocol elements:
- AgentCard: advertises capabilities
- Task: encapsulates task metadata and artifact result
- Message: JSON-wrapped communication with parts negotiation
"""
import json
import uuid
from typing import List, Dict, Any

# --- Protocol Data Structures ---
class AgentCard:
    """Agent description for capability discovery"""
    def __init__(self, name: str, capabilities: List[str]):
        self.name = name
        self.capabilities = capabilities
    def to_json(self) -> str:
        return json.dumps({"name": self.name, "capabilities": self.capabilities})
    @staticmethod
    def from_json(card_str: str) -> 'AgentCard':
        data = json.loads(card_str)
        return AgentCard(data["name"], data["capabilities"])

class Task:
    """Defines a task lifecycle and artifact"""
    def __init__(self, id: str, description: str, payload: Any):
        self.id = id
        self.description = description
        self.payload = payload
        self.status = "pending"
        self.artifact = None
    def to_json(self) -> str:
        return json.dumps({
            "id": self.id,
            "description": self.description,
            "payload": self.payload,
            "status": self.status,
            "artifact": self.artifact
        })
    @staticmethod
    def from_json(task_str: str) -> 'Task':
        data = json.loads(task_str)
        t = Task(data["id"], data["description"], data["payload"])
        t.status = data["status"]
        t.artifact = data["artifact"]
        return t

class Message:
    """Encapsulates protocol messages with optional parts"""
    def __init__(self, sender: str, receiver: str, content: Any, parts: List[Dict[str, Any]] = None):
        self.sender = sender
        self.receiver = receiver
        self.content = content
        self.parts = parts or []
    def to_json(self) -> str:
        return json.dumps({
            "sender": self.sender,
            "receiver": self.receiver,
            "content": self.content,
            "parts": self.parts
        })
    @staticmethod
    def from_json(msg_str: str) -> 'Message':
        data = json.loads(msg_str)
        return Message(data["sender"], data["receiver"], data["content"], data.get("parts", []))

# --- Agent Base Class ---
class Agent:
    def __init__(self, name: str, capabilities: List[str]):
        self.card = AgentCard(name, capabilities)

    def advertise(self) -> str:
        """Return AgentCard JSON"""
        print(f"[{self.card.name}] Advertising capabilities: {self.card.capabilities}")
        return self.card.to_json()

    def handle_message(self, msg_json: str) -> str:
        """Default echo behavior"""
        msg = Message.from_json(msg_json)
        print(f"[{self.card.name}] Received message from {msg.sender}: {msg.content}")
        resp = Message(self.card.name, msg.sender, {"ack": msg.content})
        return resp.to_json()

# --- Specific Agents ---
class SearchAgent(Agent):
    def __init__(self):
        super().__init__("SearchAgent", ["search_candidates"])

    def handle_message(self, msg_json: str) -> str:
        msg = Message.from_json(msg_json)
        payload = msg.content
        if payload.get("action") == "search":
            # Dummy candidate profiles
            criteria = payload.get("criteria", {})
            candidates = [
                {"name": "Alice", "location": criteria.get("location", "NYC"), "skills": ["python"]},
                {"name": "Bob",   "location": criteria.get("location", "NYC"), "skills": ["java"]},
                {"name": "Carol", "location": criteria.get("location", "NYC"), "skills": ["go"]},
            ]
            task = Task(payload.get("task_id"), "search", criteria)
            task.status = "completed"
            task.artifact = candidates
            response = Message(self.card.name, msg.sender, {"task": task.to_json()})
            return response.to_json()
        return super().handle_message(msg_json)

class FilterAgent(Agent):
    def __init__(self):
        super().__init__("FilterAgent", ["filter_candidates"])

    def handle_message(self, msg_json: str) -> str:
        msg = Message.from_json(msg_json)
        payload = msg.content
        if payload.get("action") == "filter":
            task = Task(payload.get("task_id"), "filter", payload)
            candidates = payload.get("candidates", [])
            skill = payload.get("skill")
            # Filter by skill
            filtered = [c for c in candidates if skill in c.get("skills", [])]
            task.status = "completed"
            task.artifact = filtered
            response = Message(self.card.name, msg.sender, {"task": task.to_json()})
            return response.to_json()
        return super().handle_message(msg_json)

class ScheduleAgent(Agent):
    def __init__(self):
        super().__init__("ScheduleAgent", ["schedule_interview"])

    def handle_message(self, msg_json: str) -> str:
        msg = Message.from_json(msg_json)
        payload = msg.content
        if payload.get("action") == "schedule":
            task = Task(payload.get("task_id"), "schedule", payload)
            candidates = payload.get("candidates", [])
            # Dummy schedule times
            schedule = {c["name"]: "2025-05-01 10:00" for c in candidates}
            task.status = "completed"
            task.artifact = schedule
            response = Message(self.card.name, msg.sender, {"task": task.to_json()})
            return response.to_json()
        return super().handle_message(msg_json)

# --- Orchestrator (Client) ---
def main():
    # Initialize agents
    agents = [SearchAgent(), FilterAgent(), ScheduleAgent()]

    # Capability discovery
    catalog: Dict[str, AgentCard] = {}
    for ag in agents:
        card = AgentCard.from_json(ag.advertise())
        catalog[card.name] = card

    # Step 1: Search candidates
    task1_id = str(uuid.uuid4())
    search_msg = Message("Orchestrator", "SearchAgent", {"action": "search", "task_id": task1_id, "criteria": {"location": "NYC"}})
    resp1 = agents[0].handle_message(search_msg.to_json())
    res1 = Message.from_json(resp1).content
    task1 = Task.from_json(res1["task"])
    candidates = task1.artifact
    print("Found candidates:", candidates)

    # Step 2: Filter by skill
    task2_id = str(uuid.uuid4())
    filter_msg = Message("Orchestrator", "FilterAgent", {"action": "filter", "task_id": task2_id, "candidates": candidates, "skill": "python"})
    resp2 = agents[1].handle_message(filter_msg.to_json())
    task2 = Task.from_json(Message.from_json(resp2).content["task"])
    filtered = task2.artifact
    print("Filtered candidates:", filtered)

    # Step 3: Schedule interviews
    task3_id = str(uuid.uuid4())
    sched_msg = Message("Orchestrator", "ScheduleAgent", {"action": "schedule", "task_id": task3_id, "candidates": filtered})
    resp3 = agents[2].handle_message(sched_msg.to_json())
    task3 = Task.from_json(Message.from_json(resp3).content["task"])
    schedule = task3.artifact
    print("Interview schedule:", schedule)

if __name__ == '__main__':
    main()