pycityagent 2.0.0a34__tar.gz → 2.0.0a36__tar.gz

{pycityagent-2.0.0a34 → pycityagent-2.0.0a36}/PKG-INFO

@@ -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 → pycityagent-2.0.0a36}/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-2.0.0a34 → pycityagent-2.0.0a36}/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-2.0.0a34 → pycityagent-2.0.0a36}/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-2.0.0a34 → pycityagent-2.0.0a36}/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-2.0.0a34 → pycityagent-2.0.0a36}/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 → pycityagent-2.0.0a36}/pyproject.toml

@@ -3,7 +3,7 @@ in-project = true
 
 [tool.poetry]
 name = "pycityagent"
-version = "2.0.0a34"
+version = "2.0.0a36"
 description = "LLM-based城市环境agent构建库"
 authors = ["Yuwei Yan <pinkgranite86@gmail.com>","Junbo Yan <yanjb20thu@gmali.com>"]
 license = "MIT"
@@ -31,7 +31,7 @@ numpy = "^1.20.0"
 openai = "^1.58.1"
 Pillow = "11.0.0"
 protobuf = "<=4.24.0,<5.0.0"
-pycitydata = "1.0.0"
+pycitydata = "^1.0.3"
 pycityproto = "^2.1.5"
 Requests = "2.32.3"
 Shapely = "2.0.6"
@@ -39,7 +39,7 @@ sidecar = "0.7.0"
 pyyaml = "^6.0.2"
 zhipuai = "^2.1.5.20230904"
 gradio = "^5.7.1"
-mosstool = "1.0.24"
+mosstool = "^1.3.0"
 ray = "^2.40.0"
 aiomqtt = "^2.3.0"
 fastavro = "^1.10.0"