protocols/a2a_protocol.py

"""
Agent-to-Agent (A2A) protocol: direct async calls between agent methods.
Used when agents need to request data synchronously from another agent
rather than through the event queue.
"""
from typing import Any, Callable, Dict, Optional
from utils.logger import get_logger

logger = get_logger("a2a_protocol")


class A2AProtocol:
    """Registry-based direct agent-to-agent call protocol."""

    def __init__(self):
        self._registry: Dict[str, Callable] = {}

    def register(self, agent_name: str, handler: Callable) -> None:
        self._registry[agent_name] = handler
        logger.debug(f"[A2A] Registered handler for '{agent_name}'")

    async def call(self, target: str, method: str, **kwargs) -> Any:
        handler = self._registry.get(target)
        if handler is None:
            raise KeyError(f"[A2A] No handler registered for '{target}'")
        result = handler(method, **kwargs)
        if hasattr(result, "__await__"):
            return await result
        return result

    def is_registered(self, agent_name: str) -> bool:
        return agent_name in self._registry