pycityagent 2.0.0a48__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a49__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/agent.py

@@ -9,7 +9,7 @@ from abc import ABC, abstractmethod
 from copy import deepcopy
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any, List, Optional, Type, get_type_hints
+from typing import Any,  Optional, Type, get_type_hints
 from uuid import UUID
 
 import fastavro
@@ -49,7 +49,7 @@ class Agent(ABC):
     """
     Agent base class
     """
-    configurable_fields: List[str] = []
+    configurable_fields: list[str] = []
     default_values: dict[str, Any] = {}
 
     def __init__(
@@ -107,7 +107,7 @@ class Agent(ABC):
         return state
 
     @classmethod
-    def export_class_config(cls) -> Dict[str, Dict]:
+    def export_class_config(cls) -> dict[str, Dict]:
         result = {
             "agent_name": cls.__name__,
             "config": {},
@@ -131,7 +131,7 @@ class Agent(ABC):
         return result
 
     @classmethod
-    def _export_subblocks(cls, block_cls: Type[Block]) -> List[Dict]:
+    def _export_subblocks(cls, block_cls: Type[Block]) -> list[Dict]:
         children = []
         hints = get_type_hints(block_cls)  # 获取类的注解
         for attr_name, attr_type in hints.items():
@@ -151,7 +151,7 @@ class Agent(ABC):
             json.dump(config, f, indent=4)
 
     @classmethod
-    def import_block_config(cls, config: Dict[str, List[Dict]]) -> "Agent":
+    def import_block_config(cls, config: dict[str, list[Dict]]) -> "Agent":
         agent = cls(name=config["agent_name"])
 
         def build_block(block_data: Dict) -> Block:
@@ -172,7 +172,7 @@ class Agent(ABC):
             config = json.load(f)
             return cls.import_block_config(config)
         
-    def load_from_config(self, config: Dict[str, List[Dict]]) -> None:
+    def load_from_config(self, config: dict[str, list[Dict]]) -> None:
         """
         使用配置更新当前Agent实例的Block层次结构。
         """
@@ -185,7 +185,7 @@ class Agent(ABC):
         # 递归更新或创建顶层Block
         for block_data in config.get("blocks", []):
             block_name = block_data["name"]
-            existing_block = getattr(self, block_name, None)
+            existing_block = getattr(self, block_name, None) # type:ignore
 
             if existing_block:
                 # 如果Block已经存在，则递归更新

pycityagent/simulation/agentgroup.py

@@ -1,9 +1,9 @@
 import asyncio
-from collections.abc import Callable
 import json
 import logging
 import time
 import uuid
+from collections.abc import Callable
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, Optional, Type, Union
@@ -34,7 +34,10 @@ class AgentGroup:
         self,
         agent_class: Union[type[Agent], list[type[Agent]]],
         number_of_agents: Union[int, list[int]],
-        memory_config_function_group: Union[Callable[[], tuple[dict, dict, dict]], list[Callable[[], tuple[dict, dict, dict]]]],
+        memory_config_function_group: Union[
+            Callable[[], tuple[dict, dict, dict]],
+            list[Callable[[], tuple[dict, dict, dict]]],
+        ],
         config: dict,
         exp_id: str | UUID,
         exp_name: str,
@@ -81,14 +84,14 @@ class AgentGroup:
         # prepare Messager
         if "mqtt" in config["simulator_request"]:
             self.messager = Messager.remote(
-                hostname=config["simulator_request"]["mqtt"]["server"], 
+                hostname=config["simulator_request"]["mqtt"]["server"],  # type:ignore
                 port=config["simulator_request"]["mqtt"]["port"],
                 username=config["simulator_request"]["mqtt"].get("username", None),
                 password=config["simulator_request"]["mqtt"].get("password", None),
             )
         else:
             self.messager = None
-        
+
         self.message_dispatch_task = None
         self._pgsql_writer = pgsql_writer
         self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
@@ -168,21 +171,21 @@ class AgentGroup:
     @property
     def agent_count(self):
         return self.number_of_agents
-    
+
     @property
     def agent_uuids(self):
         return list(self.id2agent.keys())
-    
+
     @property
     def agent_type(self):
         return self.agent_class
-    
+
     def get_agent_count(self):
         return self.agent_count
-    
+
     def get_agent_uuids(self):
         return self.agent_uuids
-    
+
     def get_agent_type(self):
         return self.agent_type
 
@@ -190,10 +193,6 @@ class AgentGroup:
         self.message_dispatch_task.cancel()  # type: ignore
         await asyncio.gather(self.message_dispatch_task, return_exceptions=True)  # type: ignore
 
-    async def __aexit__(self, exc_type, exc_value, traceback):
-        self.message_dispatch_task.cancel()  # type: ignore
-        await asyncio.gather(self.message_dispatch_task, return_exceptions=True)  # type: ignore
-
     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} ...")
@@ -201,6 +200,7 @@ 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} ...")
+        assert self.messager is not None
         await self.messager.connect.remote()
         if await self.messager.is_connected.remote():
             await self.messager.start_listening.remote()
