Upload 2 files

aworld/core/event/base.py

@@ -0,0 +1,185 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import abc
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Dict, Generic, TypeVar, List, Optional
+
+from pydantic import BaseModel
+
+from aworld.config.conf import ConfigDict
+from aworld.core.common import Config, Observation, ActionModel, TaskItem
+from aworld.core.context.base import Context
+
+
+class Constants:
+    AGENT = "agent"
+    TOOL = "tool"
+    TASK = "task"
+    OUTPUT = "output"
+    TOOL_CALLBACK = "tool_callback"
+
+
+class TopicType:
+    START = "__start"
+    FINISHED = "__finished"
+    OUTPUT = "__output"
+    ERROR = "__error"
+    RERUN = "__rerun"
+    HUMAN_CONFIRM = "__human_confirm"
+    # for dynamic subscribe
+    SUBSCRIBE_TOOL = "__subscribe_tool"
+    SUBSCRIBE_AGENT = "__subscribe_agent"
+
+
+DataType = TypeVar('DataType')
+
+
+@dataclass
+class Message(Generic[DataType]):
+    """The message structure for event transmission.
+
+    Each message has a unique ID, and the actual data is carried through the `payload` attribute,
+    peer to peer(p2p) message transmission is achieved by setting the `receiver`, and topic based
+    message transmission is achieved by setting the `topic`.
+
+    Specific message recognition and processing can be achieved through the type of `payload`
+    or by extending `Message`.
+    """
+    session_id: str
+    payload: Optional[DataType]
+    # Current caller
+    sender: str
+    # event type
+    category: str
+    # Next caller
+    receiver: str = None
+    # The previous caller
+    caller: str = None
+    id: str = uuid.uuid4().hex
+    priority: int = 0
+    # Topic of message
+    topic: str = None
+    headers: Dict[str, Any] = field(default_factory=dict)
+    timestamp: int = time.time()
+
+    def __post_init__(self):
+        context = self.headers.get("context")
+        if not context:
+            self.headers['context'] = Context.instance()
+
+    def __lt__(self, other: object) -> bool:
+        if not isinstance(other, Message):
+            raise RuntimeError
+        return self.priority < other.priority
+
+    def key(self) -> str:
+        category = self.category if self.category else ''
+        if self.topic:
+            return f'{category}_{self.topic}'
+        else:
+            return f'{category}_{self.sender if self.sender else ""}'
+
+    def is_error(self):
+        return self.topic == TopicType.ERROR
+
+    @property
+    def task_id(self):
+        return self.context.task_id
+
+    @property
+    def context(self) -> Context:
+        return self.headers.get('context')
+
+    @context.setter
+    def context(self, context: Context):
+        self.headers['context'] = context
+
+
+@dataclass
+class TaskEvent(Message[TaskItem]):
+    """Task message is oriented towards applications, can interact with third-party entities independently."""
+    category: str = 'task'
+
+
+@dataclass
+class AgentMessage(Message[Observation]):
+    """Agent event is oriented towards applications, can interact with third-party entities independently.
+
+    For example, `agent` event can interact with other agents through the A2A protocol.
+    """
+    category: str = 'agent'
+
+
+@dataclass
+class ToolMessage(Message[List[ActionModel]]):
+    """Tool event is oriented towards applications, can interact with third-party entities independently.
+
+    For example, `tool` event can interact with other tools through the MCP protocol.
+    """
+    category: str = 'tool'
+
+
+class Messageable(object):
+    """Top-level API for data reception, transmission and transformation."""
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self, conf: Config = None, **kwargs):
+        self.conf = conf
+        if isinstance(conf, Dict):
+            self.conf = ConfigDict(conf)
+        elif isinstance(conf, BaseModel):
+            # To add flexibility
+            self.conf = ConfigDict(conf.model_dump())
+
+    @abc.abstractmethod
+    async def send(self, message: Message, **kwargs):
+        """Send a message to the receiver.
+
+        Args:
+            message: Message structure that carries the data that needs to be processed.
+        """
+
+    @abc.abstractmethod
+    async def receive(self, message: Message, **kwargs):
+        """Receive a message from the sender.
+
+        Mainly used for request-driven (call), event-driven is generally handled using `Eventbus`.
+
+        Args:
+            message: Message structure that carries the data that needs to be processed.
+        """
+
+    async def transform(self, message: Message, **kwargs):
+        """Transforms a message into a standardized format  from the sender.
+
+        Args:
+            message: Message structure that carries the data that needs to be processed.
+        """
+
+
+class Recordable(Messageable):
+    """Top-level API for recording data."""
+
+    async def send(self, message: Message, **kwargs):
+        return await self.write(message, **kwargs)
+
+    async def receive(self, message: Message, **kwargs):
+        return await self.read(message, **kwargs)
+
+    @abc.abstractmethod
+    async def read(self, message: Message, **kwargs):
+        """Read a message from the store.
+
+        Args:
+            message: Message structure that carries the data that needs to be read.
+        """
+
+    @abc.abstractmethod
+    async def write(self, message: Message, **kwargs):
+        """Write a message to the store.
+
+        Args:
+            message: Message structure that carries the data that needs to be write.
+        """

aworld/core/event/event_bus.py

