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

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()

pycityagent/simulation/agentgroup.py

@@ -16,11 +16,12 @@ from langchain_core.embeddings import Embeddings
 
 from ..agent import Agent, InstitutionAgent
 from ..economy.econ_client import EconomyClient
-from ..environment.simulator import Simulator
+from ..environment import Simulator
 from ..llm.llm import LLM
 from ..llm.llmconfig import LLMConfig
 from ..memory import FaissQuery, Memory
 from ..message import Messager
+from ..metrics import MlflowClient
 from ..utils import (DIALOG_SCHEMA, INSTITUTION_STATUS_SCHEMA, PROFILE_SCHEMA,
                      STATUS_SCHEMA, SURVEY_SCHEMA)
 
@@ -38,11 +39,14 @@ class AgentGroup:
             list[Callable[[], tuple[dict, dict, dict]]],
         ],
         config: dict,
+        exp_name: str,
         exp_id: str | UUID,
         enable_avro: bool,
         avro_path: Path,
         enable_pgsql: bool,
         pgsql_writer: ray.ObjectRef,
+        message_interceptor: ray.ObjectRef,
+        mlflow_run_id: str,
         embedding_model: Embeddings,
         logging_level: int,
         agent_config_file: Optional[Union[str, list[str]]] = None,
@@ -77,6 +81,18 @@ class AgentGroup:
             }
         if self.enable_pgsql:
             pass
+        # Mlflow
+        _mlflow_config = config.get("metric_request", {}).get("mlflow")
+        if _mlflow_config:
+            logger.info(f"-----Creating Mlflow client in AgentGroup {self._uuid} ...")
+            self.mlflow_client = MlflowClient(
+                config=_mlflow_config,
+                mlflow_run_name=f"{exp_name}_{1000*int(time.time())}",
+                experiment_name=exp_name,
+                run_id=mlflow_run_id,
+            )
+        else:
+            self.mlflow_client = None
 
         # prepare Messager
         if "mqtt" in config["simulator_request"]:
@@ -91,6 +107,7 @@ class AgentGroup:
 
         self.message_dispatch_task = None
         self._pgsql_writer = pgsql_writer
+        self._message_interceptor = message_interceptor
         self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
         self.initialized = False
         self.id2agent = {}
@@ -126,13 +143,9 @@ class AgentGroup:
             for j in range(number_of_agents_i):
                 memory_config_function_group_i = memory_config_function_group[i]
                 extra_attributes, profile, base = memory_config_function_group_i()
-                memory = Memory(
-                    config=extra_attributes, 
-                    profile=profile, 
-                    base=base
-                )
+                memory = Memory(config=extra_attributes, profile=profile, base=base)
                 agent = agent_class_i(
-                    name=f"{agent_class_i.__name__}_{i}",
+                    name=f"{agent_class_i.__name__}_{i}",  # type: ignore
                     memory=memory,
                     llm_client=self.llm,
                     economy_client=self.economy_client,
@@ -141,12 +154,16 @@ class AgentGroup:
                 agent.set_exp_id(self.exp_id)  # type: ignore
                 if self.messager is not None:
                     agent.set_messager(self.messager)
+                if self.mlflow_client is not None:
+                    agent.set_mlflow_client(self.mlflow_client)  # type: ignore
                 if self.enable_avro:
                     agent.set_avro_file(self.avro_file)  # type: ignore
                 if self.enable_pgsql:
                     agent.set_pgsql_writer(self._pgsql_writer)
                 if self.agent_config_file is not None and self.agent_config_file[i]:
                     agent.load_from_file(self.agent_config_file[i])
+                if self._message_interceptor is not None:
+                    agent.set_message_interceptor(self._message_interceptor)
                 self.agents.append(agent)
                 self.id2agent[agent._uuid] = agent
 
@@ -287,11 +304,13 @@ class AgentGroup:
             embedding_tasks = []
             for agent in self.agents:
                 embedding_tasks.append(agent.memory.initialize_embeddings())
-                agent.memory.set_search_components(self.faiss_query, self.embedding_model)
+                agent.memory.set_search_components(
+                    self.faiss_query, self.embedding_model
+                )
                 agent.memory.set_simulator(self.simulator)
             await asyncio.gather(*embedding_tasks)
             logger.debug(f"-----Embedding initialized in AgentGroup {self._uuid} ...")
-        
+
         self.initialized = True
         logger.debug(f"-----AgentGroup {self._uuid} initialized")
 
@@ -324,7 +343,9 @@ class AgentGroup:
                     filtered_uuids.append(agent._uuid)
         return filtered_uuids
 
-    async def gather(self, content: str, target_agent_uuids: Optional[list[str]] = None):
+    async def gather(
+        self, content: str, target_agent_uuids: Optional[list[str]] = None
+    ):
         logger.debug(f"-----Gathering {content} from all agents in group {self._uuid}")
         results = {}
         if target_agent_uuids is None:
@@ -522,6 +543,11 @@ class AgentGroup:
                     ]:
                         if key not in _status_dict:
                             _status_dict[key] = ""
