pycityagent 2.0.0a34__py3-none-any.whl → 2.0.0a36__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/environment/simulator.py

@@ -14,7 +14,6 @@ from pycityproto.city.person.v2 import person_pb2 as person_pb2
 from pycityproto.city.person.v2 import person_service_pb2 as person_service
 from pymongo import MongoClient
 from shapely.geometry import Point
-from shapely.strtree import STRtree
 
 from .sim import CityClient, ControlSimEnv
 from .utils.const import *
@@ -87,12 +86,6 @@ class Simulator:
         - Simulator map object
         """
 
-        self.pois_matrix: dict[str, list[list[list]]] = {}
-        """
-        pois的基于区块的划分——方便快速粗略地查询poi
-        通过Simulator.set_pois_matrix()初始化
-        """
-
         self.time: int = 0
         """
         - 模拟城市当前时间
@@ -102,12 +95,11 @@ class Simulator:
         self.map_x_gap = None
         self.map_y_gap = None
         self._bbox: tuple[float, float, float, float] = (-1, -1, -1, -1)
-        self.poi_matrix_centers = []
         self._lock = asyncio.Lock()
         # poi id dict
-        self.poi_id_2_aoi_id: dict[int, int] = {}
-        # poi STRtree
-        self.set_poi_tree()
+        self.poi_id_2_aoi_id: dict[int, int] = {
+            poi["id"]: poi["aoi_id"] for _, poi in self.map.pois.items()
+        }
 
     # * Agent相关
     def find_agents_by_area(self, req: dict, status=None):
@@ -137,35 +129,21 @@ class Simulator:
             resp.motions = motions  # type: ignore
             return resp
 
-    def set_poi_tree(
-        self,
-    ):
-        """
-        初始化pois_tree
-        """
-        poi_geos = []
-        tree_id_2_poi_and_catg: dict[int, tuple[dict, str]] = {}
-        for tree_id, poi in enumerate(self.map.pois.values()):
-            tree_id_2_poi_and_catg[tree_id] = (poi, poi["category"])
-            poi_geos.append(Point([poi["position"][k] for k in ["x", "y"]]))
-            self.poi_id_2_aoi_id[poi["id"]] = poi["aoi_id"]
-        self.tree_id_2_poi_and_catg = tree_id_2_poi_and_catg
-        self.pois_tree = STRtree(poi_geos)
-
     def get_poi_categories(
         self,
         center: Optional[Union[tuple[float, float], Point]] = None,
         radius: Optional[float] = None,
     ) -> list[str]:
-        if center is not None and radius is not None:
-            if not isinstance(center, Point):
-                center = Point(center)
-            indices = self.pois_tree.query(center.buffer(radius))
-        else:
-            indices = list(self.tree_id_2_poi_and_catg.keys())
-        categories = []
-        for index in indices:
-            _, catg = self.tree_id_2_poi_and_catg[index]
+        categories: list[str] = []
+        if center is None:
+            center = (0, 0)
+        _pois: list[dict] = self.map.query_pois(  # type:ignore
+            center=center,
+            radius=radius,
+            return_distance=False,
+        )
+        for poi in _pois:
+            catg = poi["category"]
             categories.append(catg.split("|")[-1])
         return list(set(categories))
 
@@ -327,7 +305,7 @@ class Simulator:
         center: Union[tuple[float, float], Point],
         radius: float,
         poi_type: Union[str, list[str]],
-    ):
+    ) -> list[dict]:
         if isinstance(poi_type, str):
             poi_type = [poi_type]
         transformed_poi_type = []
@@ -337,14 +315,16 @@ class Simulator:
             else:
                 transformed_poi_type += self.poi_cate[t]
         poi_type_set = set(transformed_poi_type)
-        if not isinstance(center, Point):
-            center = Point(center)
         # 获取半径内的poi
-        indices = self.pois_tree.query(center.buffer(radius))
+        _pois: list[dict] = self.map.query_pois(  # type:ignore
+            center=center,
+            radius=radius,
+            return_distance=False,
+        )
         # 过滤掉不满足类别前缀的poi
         pois = []
-        for index in indices:
-            poi, catg = self.tree_id_2_poi_and_catg[index]
+        for poi in _pois:
+            catg = poi["category"]
             if catg.split("|")[-1] not in poi_type_set:
                 continue
             pois.append(poi)

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.0a36.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a34
+Version: 2.0.0a36
 Summary: LLM-based城市环境agent构建库
 License: MIT
 Author: Yuwei Yan
@@ -29,7 +29,7 @@ Requires-Dist: langchain-community (>=0.3.13,<0.4.0)
 Requires-Dist: langchain-core (>=0.3.28,<0.4.0)
 Requires-Dist: matplotlib (==3.8.3)
 Requires-Dist: mlflow (>=2.19.0,<3.0.0)
-Requires-Dist: mosstool (==1.0.24)
+Requires-Dist: mosstool (>=1.3.0,<2.0.0)
 Requires-Dist: networkx (==3.2.1)
 Requires-Dist: numpy (>=1.20.0,<2.0.0)
 Requires-Dist: openai (>=1.58.1,<2.0.0)
@@ -37,7 +37,7 @@ Requires-Dist: pandavro (>=1.8.0,<2.0.0)
 Requires-Dist: poetry (>=1.2.2)
 Requires-Dist: protobuf (<=4.24.0)
 Requires-Dist: psycopg[binary] (>=3.2.3,<4.0.0)
-Requires-Dist: pycitydata (==1.0.0)
+Requires-Dist: pycitydata (>=1.0.3,<2.0.0)
 Requires-Dist: pycityproto (>=2.1.5,<3.0.0)
 Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
 Requires-Dist: ray (>=2.40.0,<3.0.0)

{pycityagent-2.0.0a34.dist-info → pycityagent-2.0.0a36.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
@@ -21,7 +21,7 @@ pycityagent/environment/sim/person_service.py,sha256=5r1F2Itn7dKJ2U4hSLovrk5p4qy
 pycityagent/environment/sim/road_service.py,sha256=bKyn3_me0sGmaJVyF6eNeFbdU-9C1yWsa9L7pieDJzg,1285
 pycityagent/environment/sim/sim_env.py,sha256=HI1LcS_FotDKQ6vBnx0e49prXSABOfA20aU9KM-ZkCY,4625
 pycityagent/environment/sim/social_service.py,sha256=9EFJAwVdUuUQkNkFRn9qZRDfD1brh2fqkvasnXUEBhQ,2014
-pycityagent/environment/simulator.py,sha256=KVfwSwVGXPqUHQGyD9jv_RXRgGal2k7NloUFVdmWE8I,12943
+pycityagent/environment/simulator.py,sha256=1OUfODDzM4EN6Lw_Wzq4KeQb-EpcUBioZYc9fxfSPn0,12070
 pycityagent/environment/utils/__init__.py,sha256=1m4Q1EfGvNpUsa1bgQzzCyWhfkpElnskNImjjFD3Znc,237
 pycityagent/environment/utils/base64.py,sha256=hoREzQo3FXMN79pqQLO2jgsDEvudciomyKii7MWljAM,374
 pycityagent/environment/utils/const.py,sha256=1LqxnYJ8FSmq37fN5kIFlWLwycEDzFa8SFS-8plrFlU,5396
@@ -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.0a36.dist-info/METADATA,sha256=aJH2tcaEccrMdJRwjgNmRS8bYla-E_TTMh1bgevmHEM,8046
+pycityagent-2.0.0a36.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+pycityagent-2.0.0a36.dist-info/RECORD,,