@@ -0,0 +1,150 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import abc
+from asyncio import Queue, PriorityQueue
+from inspect import isfunction
+from typing import Callable, Any, Dict, List
+
+from aworld.core.singleton import InheritanceSingleton
+
+from aworld.core.common import Config
+from aworld.core.event.base import Message, Messageable
+from aworld.logs.util import logger
+
+
+class Eventbus(Messageable, InheritanceSingleton):
+    __metaclass__ = abc.ABCMeta
+
+    def __init__(self, conf: Config = None, **kwargs):
+        super().__init__(conf, **kwargs)
+        # {event_type: {topic: [handler1, handler2]}}
+        self._subscribers: Dict[str, Dict[str, List[Callable[..., Any]]]] = {}
+
+        # {event_type, handler}
+        self._transformer: Dict[str, Callable[..., Any]] = {}
+
+    async def send(self, message: Message, **kwargs):
+        return await self.publish(message, **kwargs)
+
+    async def receive(self, message: Message, **kwargs):
+        return await self.consume(message)
+
+    async def publish(self, messages: Message, **kwargs):
+        """Publish a message, equals `send`."""
+
+    async def consume(self, message: Message, **kwargs):
+        """Consume the message queue."""
+
+    async def subscribe(self, event_type: str, topic: str, handler: Callable[..., Any], **kwargs):
+        """Subscribe the handler to the event type and the topic.
+
+        NOTE: The handler list is executed sequentially in the topic, the output
+              of the previous one is the input of the next one.
+
+        Args:
+            event_type: Type of events, fixed ones(task, agent, tool, error).
+            topic: Classify messages through the topic.
+            handler: Function of handle the event type and topic message.
+            kwargs:
+                - transformer: Whether it is a transformer handler.
+                - order: Handler order in the topic.
+        """
+
+    async def unsubscribe(self, event_type: str, topic: str, handler: Callable[..., Any], **kwargs):
+        """Unsubscribe the handler to the event type and the topic.
+
+        Args:
+            event_type: Type of events, fixed ones(task, agent, tool, error).
+            topic: Classify messages through the topic.
+            handler: Function of handle the event type and topic message.
+            kwargs:
+                - transformer: Whether it is a transformer handler.
+        """
+
+    def get_handlers(self, event_type: str) -> Dict[str, List[Callable[..., Any]]]:
+        return self._subscribers.get(event_type, {})
+
+    def get_topic_handlers(self, event_type: str, topic: str) -> List[Callable[..., Any]]:
+        return self._subscribers.get(event_type, {}).get(topic, [])
+
+    def get_transform_handlers(self, key: str) -> Callable[..., Any]:
+        return self._transformer.get(key, None)
+
+    def close(self):
+        pass
+
+
+class InMemoryEventbus(Eventbus):
+    def __init__(self, conf: Config = None, **kwargs):
+        super().__init__(conf, **kwargs)
+
+        # use asyncio Queue as default
+        # use asyncio Queue as default, isolation based on session_id
+        # self._message_queue: Queue = Queue()
+        self._message_queue: Dict[str, Queue] = {}
+
+    def wait_consume_size(self, id: str) -> int:
+        return self._message_queue.get(id, Queue()).qsize()
+
+    async def publish(self, message: Message, **kwargs):
+        queue = self._message_queue.get(message.task_id)
+        if not queue:
+            queue = PriorityQueue()
+            self._message_queue[message.task_id] = queue
+        await queue.put(message)
+
+    async def consume(self, message: Message, **kwargs):
+        return await self._message_queue.get(message.task_id, PriorityQueue()).get()
+
+    async def consume_nowait(self, message: Message):
+        return self._message_queue.get(message.task_id, PriorityQueue()).get_nowait()
+
+    async def done(self, id: str):
+        while not self._message_queue.get(id, PriorityQueue()).empty():
+            self._message_queue.get(id, PriorityQueue()).get_nowait()
+        self._message_queue.get(id, PriorityQueue()).task_done()
+
+    async def subscribe(self, event_type: str, topic: str, handler: Callable[..., Any], **kwargs):
+        if kwargs.get("transformer"):
+            if event_type in self._transformer:
+                logger.warning(f"{event_type} transform already subscribe.")
+                return
+
+            if isfunction(handler):
+                self._transformer[event_type] = handler
+            else:
+                logger.warning(f"{event_type} {topic} subscribe fail, handler {handler} is not a function.")
+            return
+
+        order = kwargs.get('order', 99999)
+
+        handlers = self._subscribers.get(event_type)
+        if not handlers:
+            self._subscribers[event_type] = {}
+        topic_handlers = self._subscribers[event_type].get(topic)
+        if not topic_handlers:
+            self._subscribers[event_type][topic] = []
+
+        if order >= len(self._subscribers[event_type][topic]):
+            self._subscribers[event_type][topic].append(handler)
+        else:
+            self._subscribers[event_type][topic].insert(order, handler)
+        logger.info(f"subscribe {event_type} {topic} {handler} success.")
+        logger.info(f"{self._subscribers}")
+
+    async def unsubscribe(self, event_type: str, topic: str, handler: Callable[..., Any], **kwargs):
+        if kwargs.get("transformer"):
+            if event_type not in self._transformer:
+                logger.warning(f"{event_type} transform not subscribe.")
+                return
+
+            self._transformer.pop(event_type, None)
+            return
+
+        if event_type not in self._subscribers:
+            logger.warning(f"{event_type} handler not register.")
+            return
+
+        handlers = self._subscribers[event_type]
+        topic_handlers: List = handlers.get(topic, [])
+        topic_handlers.remove(handler)