pycityagent 2.0.0a65__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a67__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/message/message_interceptor.py

@@ -19,6 +19,12 @@ From `{from_uuid}` To `{to_uuid}` abort due to block `{block_name}`
 
 logger = logging.getLogger("message_interceptor")
 
+__all__ = [
+    "MessageBlockBase",
+    "MessageInterceptor",
+    "MessageBlockListenerBase",
+]
+
 
 class MessageBlockBase(ABC):
     """
@@ -82,7 +88,7 @@ class MessageBlockBase(ABC):
 @ray.remote
 class MessageInterceptor:
     """
-    信息拦截器
+    A class to intercept and process messages based on configured rules.
     """
 
     def __init__(
@@ -92,6 +98,15 @@ class MessageInterceptor:
         llm_config: Optional[dict] = None,
         queue: Optional[Queue] = None,
     ) -> None:
+        """
+        Initialize the MessageInterceptor with optional configuration.
+
+        - **Args**:
+            - `blocks` (Optional[list[MessageBlockBase]], optional): Initial list of message interception rules. Defaults to an empty list.
+            - `black_list` (Optional[list[tuple[str, str]]], optional): Initial blacklist of communication pairs. Defaults to an empty list.
+            - `llm_config` (Optional[dict], optional): Configuration dictionary for initializing the LLM instance. Defaults to None.
+            - `queue` (Optional[Queue], optional): Queue for message processing. Defaults to None.
+        """
         if blocks is not None:
             self._blocks: list[MessageBlockBase] = blocks
         else:
@@ -114,6 +129,18 @@ class MessageInterceptor:
     def llm(
         self,
     ) -> LLM:
+        """
+        Access the Large Language Model instance.
+
+        - **Description**:
+            - Provides access to the internal LLM instance. Raises an error if accessed before assignment.
+
+        - **Raises**:
+            - `RuntimeError`: If accessed before setting the LLM.
+
+        - **Returns**:
+            - `LLM`: The Large Language Model instance.
+        """
         if self._llm is None:
             raise RuntimeError(f"LLM access before assignment, please `set_llm` first!")
         return self._llm
@@ -122,12 +149,30 @@ class MessageInterceptor:
     async def blocks(
         self,
     ) -> list[MessageBlockBase]:
+        """
+        Retrieve the message interception rules.
+
+        - **Description**:
+            - Returns a copy of the current list of message interception rules.
+
+        - **Returns**:
+            - `list[MessageBlockBase]`: The list of message interception rules.
+        """
         return self._blocks
 
     @lock_decorator
     async def set_llm(self, llm: LLM):
         """
-        Set the llm_client of the block.
+        Set the Large Language Model instance.
+
+        - **Description**:
+            - Updates the internal LLM instance used for message processing.
+
+        - **Args**:
+            - `llm` (LLM): The LLM instance to be set.
+
+        - **Returns**:
+            - `None`
         """
         if self._llm is None:
             self._llm = llm
@@ -136,24 +181,64 @@ class MessageInterceptor:
     async def violation_counts(
         self,
     ) -> dict[str, int]:
+        """
+        Retrieve the violation counts.
+
+        - **Description**:
+            - Returns a deep copy of the violation counts to prevent external modification of the original data.
+
+        - **Returns**:
+            - `dict[str, int]`: The dictionary of violation counts.
+        """
         return deepcopy(self._violation_counts)
 
     @property
     def has_llm(
         self,
     ) -> bool:
+        """
+        Check if a Large Language Model is configured.
+
+        - **Description**:
+            - Confirms whether a Large Language Model instance has been set.
+
+        - **Returns**:
+            - `bool`: True if an LLM is set, otherwise False.
+        """
         return self._llm is not None
 
     @lock_decorator
     async def black_list(
         self,
     ) -> list[tuple[str, str]]:
+        """
+        Retrieve the blacklist.
+
+        - **Description**:
+            - Returns a deep copy of the current blacklist to protect the original data from external modifications.
+
+        - **Returns**:
+            - `list[tuple[str, str]]`: The blacklist.
+        """
         return deepcopy(self._black_list)
 
     @lock_decorator
     async def add_to_black_list(
         self, black_list: Union[list[tuple[str, str]], tuple[str, str]]
     ):
+        """
+        Add entries to the blacklist.
+
+        - **Description**:
+            - Adds one or more entries to the blacklist, ensuring each entry's uniqueness.
+
+        - **Args**:
+            - `black_list` (Union[list[tuple[str, str]], tuple[str, str]]):
+                Can be a single tuple or a list of tuples indicating the entries to add to the blacklist.
+
+        - **Returns**:
+            - `None`
+        """
         if all(isinstance(s, str) for s in black_list):
             # tuple[str,str]
             _black_list = [black_list]
@@ -167,12 +252,30 @@ class MessageInterceptor:
     def has_queue(
         self,
     ) -> bool:
+        """
+        Check if a queue is configured.
+
+        - **Description**:
+            - Confirms whether a queue instance has been set for message processing.
+
+        - **Returns**:
+            - `bool`: True if a queue is set, otherwise False.
+        """
         return self._queue is not None
 
     @lock_decorator
     async def set_queue(self, queue: Queue):
         """
         Set the queue of the MessageInterceptor.