@@ -293,10 +293,12 @@ class AgentGroup:
         self.initialized = True
         logger.debug(f"-----AgentGroup {self._uuid} initialized")
 
-    async def filter(self, 
-                     types: Optional[list[Type[Agent]]] = None, 
-                     keys: Optional[list[str]] = None, 
-                     values: Optional[list[Any]] = None) -> list[str]:
+    async def filter(
+        self,
+        types: Optional[list[Type[Agent]]] = None,
+        keys: Optional[list[str]] = None,
+        values: Optional[list[Any]] = None,
+    ) -> list[str]:
         filtered_uuids = []
         for agent in self.agents:
             add = True
@@ -304,6 +306,7 @@ class AgentGroup:
                 if agent.__class__ in types:
                     if keys:
                         for key in keys:
+                            assert values is not None
                             if not agent.memory.get(key) == values[keys.index(key)]:
                                 add = False
                                 break
@@ -311,6 +314,7 @@ class AgentGroup:
                         filtered_uuids.append(agent._uuid)
             elif keys:
                 for key in keys:
+                    assert values is not None
                     if not agent.memory.get(key) == values[keys.index(key)]:
                         add = False
                         break
@@ -335,6 +339,7 @@ class AgentGroup:
     async def message_dispatch(self):
         logger.debug(f"-----Starting message dispatch for group {self._uuid}")
         while True:
