Upload 2 files

aworld/events/manager.py

@@ -0,0 +1,115 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+from typing import Dict, Any, List, Callable
+
+from aworld.core.context.base import Context
+from aworld.core.event import eventbus
+from aworld.core.event.base import Constants, Message
+
+
+class EventManager:
+    """The event manager is now used to build an event bus instance and store the messages recently."""
+
+    def __init__(self, context: Context, **kwargs):
+        # use conf to build event bus instance
+        self.event_bus = eventbus
+        self.context = context
+        # Record events in memory for re-consume.
+        self.messages: Dict[str, List[Message]] = {'None': []}
+        self.max_len = kwargs.get('max_len', 1000)
+
+    async def emit(
+            self,
+            data: Any,
+            sender: str,
+            receiver: str = None,
+            topic: str = None,
+            session_id: str = None,
+            event_type: str = Constants.TASK
+    ):
+        """Send data to the event bus.
+
+        Args:
+            data: Message payload.
+            sender: The sender name of the message.
+            receiver: The receiver name of the message.
+            topic: The topic to which the message belongs.
+            session_id: Special session id.
+            event_type: Event type.
+        """
+        event = Message(
+            payload=data,
+            session_id=session_id if session_id else self.context.session_id,
+            sender=sender,
+            receiver=receiver,
+            topic=topic,
+            category=event_type,
+        )
+        return await self.emit_message(event)
+
+    async def emit_message(self, event: Message):
+        """Send the message to the event bus."""
+        key = event.key()
+        if key not in self.messages:
+            self.messages[key] = []
+        self.messages[key].append(event)
+        if len(self.messages) > self.max_len:
+            self.messages = self.messages[-self.max_len:]
+
+        await self.event_bus.publish(event)
+        return True
+
+    async def consume(self, nowait: bool = False):
+        msg = Message(session_id=self.context.session_id, sender="", category="", payload="")
+        msg.context = self.context
+        if nowait:
+            return await self.event_bus.consume_nowait(msg)
+        return await self.event_bus.consume(msg)
+
+    async def done(self):
+        await self.event_bus.done(self.context.task_id)
+
+    async def register(self, event_type: str, topic: str, handler: Callable[..., Any], **kwargs):
+        await self.event_bus.subscribe(event_type, topic, handler, **kwargs)
+
+    async def unregister(self, event_type: str, topic: str, handler: Callable[..., Any], **kwargs):
+        await self.event_bus.unsubscribe(event_type, topic, handler, **kwargs)
+
+    async def register_transformer(self, event_type: str, topic: str, handler: Callable[..., Any], **kwargs):
+        await self.event_bus.subscribe(event_type, topic, handler, transformer=True, **kwargs)
+
+    async def unregister_transformer(self, event_type: str, topic: str, handler: Callable[..., Any], **kwargs):
+        await self.event_bus.unsubscribe(event_type, topic, handler, transformer=True, **kwargs)
+
+    def messages_by_key(self, key: str) -> List[Message]:
+        return self.messages.get(key, [])
+
+    def messages_by_sender(self, sender: str, key: str):
+        results = []
+        for res in self.messages.get(key, []):
+            if res.sender == sender:
+                results.append(res)
+        return results
+
+    def messages_by_topic(self, topic: str, key: str):
+        results = []
+        for res in self.messages.get(key, []):
+            if res.topic == topic:
+                results.append(res)
+        return results
+
+    def session_messages(self, session_id: str) -> List[Message]:
+        return [m for k, msg in self.messages.items() for m in msg if m.session_id == session_id]
+
+    @staticmethod
+    def mark_valid(messages: List[Message]):
+        for msg in messages:
+            msg.is_valid = True
+
+    @staticmethod
+    def mark_invalid(messages: List[Message]):
+        for msg in messages:
+            msg.is_valid = False
+
+    def clear_messages(self):
+        self.messages = []

aworld/events/util.py

@@ -0,0 +1,56 @@
+# coding: utf-8
+# Copyright (c) 2025 inclusionAI.
+import asyncio
+from typing import Callable, Any
+
+from aworld.core.context.base import Context
+from aworld.core.event import eventbus
+from aworld.core.event.base import Message, Constants
+from aworld.events.manager import EventManager
+from aworld.utils.common import sync_exec
+
+
+def subscribe(key: str, category: str = None):
+    """Subscribe the special event to handle.
+
+    Examples:
+        >>> cate = Constants.TOOL or Constants.AGENT; key = "topic"
+        >>> @subscribe(category=cate, key=key)
+        >>> def example(message: Message) -> Message | None:
+        >>>     print("do something")
+
+    Args:
+         key: The index key of the handler.
+         category: Types of subscription events, the value is `agent` or `tool`.
+    """
+    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+        if category is None:
+            sync_exec(eventbus.subscribe, Constants.TOOL, key, func)
+            sync_exec(eventbus.subscribe, Constants.AGENT, key, func)
+        else:
+            sync_exec(eventbus.subscribe, category, key, func)
+        return func
+
+    return decorator
+
+async def _send_message(msg: Message) -> str:
+    context = msg.context
+    if not context:
+        context = Context()
+
+    event_mng = context.event_manager
+    if not event_mng:
+        event_mng = EventManager(context)
+
+    await event_mng.emit_message(msg)
+    return msg.id
+
+
+async def send_message(msg: Message) -> asyncio.Task:
+    """Utility function of send event.
+
+    Args:
+        msg: The content and meta information to be sent.
+    """
+    task = asyncio.create_task(_send_message(msg), name=msg.id)
+    return task