+
+        - **Description**:
+            - Assigns a queue to the MessageInterceptor for asynchronous message handling.
+
+        - **Args**:
+            - `queue` (Queue): The queue instance to be set.
+
+        - **Returns**:
+            - `None`
         """
         self._queue = queue
 
@@ -180,6 +283,19 @@ class MessageInterceptor:
     async def remove_from_black_list(
         self, to_remove_black_list: Union[list[tuple[str, str]], tuple[str, str]]
     ):
+        """
+        Remove entries from the blacklist.
+
+        - **Description**:
+            - Removes one or more entries from the blacklist, ensuring each entry's removal.
+
+        - **Args**:
+            - `to_remove_black_list` (Union[list[tuple[str, str]], tuple[str, str]]):
+                Can be a single tuple or a list of tuples indicating the entries to remove from the blacklist.
+
+        - **Returns**:
+            - `None`
+        """
         if all(isinstance(s, str) for s in to_remove_black_list):
             # tuple[str,str]
             _black_list = [to_remove_black_list]
@@ -192,6 +308,18 @@ class MessageInterceptor:
     def queue(
         self,
     ) -> Queue:
+        """
+        Access the queue used for message processing.
+
+        - **Description**:
+            - Provides access to the internal queue. Raises an error if accessed before assignment.
+
+        - **Raises**:
+            - `RuntimeError`: If accessed before setting the queue.
+
+        - **Returns**:
+            - `Queue`: The queue instance.
+        """
         if self._queue is None:
             raise RuntimeError(
                 f"Queue access before assignment, please `set_queue` first!"
@@ -200,12 +328,37 @@ class MessageInterceptor:
 
     @lock_decorator
     async def insert_block(self, block: MessageBlockBase, index: Optional[int] = None):
+        """
+        Insert a message block into the blocks list at a specified position.
+
+        - **Description**:
+            - Inserts a new message interception rule into the list at the specified index or appends it if no index is provided.
+
+        - **Args**:
+            - `block` (MessageBlockBase): The message block to insert.
+            - `index` (Optional[int], optional): The position at which to insert the block. Defaults to appending at the end.
+
+        - **Returns**:
+            - `None`
+        """
         if index is None:
             index = len(self._blocks)
         self._blocks.insert(index, block)
 
     @lock_decorator
     async def pop_block(self, index: Optional[int] = None) -> MessageBlockBase:
+        """
+        Remove and return a message block from the blocks list.
+
+        - **Description**:
+            - Removes and returns the message block at the specified index or the last one if no index is provided.
+
+        - **Args**:
+            - `index` (Optional[int], optional): The position of the block to remove. Defaults to removing the last element.
+
+        - **Returns**:
+            - `MessageBlockBase`: The removed message block.
+        """
         if index is None:
             index = -1
         return self._blocks.pop(index)
@@ -214,6 +367,19 @@ class MessageInterceptor:
     async def set_black_list(
         self, black_list: Union[list[tuple[str, str]], tuple[str, str]]
     ):
+        """
+        Set the blacklist with new entries.
+
+        - **Description**:
+            - Updates the blacklist with new entries, ensuring each entry's uniqueness.
+
+        - **Args**:
+            - `black_list` (Union[list[tuple[str, str]], tuple[str, str]]):
+                Can be a single tuple or a list of tuples indicating the new blacklist entries.
+
+        - **Returns**:
+            - `None`
+        """
         if all(isinstance(s, str) for s in black_list):
             # tuple[str,str]
             _black_list = [black_list]
@@ -224,6 +390,18 @@ class MessageInterceptor:
 
     @lock_decorator
     async def set_blocks(self, blocks: list[MessageBlockBase]):
+        """
+        Replace the current blocks list with a new list of message blocks.
+
+        - **Description**:
+            - Sets a new list of message interception rules, replacing the existing list.
+
+        - **Args**:
+            - `blocks` (list[MessageBlockBase]): The new list of message blocks to set.
+
+        - **Returns**:
+            - `None`
+        """
         self._blocks = blocks
 
     @lock_decorator
@@ -233,6 +411,20 @@ class MessageInterceptor:
         to_uuid: str,
         msg: str,
     ):
+        """
+        Forward a message through all message blocks.
+
+        - **Description**:
+            - Processes a message by passing it through all configured message blocks. Each block can modify the message or prevent its forwarding based on implemented logic.
+
+        - **Args**:
+            - `from_uuid` (str): The UUID of the sender.
+            - `to_uuid` (str): The UUID of the recipient.
+            - `msg` (str): The message content to forward.
+
+        - **Returns**:
+            - `bool`: True if the message was successfully processed by all blocks, otherwise False.
+        """
         for _block in self._blocks:
             if not _block.has_llm and self.has_llm:
                 await _block.set_llm(self.llm)
@@ -274,11 +466,29 @@ class MessageInterceptor:
 
 
 class MessageBlockListenerBase(ABC):
+    """
+    Base class for message block listeners that can listen to a queue and process items.
+
+    - **Attributes**:
+        - `_queue` (Optional[Queue]): Queue from which the listener retrieves items.
+        - `_lock` (asyncio.Lock): Lock for thread-safe access in asynchronous environments.
+        - `_values_from_queue` (list[Any]): List of values retrieved from the queue if saving is enabled.
+        - `_save_queue_values` (bool): Flag indicating whether to save values from the queue.
+        - `_get_queue_period` (float): Period in seconds between queue retrieval attempts.
+    """
+
     def __init__(
         self,
         save_queue_values: bool = False,
         get_queue_period: float = 0.1,
     ) -> None:
+        """
+        Initialize the MessageBlockListenerBase with optional configuration.
+
+        - **Args**:
+            - `save_queue_values` (bool, optional): Whether to save values retrieved from the queue. Defaults to False.
+            - `get_queue_period` (float, optional): Time period in seconds between queue retrieval attempts. Defaults to 0.1.
+        """
         self._queue = None
         self._lock = asyncio.Lock()
         self._values_from_queue: list[Any] = []
@@ -289,6 +499,18 @@ class MessageBlockListenerBase(ABC):
     def queue(
         self,
     ) -> Queue:
+        """
+        Access the queue used by the listener.
+
+        - **Description**:
+            - Provides access to the internal queue. Raises an error if accessed before assignment.
+
+        - **Raises**:
+            - `RuntimeError`: If accessed before setting the queue.
+
+        - **Returns**:
+            - `Queue`: The queue instance.
+        """
         if self._queue is None:
             raise RuntimeError(
                 f"Queue access before assignment, please `set_queue` first!"
@@ -299,12 +521,30 @@ class MessageBlockListenerBase(ABC):
     def has_queue(
         self,
     ) -> bool:
+        """
+        Check if a queue is configured.
+
+        - **Description**:
+            - Confirms whether a queue instance has been set for the listener.
+
+        - **Returns**:
+            - `b
+        """
         return self._queue is not None
 
     @lock_decorator
     async def set_queue(self, queue: Queue):
         """
-        Set the queue of the MessageBlockListenerBase.
+        Set the queue for the listener.
+
+        - **Description**:
+            - Assigns a queue to the listener for asynchronous item processing.
+
+        - **Args**:
+            - `queue` (Queue): The queue instance to be set.
+
+        - **Returns**:
+            - `None`
         """
         self._queue = queue
 
@@ -312,6 +552,16 @@ class MessageBlockListenerBase(ABC):
     async def forward(
         self,
     ):
+        """
+        Continuously retrieve items from the queue and process them.
+
+        - **Description**:
+            - Listens to the queue, retrieves items at intervals defined by `_get_queue_period`,
+              and processes each item. If `_save_queue_values` is True, it saves the items in `_values_from_queue`.
+
+        - **Returns**:
+            - `None`
+        """
         while True:
             if self.has_queue:
                 value = await self.queue.get_async()  # type: ignore

pycityagent/message/messager.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+import time
 from typing import Any, Optional, Union
 
 import ray
@@ -8,11 +9,26 @@ from aiomqtt import Client
 
 from .message_interceptor import MessageInterceptor
 
+__all__ = [
+    "Messager",
+]
+
 logger = logging.getLogger("pycityagent")
 
 
 @ray.remote
 class Messager:
+    """
+    A class to manage message sending and receiving using an MQTT protocol.
+
+    - **Attributes**:
+        - `client` (Client): An instance of the MQTT client.
+        - `connected` (bool): Indicates whether the connection to the broker is established.
+        - `message_queue` (asyncio.Queue): Queue for storing received messages.
+        - `receive_messages_task` (Optional[Task]): Task for listening to incoming messages.
+        - `_message_interceptor` (Optional[ray.ObjectRef]): Reference to a remote message interceptor object.
+    """
+
     def __init__(
         self,
         hostname: str,
@@ -22,6 +38,17 @@ class Messager:
         timeout=60,
         message_interceptor: Optional[ray.ObjectRef] = None,
     ):
+        """
+        Initialize the Messager with connection parameters.
+
+        - **Args**:
+            - `hostname` (str): The hostname or IP address of the MQTT broker.
+            - `port` (int, optional): Port number of the MQTT broker. Defaults to 1883.
+            - `username` (str, optional): Username for broker authentication.
+            - `password` (str, optional): Password for broker authentication.
+            - `timeout` (int, optional): Connection timeout in seconds. Defaults to 60.
+            - `message_interceptor` (Optional[ray.ObjectRef], optional): Reference to a message interceptor object.
+        """
         self.client = Client(
             hostname, port=port, username=username, password=password, timeout=timeout
         )
@@ -29,23 +56,61 @@ class Messager:
         self.message_queue = asyncio.Queue()  # 用于存储接收到的消息
         self.receive_messages_task = None
         self._message_interceptor = message_interceptor
+        self._log_list = []
 
     @property
     def message_interceptor(
         self,
     ) -> Union[None, ray.ObjectRef]:
+        """
+        Access the message interceptor reference.
+
+        - **Returns**:
+            - `Union[None, ray.ObjectRef]`: The message interceptor reference.
+        """
         return self._message_interceptor
+    
+    def get_log_list(self):
+        return self._log_list
+    
+    def clear_log_list(self):
+        self._log_list = []
 
     def set_message_interceptor(self, message_interceptor: ray.ObjectRef):
+        """
+        Set the message interceptor reference.
+
+        - **Args**:
+            - `message_interceptor` (ray.ObjectRef): The message interceptor reference to be set.
+        """
         self._message_interceptor = message_interceptor
 
     async def __aexit__(self, exc_type, exc_value, traceback):
+        """
+        Asynchronous exit method to ensure proper cleanup when used in an async context manager.
+
+        - **Args**:
+            - `exc_type`, `exc_value`, `traceback`: Exception information if an exception occurred.
+        """
         await self.stop()
 
     async def ping(self):
+        """
+        Send a ping message to the MQTT broker.
+
+        - **Description**:
+            - Publishes a 'ping' message on the 'ping' topic with QoS level 1.
+        """
         await self.client.publish(topic="ping", payload="ping", qos=1)
 
     async def connect(self):
+        """
+        Attempt to connect to the MQTT broker up to three times.
+
+        - **Description**:
+            - Tries to establish a connection to the MQTT broker. Retries up to three times with delays between attempts.
+            - Logs success or failure accordingly.
+        """
         for i in range(3):
             try:
                 await self.client.__aenter__()
@@ -59,18 +124,35 @@ class Messager:
         logger.error("All connection attempts failed.")
 
     async def disconnect(self):
+        """
+        Disconnect from the MQTT broker.
+
+        - **Description**:
+            - Closes the connection to the MQTT broker and logs the disconnection.
+        """
         await self.client.__aexit__(None, None, None)
         self.connected = False
         logger.info("Disconnected from MQTT Broker")
 
     async def is_connected(self):
-        """检查是否成功连接到 Broker"""
+        """
+        Check if the connection to the broker is established.
+
+        - **Returns**:
+            - `bool`: True if connected, otherwise False.
+        """
         return self.connected
 
-    # TODO:add message interceptor
     async def subscribe(
         self, topics: Union[str, list[str]], agents: Union[Any, list[Any]]
     ):
+        """
+        Subscribe to one or more MQTT topics.
+
+        - **Args**:
+            - `topics` (Union[str, list[str]]): Topic or list of topics to subscribe to.
+            - `agents` (Union[Any, list[Any]]): Agents or list of agents associated with the subscription.
+        """
         if not await self.is_connected():
             logger.error(
                 f"Cannot subscribe to {topics} because not connected to the Broker."
@@ -83,12 +165,22 @@ class Messager:
         await self.client.subscribe(topics, qos=1)  # type: ignore
 
     async def receive_messages(self):
-        """监听并将消息存入队列"""
+        """
+        Listen for incoming messages and store them in the queue.
+
+        - **Description**:
+            - Continuously listens for incoming messages and puts them into the message queue.
+        """
         async for message in self.client.messages:
             await self.message_queue.put(message)
 
     async def fetch_messages(self):
-        """从队列中批量获取消息"""
+        """
+        Retrieve all messages currently in the queue.
+
+        - **Returns**:
+            - `list[Any]`: List of messages retrieved from the queue.
+        """
         messages = []
         while not self.message_queue.empty():
             messages.append(await self.message_queue.get())
@@ -101,7 +193,28 @@ class Messager:
         from_uuid: Optional[str] = None,
         to_uuid: Optional[str] = None,
     ):
-        """通过 Messager 发送消息"""
+        """
+        Send a message through the MQTT broker.
+
+        - **Args**:
+            - `topic` (str): Topic to which the message should be published.
+            - `payload` (dict): Payload of the message to send.
+            - `from_uuid` (Optional[str], optional): UUID of the sender. Required for interception.
+            - `to_uuid` (Optional[str], optional): UUID of the recipient. Required for interception.
+
+        - **Description**:
+            - Serializes the payload to JSON, checks it against the message interceptor (if any),
+              and publishes the message to the specified topic if valid.
+        """
+        start_time = time.time()
+        log = {
+            "topic": topic,
+            "payload": payload,
+            "from_uuid": from_uuid,
+            "to_uuid": to_uuid,
+            "start_time": start_time,
+            "consumption": 0
+        }
         message = json.dumps(payload, default=str)
         interceptor = self.message_interceptor
         is_valid: bool = True
@@ -114,15 +227,29 @@ class Messager:
             logger.info(f"Message sent to {topic}: {message}")
         else:
             logger.info(f"Message not sent to {topic}: {message} due to interceptor")
+        log["consumption"] = time.time() - start_time
+        self._log_list.append(log)
 
     async def start_listening(self):
-        """启动消息监听任务"""
+        """
+        Start the task for listening to incoming messages.
+
+        - **Description**:
+            - Starts a task that listens for incoming messages and puts them into the queue.
+            - Only starts the task if the connection to the broker is active.
+        """
         if await self.is_connected():
             self.receive_messages_task = asyncio.create_task(self.receive_messages())
         else:
             logger.error("Cannot start listening because not connected to the Broker.")
 
     async def stop(self):
+        """
+        Stop the listener and disconnect from the MQTT broker.
+
+        - **Description**:
+            - Cancels the receive_messages_task and ensures the MQTT broker connection is closed.
+        """
         assert self.receive_messages_task is not None
         self.receive_messages_task.cancel()
         await asyncio.gather(self.receive_messages_task, return_exceptions=True)