pycityagent 2.0.0a34__py3-none-any.whl → 2.0.0a35__py3-none-any.whl

pycityagent/agent.py

@@ -687,6 +687,7 @@ class InstitutionAgent(Agent):
             # TODO: More general id generation
             _id = random.randint(100000, 999999)
             self._agent_id = _id
+            self.memory.set_agent_id(_id)
             await self.memory.update("id", _id, protect_llm_read_only_fields=False)
             try:
                 await self._economy_client.remove_orgs([self._agent_id])

pycityagent/message/messager.py

@@ -1,8 +1,8 @@
 import asyncio
-from collections import defaultdict
 import json
 import logging
 import math
+from typing import Any, List, Union
 from aiomqtt import Client
 
 logger = logging.getLogger("pycityagent")
@@ -17,6 +17,10 @@ class Messager:
         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:
@@ -36,15 +40,20 @@ class Messager:
         """检查是否成功连接到 Broker"""
         return self.connected
 
-    async def subscribe(self, topic, agent):
+    async def subscribe(self, topics: Union[str, List[str]], agents: Union[Any, List[Any]]):
         if not self.is_connected():
             logger.error(
-                f"Cannot subscribe to {topic} because not connected to the Broker."
+                f"Cannot subscribe to {topics} because not connected to the Broker."
             )
             return
-        await self.client.subscribe(topic)
-        self.subscribers[topic] = agent
-        logger.info(f"Subscribed to {topic} for Agent {agent._uuid}")
+        if not isinstance(topics, list):
+            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):
         """监听并将消息存入队列"""
@@ -61,12 +70,17 @@ class Messager:
     async def send_message(self, topic: str, payload: dict):
         """通过 Messager 发送消息"""
         message = json.dumps(payload, default=str)
-        await self.client.publish(topic, message)
+        await self.client.publish(topic=topic, payload=message, qos=1)
         logger.info(f"Message sent to {topic}: {message}")
 
     async def start_listening(self):
         """启动消息监听任务"""
         if self.is_connected():
-            asyncio.create_task(self.receive_messages())
+            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):
+        self.receive_messages_task.cancel()
+        await asyncio.gather(self.receive_messages_task, return_exceptions=True)
+        await self.disconnect()

pycityagent/simulation/agentgroup.py

@@ -69,6 +69,7 @@ class AgentGroup:
             username=config["simulator_request"]["mqtt"].get("username", None),
             password=config["simulator_request"]["mqtt"].get("password", None),
         )
+        self.message_dispatch_task = None
         self._pgsql_writer = pgsql_writer
         self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
         self.initialized = False
@@ -132,6 +133,10 @@ class AgentGroup:
             if self.embedding_model is not None:
                 agent.memory.set_embedding_model(self.embedding_model)
 
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self.message_dispatch_task.cancel()
+        await asyncio.gather(self.message_dispatch_task, return_exceptions=True)
+
     async def init_agents(self):
         logger.debug(f"-----Initializing Agents in AgentGroup {self._uuid} ...")
         logger.debug(f"-----Binding Agents to Simulator in AgentGroup {self._uuid} ...")
@@ -142,16 +147,14 @@ class AgentGroup:
         await self.messager.connect()
         if self.messager.is_connected():
             await self.messager.start_listening()
+            topics = []
+            agents = []
             for agent in self.agents:
                 agent.set_messager(self.messager)
-                topic = f"exps/{self.exp_id}/agents/{agent._uuid}/agent-chat"
-                await self.messager.subscribe(topic, agent)
-                topic = f"exps/{self.exp_id}/agents/{agent._uuid}/user-chat"
-                await self.messager.subscribe(topic, agent)
-                topic = f"exps/{self.exp_id}/agents/{agent._uuid}/user-survey"
-                await self.messager.subscribe(topic, agent)
-                topic = f"exps/{self.exp_id}/agents/{agent._uuid}/gather"
-                await self.messager.subscribe(topic, agent)
+                topic = (f"exps/{self.exp_id}/agents/{agent._uuid}/#", 1)
+                topics.append(topic)
+                agents.append(agent)
+            await self.messager.subscribe(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} ...")
@@ -225,6 +228,7 @@ class AgentGroup:
                 logger.warning(
                     "Messager is not connected. Skipping message processing."
                 )
