pycityagent 2.0.0a40__tar.gz → 2.0.0a41__tar.gz

{pycityagent-2.0.0a40 → pycityagent-2.0.0a41}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a40
+Version: 2.0.0a41
 Summary: LLM-based城市环境agent构建库
 License: MIT
 Author: Yuwei Yan

{pycityagent-2.0.0a40 → pycityagent-2.0.0a41}/pycityagent/agent.py

@@ -464,7 +464,7 @@ class Agent(ABC):
         if self._messager is None:
             raise RuntimeError("Messager is not set")
         topic = f"exps/{self._exp_id}/agents/{to_agent_uuid}/{sub_topic}"
-        await self._messager.send_message(topic, payload)
+        await self._messager.send_message.remote(topic, payload)
 
     async def send_message_to_agent(
         self, to_agent_uuid: str, content: str, type: str = "social"

{pycityagent-2.0.0a40 → pycityagent-2.0.0a41}/pycityagent/message/messager.py

@@ -4,32 +4,37 @@ import logging
 import math
 from typing import Any, List, Union
 from aiomqtt import Client
+import ray
 
 logger = logging.getLogger("pycityagent")
 
+@ray.remote
 class Messager:
     def __init__(
-        self, hostname:str, port:int=1883, username=None, password=None, timeout=math.inf
+        self, hostname:str, port:int=1883, username=None, password=None, timeout=60
     ):
         self.client = Client(
             hostname, port=port, username=username, password=password, timeout=timeout
         )
         self.connected = False  # 是否已连接标志
         self.message_queue = asyncio.Queue()  # 用于存储接收到的消息
-        self.subscribers = {}  # 订阅者信息，topic -> Agent 映射
         self.receive_messages_task = None
     
     async def __aexit__(self, exc_type, exc_value, traceback):
         await self.stop()
 
     async def connect(self):
-        try:
-            await self.client.__aenter__()
-            self.connected = True
-            logger.info("Connected to MQTT Broker")
-        except Exception as e:
-            self.connected = False
-            logger.error(f"Failed to connect to MQTT Broker: {e}")
+        for i in range(3):
+            try:
+                await self.client.__aenter__()
+                self.connected = True
+                logger.info("Connected to MQTT Broker")
+                return
+            except Exception as e:
+                logger.error(f"Attempt {i+1}: Failed to connect to MQTT Broker: {e}")
+                await asyncio.sleep(10)
+        self.connected = False
+        logger.error("All connection attempts failed.")
 
     async def disconnect(self):
         await self.client.__aexit__(None, None, None)
@@ -50,9 +55,6 @@ class Messager:
             topics = [topics]
         if not isinstance(agents, list):
             agents = [agents]
-        for topic, agent in zip(topics, agents):
-            self.subscribers[topic] = agent
-            logger.info(f"Subscribed to {topic} for Agent {agent._uuid}")
         await self.client.subscribe(topics, qos=1)
 
     async def receive_messages(self):

{pycityagent-2.0.0a40 → pycityagent-2.0.0a41}/pycityagent/simulation/agentgroup.py

@@ -63,7 +63,7 @@ class AgentGroup:
         if self.enable_pgsql:
             pass
 
-        self.messager = Messager(
+        self.messager = Messager.remote(
             hostname=config["simulator_request"]["mqtt"]["server"],
             port=config["simulator_request"]["mqtt"]["port"],
             username=config["simulator_request"]["mqtt"].get("username", None),
@@ -144,17 +144,17 @@ class AgentGroup:
             await agent.bind_to_simulator()  # type: ignore
         self.id2agent = {agent._uuid: agent for agent in self.agents}
         logger.debug(f"-----Binding Agents to Messager in AgentGroup {self._uuid} ...")
-        await self.messager.connect()
-        if self.messager.is_connected():
-            await self.messager.start_listening()
+        await self.messager.connect.remote()
+        if ray.get(self.messager.is_connected.remote()):
+            await self.messager.start_listening.remote()
             topics = []
             agents = []
             for agent in self.agents:
                 agent.set_messager(self.messager)
                 topic = (f"exps/{self.exp_id}/agents/{agent._uuid}/#", 1)
                 topics.append(topic)
-                agents.append(agent)
-            await self.messager.subscribe(topics, agents)
+                agents.append(agent.uuid)
+            await self.messager.subscribe.remote(topics, agents)
         self.message_dispatch_task = asyncio.create_task(self.message_dispatch())
         if self.enable_avro:
             logger.debug(f"-----Creating Avro files in AgentGroup {self._uuid} ...")
@@ -253,14 +253,14 @@ class AgentGroup:
     async def message_dispatch(self):
         logger.debug(f"-----Starting message dispatch for group {self._uuid}")
         while True:
-            if not self.messager.is_connected():
+            if not ray.get(self.messager.is_connected.remote()):
                 logger.warning(
                     "Messager is not connected. Skipping message processing."
                 )
                 break
 
             # Step 1: 获取消息
-            messages = await self.messager.fetch_messages()
+            messages = await self.messager.fetch_messages.remote()
             logger.info(f"Group {self._uuid} received {len(messages)} messages")
 
             # Step 2: 分发消息到对应的 Agent

{pycityagent-2.0.0a40 → pycityagent-2.0.0a41}/pycityagent/simulation/simulation.py

@@ -69,7 +69,7 @@ class AgentSimulation:
         self._loop = asyncio.get_event_loop()
         # self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
 
-        self._messager = Messager(
+        self._messager = Messager.remote(
             hostname=config["simulator_request"]["mqtt"]["server"],
             port=config["simulator_request"]["mqtt"]["port"],
             username=config["simulator_request"]["mqtt"].get("username", None),
@@ -206,7 +206,7 @@ class AgentSimulation:
             group_size: 每个组的智能体数量，每一个组为一个独立的ray actor
             memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组, 如果为列表，则每个元素表示一个智能体类创建的Memory配置函数
         """
-        await self._messager.connect()
+        await self._messager.connect.remote()
         if not isinstance(agent_count, list):
             agent_count = [agent_count]
 
@@ -499,7 +499,7 @@ class AgentSimulation:
         }
         for uuid in agent_uuids:
             topic = self._user_survey_topics[uuid]
-            await self._messager.send_message(topic, payload)
+            await self._messager.send_message.remote(topic, payload)
 
     async def send_interview_message(
         self, content: str, agent_uuids: Union[uuid.UUID, list[uuid.UUID]]
@@ -516,7 +516,7 @@ class AgentSimulation:
             agent_uuids = [agent_uuids]
         for uuid in agent_uuids:
             topic = self._user_chat_topics[uuid]
-            await self._messager.send_message(topic, payload)
+            await self._messager.send_message.remote(topic, payload)
 
     async def step(self):
         """运行一步, 即每个智能体执行一次forward"""

{pycityagent-2.0.0a40 → pycityagent-2.0.0a41}/pyproject.toml

@@ -3,7 +3,7 @@ in-project = true
 
 [tool.poetry]
 name = "pycityagent"
-version = "2.0.0a40"
+version = "2.0.0a41"
 description = "LLM-based城市环境agent构建库"
 authors = ["Yuwei Yan <pinkgranite86@gmail.com>","Junbo Yan <yanjb20thu@gmali.com>"]
 license = "MIT"