+                    for key in [
+                        "friend_ids",
+                    ]:
+                        if key not in _status_dict:
+                            _status_dict[key] = []
                     _status_dict["created_at"] = _date_time
             else:
                 if not issubclass(type(self.agents[0]), InstitutionAgent):
@@ -537,10 +563,19 @@ class AgentGroup:
                             parent_id = position["lane_position"]["lane_id"]
                         else:
                             parent_id = -1
-                        hunger_satisfaction = await agent.status.get("hunger_satisfaction")
-                        energy_satisfaction = await agent.status.get("energy_satisfaction")
-                        safety_satisfaction = await agent.status.get("safety_satisfaction")
-                        social_satisfaction = await agent.status.get("social_satisfaction")
+                        hunger_satisfaction = await agent.status.get(
+                            "hunger_satisfaction"
+                        )
+                        energy_satisfaction = await agent.status.get(
+                            "energy_satisfaction"
+                        )
+                        safety_satisfaction = await agent.status.get(
+                            "safety_satisfaction"
+                        )
+                        social_satisfaction = await agent.status.get(
+                            "social_satisfaction"
+                        )
+                        friend_ids = await agent.status.get("friends")
                         action = await agent.status.get("current_step")
                         action = action["intention"]
                         _status_dict = {
@@ -550,6 +585,9 @@ class AgentGroup:
                             "lng": lng,
                             "lat": lat,
                             "parent_id": parent_id,
+                            "friend_ids": [
+                                str(_friend_id) for _friend_id in friend_ids
+                            ],
                             "action": action,
                             "hungry": hunger_satisfaction,
                             "tired": energy_satisfaction,
@@ -612,6 +650,10 @@ class AgentGroup:
                             employees = await agent.status.get("employees")
                         except:
                             employees = []
+                        try:
+                            friend_ids = await agent.status.get("friends")
+                        except:
+                            friend_ids = []
                         _status_dict = {
                             "id": agent._uuid,
                             "day": _day,
@@ -619,6 +661,9 @@ class AgentGroup:
                             "lng": lng,
                             "lat": lat,
                             "parent_id": parent_id,
+                            "friend_ids": [
+                                str(_friend_id) for _friend_id in friend_ids
+                            ],
                             "action": "",
                             "type": await agent.status.get("type"),
                             "nominal_gdp": nominal_gdp,
@@ -644,6 +689,7 @@ class AgentGroup:
                     "lng",
                     "lat",
                     "parent_id",
+                    "friend_ids",
                     "action",
                     "created_at",
                 ]
@@ -668,6 +714,7 @@ class AgentGroup:
             await asyncio.gather(*tasks)
         except Exception as e:
             import traceback
+
             logger.error(f"模拟器运行错误: {str(e)}\n{traceback.format_exc()}")
             raise RuntimeError(str(e)) from e
 
@@ -676,5 +723,6 @@ class AgentGroup:
             await self.save_status(day, t)
         except Exception as e:
             import traceback
+
             logger.error(f"模拟器运行错误: {str(e)}\n{traceback.format_exc()}")
             raise RuntimeError(str(e)) from e