+                break
 
             # Step 1: 获取消息
             messages = await self.messager.fetch_messages()

pycityagent/simulation/simulation.py

@@ -74,7 +74,6 @@ class AgentSimulation:
             username=config["simulator_request"]["mqtt"].get("username", None),
             password=config["simulator_request"]["mqtt"].get("password", None),
         )
-        asyncio.create_task(self._messager.connect())
 
         # storage
         _storage_config: dict[str, Any] = config.get("storage", {})
@@ -202,6 +201,7 @@ class AgentSimulation:
             group_size: 每个组的智能体数量，每一个组为一个独立的ray actor
             memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组, 如果为列表，则每个元素表示一个智能体类创建的Memory配置函数
         """
+        await self._messager.connect()
         if not isinstance(agent_count, list):
             agent_count = [agent_count]
 
@@ -535,13 +535,6 @@ class AgentSimulation:
         try:
             if self.enable_pgsql:
                 worker: ray.ObjectRef = self._pgsql_writers[0]  # type:ignore
-                # if self._last_asyncio_pg_task is not None:
-                #     await self._last_asyncio_pg_task
-                # self._last_asyncio_pg_task = (
-                #     worker.async_update_exp_info.remote(  # type:ignore
-                #         pg_exp_info
-                #     )
-                # )
                 pg_exp_info = {k: v for k, v in self._exp_info.items()}
                 pg_exp_info["created_at"] = self._exp_created_time
                 pg_exp_info["updated_at"] = self._exp_updated_time

{pycityagent-2.0.0a34.dist-info → pycityagent-2.0.0a35.dist-info}/METADATA

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

{pycityagent-2.0.0a34.dist-info → pycityagent-2.0.0a35.dist-info}/RECORD

@@ -1,5 +1,5 @@
 pycityagent/__init__.py,sha256=fv0mzNGbHBF6m550yYqnuUpB8iQPWS-7EatYRK7DO4s,693
-pycityagent/agent.py,sha256=l8Oa95_K5JBWKzvZmbQe_QM_E_vaG-YstuuR55kgC6Y,29005
+pycityagent/agent.py,sha256=r1uG4ib6fAmqzZXqGUiwswYdIOS8ybTFdbVqJWjhwuI,29047
 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,13 +45,13 @@ 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=W_OVlNGcreHSBf6v-DrEnfNCXExB78ySr0w26MSncfU,2541
+pycityagent/message/messager.py,sha256=gz-EZOGakgwQH8ZKabAGr1pT43E8B9z-s4dbNx-mxr4,3167
 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=QwVbgqKYp42_wRS3K6C6E7Aq8cYBacabHmXIKRaxYyw,23955
-pycityagent/simulation/simulation.py,sha256=9kkdgXSEOAN8wiewVFyORksti4IdVNU0opObV6ZYa9k,23344
+pycityagent/simulation/agentgroup.py,sha256=mZXznSf7VEHNjU0KM5TuPB-Z09DqMZ6VeQib7ciTLpo,23914
+pycityagent/simulation/simulation.py,sha256=OVflF_Z_kZfE6c4z_6tHTGD53jq89Hewg8Yc5KflRI0,23008
 pycityagent/simulation/storage/pg.py,sha256=qGrYzJIAzjv8-d3-cle0rY0AN6XB6MgnHkFLBoLmKWU,7251
 pycityagent/survey/__init__.py,sha256=rxwou8U9KeFSP7rMzXtmtp2fVFZxK4Trzi-psx9LPIs,153
 pycityagent/survey/manager.py,sha256=S5IkwTdelsdtZETChRcfCEczzwSrry_Fly9MY4s3rbk,1681
@@ -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.0a34.dist-info/METADATA,sha256=za2dplTaxTwxYovvgF_HioNdjMXvmpTdg6v4Nx8Q2vI,8033
-pycityagent-2.0.0a34.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-pycityagent-2.0.0a34.dist-info/RECORD,,
+pycityagent-2.0.0a35.dist-info/METADATA,sha256=8uify9sg-cHlR4Vuw_N7ySp46ZoTmV1czHaD51Ng-7k,8033
+pycityagent-2.0.0a35.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+pycityagent-2.0.0a35.dist-info/RECORD,,