RPC capability

[ahilloffis]

Communication between two agents often times contain request-reply structures, as also shown in the tutorial.

It would be nice to have an inbuilt way RPCs, which could also facilitate await. This way, the control flow is easier to track.

[ahilloffis]

Just for you, maybe for an internal discussion, our current implementation of such an RPC structure:

import uuid
import asyncio
import logging
from typing import Dict, Any, Coroutine, Set
from enum import Enum
from dataclasses import dataclass

from mango import Role, RoleContext, AgentAddress, sender_addr, json_serializable

logging.basicConfig(level=logging.INFO)

LOGGER = logging.getLogger(__name__)

@json_serializable
@dataclass
class RequestMessage:
    request: Any

@json_serializable
@dataclass
class ResponseMessage:
    content: Any

class Response:
    corr_id: str
    future: asyncio.Future
    request: Any
    content: Any

    def __init__(self, request: Any = None, future: asyncio.Future = None, corr_id: str = None):
        self.request = request
        self.future = future if future is not None else asyncio.Future()
        self.corr_id = corr_id if corr_id is not None else str(uuid.uuid4())

class RPCContainer:
    class SCHEDULER(Enum):
        ASYNCIO = 1
        MANGO = 2

    def __init__(self, context: RoleContext, timeout=None):
        self.context = context
        self._requests_container: Dict[str, Response] = {}
        self.timeout = timeout
        self._registered_roles: Set[int] = {}
        
    def handle_response(self, response_message: ResponseMessage, meta):
        corr_id = meta["corr_id"]
        response = self._requests_container.pop(corr_id)
        response.content = response_message.content
        response.future.set_result(response.content)

    async def send_rpc(self, src: Role, message: Any, receiver_addr: AgentAddress, timeout: float=None, scheduler: SCHEDULER=SCHEDULER.ASYNCIO):
        LOGGER.debug(f"Send rpc to {receiver_addr}")
        timeout = self.timeout if timeout is None else timeout
        if id(src) not in self._registered_roles:
            src.context.subscribe_message(src, self.handle_response, lambda content, meta: isinstance(content, ResponseMessage) and meta.get("corr_id", "") in self._requests_container)

        response = Response(message)
        corr_id = response.corr_id
        self._requests_container[corr_id] = response
        if timeout is not None and timeout > 0:
            def at_timeout():
                if not response.future.done() and not response.future.cancelled():
                    LOGGER.info(f"Timeout of rpc of correlation id {corr_id}")
                    response.future.set_exception(TimeoutError(f"Request not successfull after {timeout} seconds for correlation id {corr_id}."))
            
            match scheduler:
                case RPCContainer.SCHEDULER.MANGO:
                    async def at_timeout_async():
                        at_timeout()
                    self.context.schedule_timestamp_task(at_timeout_async(), self.context.current_timestamp+timeout)
                case RPCContainer.SCHEDULER.ASYNCIO: # ToDo: Just use asyncio.wait_for(coro, timeout)
                    asyncio.get_event_loop().call_later(timeout, at_timeout)
                case _:
                    raise AssertionError("Scheduler not found") # ToDo: Allow additional schedulers
            
        await self.context.send_message(content=RequestMessage(message), receiver_addr=receiver_addr, corr_id=corr_id)
        LOGGER.debug(f"Await rpc for correlation id {corr_id}")
        return await response.future
    
    async def handle_request_async(self, coro, content, meta):
        response = await coro(content, meta)
        if response is not None:
            if not isinstance(response, ResponseMessage):
                response = ResponseMessage(response)
            self.context.schedule_instant_message(response, sender_addr(meta), corr_id=meta["corr_id"])

    def handle_request(self, coro, content, meta):
        LOGGER.info(f"Got request for {meta["corr_id"]}: Will schedule instant task")
        self.context.schedule_instant_task(self.handle_request_async(coro, content.request, meta))

    def register_response(self, coro: Coroutine, role: Role, request_class):
        role.context.subscribe_message(role,
                                       lambda content, meta: self.handle_request(coro, content, meta), # Wrapper for receiving the data
                                       lambda content, meta: isinstance(content, RequestMessage) and isinstance(content.request, request_class) # Filter for request response pattern
                                       )
        # ToDo: Think about making a more general subscription and then handle request in here instead, to also handle invalid requests

And usage:

rpc = RPCContainer(self.context) # rpc container should be part of the agent and not initialized every time
msg = MyCustomRequest(data)
try:
    response = await rpc.send_rpc(src=self, message=msg, receiver_addr=self.history_agent[key])
    LOGGER.info(f"Answer: {response}")
except TimeoutError as e:
    LOGGER.info(f"Got a timeout error {e}")

If MQTT is used, version 5 now also supports extra options that can be used for RPCs.

[rcschrg]

As a quick note for myself, this is very similar to the tracked messages we implemented in Mango.jl, so maybe we should just transfer this feature to mango-python.

[maurerle]

I do also have one or two use cases where this would be handy.
Currently, I have a feature to await a response to a message implemented in the AgentRole itself.

Regarding the correlation id - I think that the FIPA ACL reply_with and in_reply_to would be beneficial, but I also do see arguments for not using FIPA ACL for this feature..

[ahilloffis]

I agree that correlation id may be the wrong parameter name, as it is borrowed from AMQP/MQTT RPC scheme.

[rcschrg]

#174 will implement this. We will also ship a possibility to handle messages in an async context. It will also not be limited to roles, generally we want to align agents and roles such that the decision to use roles as modeling approach is only based on the architectural implications rather than the available feature set.

from mango import Role, sender_addr

class Coordinator(Role):
    async def poll_prices(self, sellers):
        replies = await self.context.gather(
            PriceRequest(),
            sellers,                 # AgentAddress or iterable of AgentAddress
            reply_type=PriceOffer,   # (optional) ignore any non-PriceOffer traffic
            timeout=2.0,             # (optional) scheduler-clock seconds
        )
        # replies: {AgentAddress: PriceOffer}
        best = min(replies.values(), key=lambda offer: offer.price)
        return best

class Responder(Role):
    @on_message(PriceRequest)
    async def on_price_request(self, content, meta):
       ...