+            assert self.messager is not None
             if not await self.messager.is_connected.remote():
                 logger.warning(
                     "Messager is not connected. Skipping message processing."

pycityagent/workflow/tool.py

@@ -1,10 +1,16 @@
-from typing import Any, Optional, Union
+import asyncio
+import time
 from collections import defaultdict
 from collections.abc import Callable, Sequence
+from typing import Any, Optional, Union
+
 from mlflow.entities import Metric
-import time
 
-from ..environment import LEVEL_ONE_PRE, POI_TYPE_DICT
+from ..agent import Agent
+from ..utils.decorators import lock_decorator
+from ..environment import (LEVEL_ONE_PRE, POI_TYPE_DICT, AoiService,
+                           PersonService)
+from ..workflow import Block
 
 
 class Tool:
@@ -34,31 +40,23 @@ class Tool:
         raise NotImplementedError
 
     @property
-    def agent(self):
+    def agent(self) -> Agent:
         instance = self._instance  # type:ignore
-        if not isinstance(instance, self._get_agent_class()):
+        if not isinstance(instance, Agent):
             raise RuntimeError(
                 f"Tool bind to object `{type(instance).__name__}`, not an `Agent` object!"
             )
         return instance
 
     @property
-    def block(self):
+    def block(self) -> Block:
         instance = self._instance  # type:ignore
-        if not isinstance(instance, self._get_block_class()):
+        if not isinstance(instance, Block):
             raise RuntimeError(
                 f"Tool bind to object `{type(instance).__name__}`, not an `Block` object!"
             )
         return instance
 
-    def _get_agent_class(self):
-        from ..agent import Agent
-        return Agent
-
-    def _get_block_class(self):
-        from ..workflow import Block
-        return Block
-
 
 class GetMap(Tool):
     """Retrieve the map from the simulator. Can be bound only to an `Agent` instance."""
@@ -140,7 +138,7 @@ class SencePOI(Tool):
 
 class UpdateWithSimulator(Tool):
     def __init__(self) -> None:
-        pass
+        self._lock = asyncio.Lock()
 
     async def _update_motion_with_sim(
         self,
@@ -164,6 +162,7 @@ class UpdateWithSimulator(Tool):
             except KeyError as e:
                 continue
 
+    @lock_decorator
     async def __call__(
         self,
     ):
@@ -173,8 +172,9 @@ class UpdateWithSimulator(Tool):
 
 class ResetAgentPosition(Tool):
     def __init__(self) -> None:
-        pass
+        self._lock = asyncio.Lock()
 
+    @lock_decorator
     async def __call__(
         self,
         aoi_id: Optional[int] = None,
@@ -198,7 +198,8 @@ class ExportMlflowMetrics(Tool):
         self._log_batch_size = log_batch_size
         # TODO: support other log types
         self.metric_log_cache: dict[str, list[Metric]] = defaultdict(list)
-
+        self._lock = asyncio.Lock()
+    @lock_decorator
     async def __call__(
         self,
         metric: Union[Sequence[Union[Metric, dict]], Union[Metric, dict]],
@@ -230,7 +231,7 @@ class ExportMlflowMetrics(Tool):
                 _cache = _cache[batch_size:]
         if clear_cache:
             await self._clear_cache()
-
+    @lock_decorator
     async def _clear_cache(
         self,
     ):

{pycityagent-2.0.0a48.dist-info → pycityagent-2.0.0a49.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a48
+Version: 2.0.0a49
 Summary: LLM-based city environment agent building library
 Author-email: Yuwei Yan <pinkgranite86@gmail.com>, Junbo Yan <yanjb20thu@gmali.com>, Jun Zhang <zhangjun990222@gmali.com>
 License: MIT License

{pycityagent-2.0.0a48.dist-info → pycityagent-2.0.0a49.dist-info}/RECORD

@@ -1,7 +1,7 @@
 pycityagent/pycityagent-sim,sha256=n96jlVZRINlBec5SPOGAdUmeLWMoEKGgoH29iOVJ0wE,34081890
 pycityagent/__init__.py,sha256=9sHyzo8monlBPDe7LTlxCIHeNPYj8HJJDCeuJrYJLsQ,750
 pycityagent/pycityagent-ui,sha256=cHZjqtrQ4Fh4qtRahFNCNbT2DNHLmUexiDAa-72Z3RQ,40333378
-pycityagent/agent.py,sha256=PgrtwgcUwTrBtjQPvXxnE2GnYiG2xtG4Ev0oRzd4eec,34017
+pycityagent/agent.py,sha256=iaF83y344V29-jghx5bJD3yY8kMuSzInBPNi7cVxYiA,34026
 pycityagent/metrics/mlflow_client.py,sha256=g_tHxWkWTDijtbGL74-HmiYzWVKb1y8-w12QrY9jL30,4449
 pycityagent/metrics/__init__.py,sha256=X08PaBbGVAd7_PRGLREXWxaqm7nS82WBQpD1zvQzcqc,128
 pycityagent/metrics/utils/const.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -23,7 +23,7 @@ pycityagent/memory/faiss_query.py,sha256=Z0JS4udyPYCIzHMq464QtHscnswu35gh9fQptik
 pycityagent/memory/state.py,sha256=TYItiyDtehMEQaSBN7PpNrnNxdDM5jGppr9R9Ufv3kA,5134
 pycityagent/simulation/simulation.py,sha256=7Y1r5DAFaB1vuDd3l_85NWpHQUy2GLkFEb800xsJr0E,25565
 pycityagent/simulation/__init__.py,sha256=P5czbcg2d8S0nbbnsQXFIhwzO4CennAhZM8OmKvAeYw,194
-pycityagent/simulation/agentgroup.py,sha256=MwvsWn3UxoY9qbTSCI0SdG9OnQ22Tk1czIfm-0vRbzI,29343
+pycityagent/simulation/agentgroup.py,sha256=DGN6q9S65K9gSv_ihBeT6iJ__LJhl27yXBFyLEk55t0,29302
 pycityagent/simulation/storage/pg.py,sha256=5itxKOkNPlOzN7z2_3oKU1ZK0uLTDugfld8ZkRbD69I,8407
 pycityagent/message/__init__.py,sha256=TCjazxqb5DVwbTu1fF0sNvaH_EPXVuj2XQ0p6W-QCLU,55
 pycityagent/message/messager.py,sha256=78K31EPKfC5IxbISc-Lc2babC7VOh9Vbe3c0sO-YDLA,3333
@@ -40,7 +40,7 @@ pycityagent/cli/wrapper.py,sha256=2Tb52gOlEVgn11Ekt6ZkRXr_BGzte-EPyBKnR6g6hQ4,11
 pycityagent/workflow/__init__.py,sha256=QNkUV-9mACMrR8c0cSKna2gC1mMZdxXbxWzjE-Uods0,621
 pycityagent/workflow/prompt.py,sha256=6jI0Rq54JLv3-IXqZLYug62vse10wTI83xvf4ZX42nk,2929
 pycityagent/workflow/block.py,sha256=8pK17DGX8MnFJHq5lJ4susdZ-CDAEIjlTV_2mJxhbbw,9600
-pycityagent/workflow/tool.py,sha256=jSO8lruDXPLkt4sJJuXvAuaV2XBMXZee_MfQqPo66AQ,8640
+pycityagent/workflow/tool.py,sha256=ChuN5Sm-yltesemewjnTOYaxN6NUjfi2MezVuclf6Tc,8776
 pycityagent/workflow/trigger.py,sha256=Df-MOBEDWBbM-v0dFLQLXteLsipymT4n8vqexmK2GiQ,5643
 pycityagent/environment/__init__.py,sha256=awHxlOud-btWbk0FCS4RmGJ13W84oVCkbGfcrhKqihA,240
 pycityagent/environment/simulator.py,sha256=1OUfODDzM4EN6Lw_Wzq4KeQb-EpcUBioZYc9fxfSPn0,12070
@@ -92,9 +92,9 @@ pycityagent/cityagent/blocks/mobility_block.py,sha256=ZebloDhvZJYbV0YSWhAbnOLToi
 pycityagent/survey/models.py,sha256=YE50UUt5qJ0O_lIUsSY6XFCGUTkJVNu_L1gAhaCJ2fs,3546
 pycityagent/survey/__init__.py,sha256=rxwou8U9KeFSP7rMzXtmtp2fVFZxK4Trzi-psx9LPIs,153
 pycityagent/survey/manager.py,sha256=S5IkwTdelsdtZETChRcfCEczzwSrry_Fly9MY4s3rbk,1681
-pycityagent-2.0.0a48.dist-info/RECORD,,
-pycityagent-2.0.0a48.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
-pycityagent-2.0.0a48.dist-info/WHEEL,sha256=76LiH1wuDHqtCdrF20bc1KNGPhF-AWvwbfgCxtm-4DY,109
-pycityagent-2.0.0a48.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
-pycityagent-2.0.0a48.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
-pycityagent-2.0.0a48.dist-info/METADATA,sha256=FtexuISKTOX_w9zQTOv-pXMHQcrkoJWHeNnMgxz37hY,9139
+pycityagent-2.0.0a49.dist-info/RECORD,,
+pycityagent-2.0.0a49.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
+pycityagent-2.0.0a49.dist-info/WHEEL,sha256=76LiH1wuDHqtCdrF20bc1KNGPhF-AWvwbfgCxtm-4DY,109
+pycityagent-2.0.0a49.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
+pycityagent-2.0.0a49.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
+pycityagent-2.0.0a49.dist-info/METADATA,sha256=AArJCX11voHGW5AlMrin1qZoj8iVGdOD9XLODRcQwUY,9139