pycityagent 2.0.0a39__py3-none-any.whl → 2.0.0a41__py3-none-any.whl

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/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/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/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/simulation/storage/pg.py

@@ -139,7 +139,12 @@ class PgWriter:
                 exec_str = "SELECT * FROM {table_name} WHERE id=%s".format(
                     table_name=table_name
                 ), (self.exp_id,)
-                await cur.execute(exec_str)  # type:ignore
+                await cur.execute(
+                    "SELECT * FROM {table_name} WHERE id=%s".format(
+                        table_name=table_name
+                    ),
+                    (self.exp_id,),
+                )  # type:ignore
                 logger.debug(f"table:{table_name} sql: {exec_str}")
                 record_exists = await cur.fetchall()
                 if record_exists:

{pycityagent-2.0.0a39.dist-info → pycityagent-2.0.0a41.dist-info}/METADATA

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

{pycityagent-2.0.0a39.dist-info → pycityagent-2.0.0a41.dist-info}/RECORD

@@ -1,5 +1,5 @@
 pycityagent/__init__.py,sha256=fv0mzNGbHBF6m550yYqnuUpB8iQPWS-7EatYRK7DO4s,693
-pycityagent/agent.py,sha256=_klNenWHt_AB6go57pjLqw4TcGW6AjoyZTUkmADT6q4,29910
+pycityagent/agent.py,sha256=KO8yJiVOZTWlYAd57gLEgbSNzQ26cw27k1NT8pAFHNY,29917
 pycityagent/economy/__init__.py,sha256=aonY4WHnx-6EGJ4WKrx4S-2jAkYNLtqUA04jp6q8B7w,75
 pycityagent/economy/econ_client.py,sha256=GuHK9ZBnhqW3Z7F8ViDJn_iN73yOBbbwFyJv1wLEBDk,12211
 pycityagent/environment/__init__.py,sha256=awHxlOud-btWbk0FCS4RmGJ13W84oVCkbGfcrhKqihA,240
@@ -45,14 +45,14 @@ pycityagent/memory/self_define.py,sha256=vpZ6CIxR2grNXEIOScdpsSc59FBg0mOKelwQuTE
 pycityagent/memory/state.py,sha256=TYItiyDtehMEQaSBN7PpNrnNxdDM5jGppr9R9Ufv3kA,5134
 pycityagent/memory/utils.py,sha256=oJWLdPeJy_jcdKcDTo9JAH9kDZhqjoQhhv_zT9qWC0w,877
 pycityagent/message/__init__.py,sha256=TCjazxqb5DVwbTu1fF0sNvaH_EPXVuj2XQ0p6W-QCLU,55
-pycityagent/message/messager.py,sha256=gz-EZOGakgwQH8ZKabAGr1pT43E8B9z-s4dbNx-mxr4,3167
+pycityagent/message/messager.py,sha256=ePu1LDZZBDMhxoVoX4-LGTcizCVsOuw4T1GRBXHdM7E,3125
 pycityagent/metrics/__init__.py,sha256=X08PaBbGVAd7_PRGLREXWxaqm7nS82WBQpD1zvQzcqc,128
 pycityagent/metrics/mlflow_client.py,sha256=g_tHxWkWTDijtbGL74-HmiYzWVKb1y8-w12QrY9jL30,4449
 pycityagent/metrics/utils/const.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycityagent/simulation/__init__.py,sha256=P5czbcg2d8S0nbbnsQXFIhwzO4CennAhZM8OmKvAeYw,194
-pycityagent/simulation/agentgroup.py,sha256=ZPzgHgin-XOz3owxMbyAxcz8AiU4I6XM-ppOUsFsAIE,25284
-pycityagent/simulation/simulation.py,sha256=aE-U9wEay9ksheSOotxyLfZOGxlKxP0OFOuRt3mDmSY,23246
-pycityagent/simulation/storage/pg.py,sha256=UuBO0J7e1Oi2LsWygIiSHQ9W-RQO4X7Vpy9Ci0aMP8s,8272
+pycityagent/simulation/agentgroup.py,sha256=3UIMbA-CpoHxtxkFWkzEpI6LMclD3CwvxS2WjpqjaWE,25356
+pycityagent/simulation/simulation.py,sha256=LGWcooTa5n3YVUeChv3M1T6keUJK1a4Ddo-g32WK1o0,23274
+pycityagent/simulation/storage/pg.py,sha256=jjdYvqKmDms3weqALqOO__gxWF-Z0YcqeD85XbP4Qks,8455
 pycityagent/survey/__init__.py,sha256=rxwou8U9KeFSP7rMzXtmtp2fVFZxK4Trzi-psx9LPIs,153
 pycityagent/survey/manager.py,sha256=S5IkwTdelsdtZETChRcfCEczzwSrry_Fly9MY4s3rbk,1681
 pycityagent/survey/models.py,sha256=YE50UUt5qJ0O_lIUsSY6XFCGUTkJVNu_L1gAhaCJ2fs,3546
@@ -70,6 +70,6 @@ pycityagent/workflow/block.py,sha256=C2aWdVRffb3LknP955GvPcBMsm3VPXN9ZuAtCgITFTo
 pycityagent/workflow/prompt.py,sha256=6jI0Rq54JLv3-IXqZLYug62vse10wTI83xvf4ZX42nk,2929
 pycityagent/workflow/tool.py,sha256=xADxhNgVsjNiMxlhdwn3xGUstFOkLEG8P67ez8VmwSI,8555
 pycityagent/workflow/trigger.py,sha256=Df-MOBEDWBbM-v0dFLQLXteLsipymT4n8vqexmK2GiQ,5643
-pycityagent-2.0.0a39.dist-info/METADATA,sha256=4rGvl0VT73oe6gw5m7_xvj_Wf8pRFGbSSaXRbAvWW3E,8046
-pycityagent-2.0.0a39.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-pycityagent-2.0.0a39.dist-info/RECORD,,
+pycityagent-2.0.0a41.dist-info/METADATA,sha256=jR-wJCsCj1qthzVzdEcXZOrFpG_v7TXyZ8Sgh-2yTHU,8046
+pycityagent-2.0.0a41.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+pycityagent-2.0.0a41.dist-info/RECORD,,