pycityagent 2.0.0a52__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a54__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/message/__init__.py

@@ -1,3 +1,10 @@
+from .message_interceptor import (MessageBlockBase, MessageBlockListenerBase,
+                                  MessageInterceptor)
 from .messager import Messager
 
-__all__ = ["Messager"]
+__all__ = [
+    "Messager",
+    "MessageBlockBase",
+    "MessageBlockListenerBase",
+    "MessageInterceptor",
+]

pycityagent/message/message_interceptor.py

@@ -0,0 +1,322 @@
+import asyncio
+import inspect
+import logging
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from collections.abc import Callable, Sequence
+from copy import deepcopy
+from typing import Any, Optional, Union
+
+import ray
+from ray.util.queue import Queue
+
+from ..llm import LLM, LLMConfig
+from ..utils.decorators import lock_decorator
+
+DEFAULT_ERROR_STRING = """
+From `{from_uuid}` To `{to_uuid}` abort due to block `{block_name}`
+"""
+
+logger = logging.getLogger("message_interceptor")
+
+
+class MessageBlockBase(ABC):
+    """
+    用于过滤的block
+    """
+
+    def __init__(
+        self,
+        name: str = "",
+    ) -> None:
+        self._name = name
+        self._llm = None
+        self._lock = asyncio.Lock()
+
+    @property
+    def llm(
+        self,
+    ) -> LLM:
+        if self._llm is None:
+            raise RuntimeError(f"LLM access before assignment, please `set_llm` first!")
+        return self._llm
+
+    @property
+    def name(
+        self,
+    ):
+        return self._name
+
+    @property
+    def has_llm(
+        self,
+    ) -> bool:
+        return self._llm is not None
+
+    @lock_decorator
+    async def set_llm(self, llm: LLM):
+        """
+        Set the llm_client of the block.
+        """
+        self._llm = llm
+
+    @lock_decorator
+    async def set_name(self, name: str):
+        """
+        Set the name of the block.
+        """
+        self._name = name
+
+    @lock_decorator
+    async def forward(
+        self,
+        from_uuid: str,
+        to_uuid: str,
+        msg: str,
+        violation_counts: dict[str, int],
+        black_list: list[tuple[str, str]],
+    ) -> tuple[bool, str]:
+        return True, ""
+
+
+@ray.remote
+class MessageInterceptor:
+    """
+    信息拦截器
+    """
+
+    def __init__(
+        self,
+        blocks: Optional[list[MessageBlockBase]] = None,
+        black_list: Optional[list[tuple[str, str]]] = None,
+        llm_config: Optional[dict] = None,
+        queue: Optional[Queue] = None,
+    ) -> None:
+        if blocks is not None:
+            self._blocks: list[MessageBlockBase] = blocks
+        else:
+            self._blocks: list[MessageBlockBase] = []
+        self._violation_counts: dict[str, int] = defaultdict(int)
+        if black_list is not None:
+            self._black_list: list[tuple[str, str]] = black_list
+        else:
+            self._black_list: list[tuple[str, str]] = (
+                []
+            )  # list[tuple(from_uuid, to_uuid)] `None` means forbidden for everyone.
+        if llm_config:
+            self._llm = LLM(LLMConfig(llm_config))
+        else:
+            self._llm = None
+        self._queue = queue
+        self._lock = asyncio.Lock()
+
+    @property
+    def llm(
+        self,
+    ) -> LLM:
+        if self._llm is None:
+            raise RuntimeError(f"LLM access before assignment, please `set_llm` first!")
+        return self._llm
+
+    @lock_decorator
+    async def blocks(
+        self,
+    ) -> list[MessageBlockBase]:
+        return self._blocks
+
+    @lock_decorator
+    async def set_llm(self, llm: LLM):
+        """
+        Set the llm_client of the block.
+        """
+        if self._llm is None:
+            self._llm = llm
+
+    @lock_decorator
+    async def violation_counts(
+        self,
+    ) -> dict[str, int]:
+        return deepcopy(self._violation_counts)
+
+    @property
+    def has_llm(
+        self,
+    ) -> bool:
+        return self._llm is not None
+
+    @lock_decorator
+    async def black_list(
+        self,
+    ) -> list[tuple[str, str]]:
+        return deepcopy(self._black_list)
+
+    @lock_decorator
+    async def add_to_black_list(
+        self, black_list: Union[list[tuple[str, str]], tuple[str, str]]
+    ):
+        if all(isinstance(s, str) for s in black_list):
+            # tuple[str,str]
+            _black_list = [black_list]
+        else:
+            _black_list = black_list
+        _black_list = [tuple(p) for p in _black_list]
+        self._black_list.extend(_black_list)  # type: ignore
+        self._black_list = list(set(self._black_list))
+
+    @property
+    def has_queue(
+        self,
+    ) -> bool:
+        return self._queue is not None
+
+    @lock_decorator
+    async def set_queue(self, queue: Queue):
+        """
+        Set the queue of the MessageInterceptor.
+        """
+        self._queue = queue
+
+    @lock_decorator
+    async def remove_from_black_list(
+        self, to_remove_black_list: Union[list[tuple[str, str]], tuple[str, str]]
+    ):
+        if all(isinstance(s, str) for s in to_remove_black_list):
+            # tuple[str,str]
+            _black_list = [to_remove_black_list]
+        else:
+            _black_list = to_remove_black_list
+        _black_list_set = {tuple(p) for p in _black_list}
+        self._black_list = [p for p in self._black_list if p not in _black_list_set]
+
+    @property
+    def queue(
+        self,
+    ) -> Queue:
+        if self._queue is None:
+            raise RuntimeError(
+                f"Queue access before assignment, please `set_queue` first!"
+            )
+        return self._queue
+
+    @lock_decorator
+    async def insert_block(self, block: MessageBlockBase, index: Optional[int] = 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:
+        if index is None:
+            index = -1
+        return self._blocks.pop(index)
+
+    @lock_decorator
+    async def set_black_list(
+        self, black_list: Union[list[tuple[str, str]], tuple[str, str]]
+    ):
+        if all(isinstance(s, str) for s in black_list):
+            # tuple[str,str]
+            _black_list = [black_list]
+        else:
+            _black_list = black_list
+        _black_list = [tuple(p) for p in _black_list]
+        self._black_list = list(set(_black_list))  # type: ignore
+
+    @lock_decorator
+    async def set_blocks(self, blocks: list[MessageBlockBase]):
+        self._blocks = blocks
+
+    @lock_decorator
+    async def forward(
+        self,
+        from_uuid: str,
+        to_uuid: str,
+        msg: str,
+    ):
+        for _block in self._blocks:
+            if not _block.has_llm and self.has_llm:
+                await _block.set_llm(self.llm)
+            func_params = inspect.signature(_block.forward).parameters
+            _args = {
+                "from_uuid": from_uuid,
+                "to_uuid": to_uuid,
+                "msg": msg,
+                "violation_counts": self._violation_counts,
+                "black_list": self._black_list,
+            }
+            _required_args = {k: v for k, v in _args.items() if k in func_params}
+            res = await _block.forward(**_required_args)
+            try:
+                is_valid, err = res
+            except TypeError as e:
+                is_valid: bool = res  # type:ignore
+                err = (
+                    DEFAULT_ERROR_STRING.format(
+                        from_uuid=from_uuid,
+                        to_uuid=to_uuid,
+                        block_name=f"{_block.__class__.__name__} `{_block.name}`",
+                    )
+                    if not is_valid
+                    else ""
+                )
+            if not is_valid:
+                if self.has_queue:
+                    logger.debug(f"put `{err}` into queue")
+                    await self.queue.put_async(err)  # type:ignore
+                self._violation_counts[from_uuid] += 1
+                print(self._black_list)
+                return False
+            else:
+                # valid
+                pass
+        print(self._black_list)
+        return True
+
+
+class MessageBlockListenerBase(ABC):
+    def __init__(
+        self,
+        save_queue_values: bool = False,
+        get_queue_period: float = 0.1,
+    ) -> None:
+        self._queue = None
+        self._lock = asyncio.Lock()
+        self._values_from_queue: list[Any] = []
+        self._save_queue_values = save_queue_values
+        self._get_queue_period = get_queue_period
+
+    @property
+    def queue(
+        self,
+    ) -> Queue:
+        if self._queue is None:
+            raise RuntimeError(
+                f"Queue access before assignment, please `set_queue` first!"
+            )
+        return self._queue
+
+    @property
+    def has_queue(
+        self,
+    ) -> bool:
+        return self._queue is not None
+
+    @lock_decorator
+    async def set_queue(self, queue: Queue):
+        """
+        Set the queue of the MessageBlockListenerBase.
+        """
+        self._queue = queue
+
+    @lock_decorator
+    async def forward(
+        self,
+    ):
+        while True:
+            if self.has_queue:
+                value = await self.queue.get_async()  # type: ignore
+                if self._save_queue_values:
+                    self._values_from_queue.append(value)
+                logger.debug(f"get `{value}` from queue")
+                # do something with the value
+            await asyncio.sleep(self._get_queue_period)

pycityagent/message/messager.py

@@ -1,19 +1,26 @@
 import asyncio
 import json
 import logging
-import math
-from typing import Any, List, Union
+from typing import Any, Optional, Union
 
 import ray
 from aiomqtt import Client
 
+from .message_interceptor import MessageInterceptor
+
 logger = logging.getLogger("pycityagent")
 
 
 @ray.remote
 class Messager:
     def __init__(
-        self, hostname: str, port: int = 1883, username=None, password=None, timeout=60
+        self,
+        hostname: str,
+        port: int = 1883,
+        username=None,
+        password=None,
+        timeout=60,
+        message_interceptor: Optional[ray.ObjectRef] = None,
     ):
         self.client = Client(
             hostname, port=port, username=username, password=password, timeout=timeout
@@ -21,6 +28,16 @@ class Messager:
         self.connected = False  # 是否已连接标志
         self.message_queue = asyncio.Queue()  # 用于存储接收到的消息
         self.receive_messages_task = None
+        self._message_interceptor = message_interceptor
+
+    @property
+    def message_interceptor(
+        self,
+    ) -> Union[None, ray.ObjectRef]:
+        return self._message_interceptor
+
+    def set_message_interceptor(self, message_interceptor: ray.ObjectRef):
+        self._message_interceptor = message_interceptor
 
     async def __aexit__(self, exc_type, exc_value, traceback):
         await self.stop()
@@ -50,8 +67,9 @@ class Messager:
         """检查是否成功连接到 Broker"""
         return self.connected
 
+    # TODO:add message interceptor
     async def subscribe(
-        self, topics: Union[str, List[str]], agents: Union[Any, List[Any]]
+        self, topics: Union[str, list[str]], agents: Union[Any, list[Any]]
     ):
         if not await self.is_connected():
             logger.error(
@@ -62,7 +80,7 @@ class Messager:
             topics = [topics]
         if not isinstance(agents, list):
             agents = [agents]
-        await self.client.subscribe(topics, qos=1)
+        await self.client.subscribe(topics, qos=1)  # type: ignore
 
     async def receive_messages(self):
         """监听并将消息存入队列"""
@@ -76,11 +94,26 @@ class Messager:
             messages.append(await self.message_queue.get())
         return messages
 
-    async def send_message(self, topic: str, payload: dict):
+    async def send_message(
+        self,
+        topic: str,
+        payload: dict,
+        from_uuid: Optional[str] = None,
+        to_uuid: Optional[str] = None,
+    ):
         """通过 Messager 发送消息"""
         message = json.dumps(payload, default=str)
-        await self.client.publish(topic=topic, payload=message, qos=1)
-        logger.info(f"Message sent to {topic}: {message}")
+        interceptor = self.message_interceptor
+        is_valid: bool = True
+        if interceptor is not None and (from_uuid is not None and to_uuid is not None):
+            is_valid = await interceptor.forward.remote(  # type:ignore
+                from_uuid, to_uuid, message
+            )
+        if is_valid:
+            await self.client.publish(topic=topic, payload=message, qos=1)
+            logger.info(f"Message sent to {topic}: {message}")
+        else:
+            logger.info(f"Message not sent to {topic}: {message} due to interceptor")
 
     async def start_listening(self):
         """启动消息监听任务"""
@@ -92,7 +125,5 @@ class Messager:
     async def stop(self):
         assert self.receive_messages_task is not None
         self.receive_messages_task.cancel()
-        await asyncio.gather(
-            self.receive_messages_task, return_exceptions=True
-        )
+        await asyncio.gather(self.receive_messages_task, return_exceptions=True)
         await self.disconnect()