pycityagent 2.0.0a48__cp310-cp310-macosx_11_0_arm64.whl → 2.0.0a50__cp310-cp310-macosx_11_0_arm64.whl

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,
@@ -45,7 +48,7 @@ class AgentGroup:
         mlflow_run_id: str,
         embedding_model: Embeddings,
         logging_level: int,
-        agent_config_file: Union[str, list[str]] = None,
+        agent_config_file: Optional[Union[str, list[str]]] = None,
     ):
         logger.setLevel(logging_level)
         self._uuid = str(uuid.uuid4())
@@ -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隐藏到计算任务后
@@ -160,7 +163,7 @@ class AgentGroup:
                     agent.memory.set_faiss_query(self.faiss_query)
                 if self.embedding_model is not None:
                     agent.memory.set_embedding_model(self.embedding_model)
-                if self.agent_config_file[i]:
+                if self.agent_config_file is not None and self.agent_config_file[i]:
                     agent.load_from_file(self.agent_config_file[i])
                 self.agents.append(agent)
                 self.id2agent[agent._uuid] = agent
@@ -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/simulation/simulation.py

@@ -13,6 +13,11 @@ import yaml
 from langchain_core.embeddings import Embeddings
 
 from ..agent import Agent, InstitutionAgent
+from ..cityagent import (BankAgent, FirmAgent, GovernmentAgent, NBSAgent,
+                         SocietyAgent, memory_config_bank, memory_config_firm,
+                         memory_config_government, memory_config_nbs,
+                         memory_config_societyagent)
+from ..cityagent.initial import bind_agent_info, initialize_social_network
 from ..environment.simulator import Simulator
 from ..llm import SimpleEmbedding
 from ..memory import Memory
@@ -22,11 +27,10 @@ from ..survey import Survey
 from ..utils import TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
 from .agentgroup import AgentGroup
 from .storage.pg import PgWriter, create_pg_tables
-from ..cityagent import SocietyAgent, FirmAgent, BankAgent, NBSAgent, GovernmentAgent, memory_config_societyagent, memory_config_government, memory_config_firm, memory_config_bank, memory_config_nbs
-from ..cityagent.initial import bind_agent_info, initialize_social_network
 
 logger = logging.getLogger("pycityagent")
 
+
 class AgentSimulation:
     """城市智能体模拟器"""
 
@@ -52,7 +56,13 @@ class AgentSimulation:
             self.agent_class = agent_class
         elif agent_class is None:
             if enable_economy:
-                self.agent_class = [SocietyAgent, FirmAgent, BankAgent, NBSAgent, GovernmentAgent]
+                self.agent_class = [
+                    SocietyAgent,
+                    FirmAgent,
+                    BankAgent,
+                    NBSAgent,
+                    GovernmentAgent,
+                ]
                 self.default_memory_config_func = [
                     memory_config_societyagent,
                     memory_config_firm,
@@ -82,7 +92,7 @@ class AgentSimulation:
         # self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
 
         self._messager = Messager.remote(
-            hostname=config["simulator_request"]["mqtt"]["server"], # type:ignore
+            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),
@@ -143,7 +153,7 @@ class AgentSimulation:
         """Directly run from config file
         Basic config file should contain:
         - simulation_config: file_path
-        - agent_config: 
+        - agent_config:
             - agent_config_file: Optional[dict]
             - memory_config_func: Optional[Union[Callable, list[Callable]]]
             - init_func: Optional[list[Callable[AgentSimulation, None]]]
@@ -173,13 +183,14 @@ class AgentSimulation:
         if "workflow" not in config:
             raise ValueError("workflow is required")
         import yaml
+
         logger.info("Loading config file...")
         with open(config["simulation_config"], "r") as f:
             simulation_config = yaml.safe_load(f)
         logger.info("Creating AgentSimulation Task...")
         simulation = cls(
-            config=simulation_config, 
-            agent_config_file=config["agent_config"].get("agent_config_file", None), 
+            config=simulation_config,
+            agent_config_file=config["agent_config"].get("agent_config_file", None),
             exp_name=config.get("exp_name", "default_experiment"),
             logging_level=config.get("logging_level", logging.WARNING),
         )
@@ -193,21 +204,28 @@ class AgentSimulation:
         await simulation.init_agents(
             agent_count=agent_count,
             group_size=config["agent_config"].get("group_size", 10000),
-            embedding_model=config["agent_config"].get("embedding_model", SimpleEmbedding()),
+            embedding_model=config["agent_config"].get(
+                "embedding_model", SimpleEmbedding()
+            ),
             memory_config_func=config["agent_config"].get("memory_config_func", None),
         )
         logger.info("Running Init Functions...")
-        for init_func in config["agent_config"].get("init_func", [bind_agent_info, initialize_social_network]):
+        for init_func in config["agent_config"].get(
+            "init_func", [bind_agent_info, initialize_social_network]
+        ):
             await init_func(simulation)
         logger.info("Starting Simulation...")
         for step in config["workflow"]:
-            logger.info(f"Running step: type: {step['type']} - description: {step.get('description', 'no description')}")
+            logger.info(
+                f"Running step: type: {step['type']} - description: {step.get('description', 'no description')}"
+            )
             if step["type"] not in ["run", "step", "interview", "survey", "intervene"]:
                 raise ValueError(f"Invalid step type: {step['type']}")
             if step["type"] == "run":
                 await simulation.run(step.get("day", 1))
             elif step["type"] == "step":
-                await simulation.step(step.get("time", 1))
+                # await simulation.step(step.get("time", 1))
+                await simulation.step()
             else:
                 await step["step_func"](simulation)
         logger.info("Simulation finished")
@@ -241,11 +259,11 @@ class AgentSimulation:
     @property
     def agent_uuid2group(self):
         return self._agent_uuid2group
-    
+
     @property
     def messager(self):
         return self._messager
-    
+
     async def _save_exp_info(self) -> None:
         """异步保存实验信息到YAML文件"""
         try:
@@ -354,38 +372,44 @@ class AgentSimulation:
         # 分别处理机构智能体和普通智能体
         institution_params = []
         citizen_params = []
-        
+
         # 收集所有参数
         for i in range(len(self.agent_class)):
             agent_class = self.agent_class[i]
             agent_count_i = agent_count[i]
             memory_config_func_i = memory_config_func[i]
-            
+
             if self.agent_config_file is not None:
-                config_file = self.agent_config_file.get(agent_class, None) 
+                config_file = self.agent_config_file.get(agent_class, None)
             else:
                 config_file = None
-                
+
             if issubclass(agent_class, InstitutionAgent):
-                institution_params.append((agent_class, agent_count_i, memory_config_func_i, config_file))
+                institution_params.append(
+                    (agent_class, agent_count_i, memory_config_func_i, config_file)
+                )
             else:
-                citizen_params.append((agent_class, agent_count_i, memory_config_func_i, config_file))
+                citizen_params.append(
+                    (agent_class, agent_count_i, memory_config_func_i, config_file)
+                )
 
         # 处理机构智能体组
         if institution_params:
             total_institution_count = sum(p[1] for p in institution_params)
-            num_institution_groups = (total_institution_count + group_size - 1) // group_size
-            
+            num_institution_groups = (
+                total_institution_count + group_size - 1
+            ) // group_size
+
             for k in range(num_institution_groups):
                 start_idx = k * group_size
                 remaining = total_institution_count - start_idx
                 number_of_agents = min(remaining, group_size)
-                
+
                 agent_classes = []
                 agent_counts = []
                 memory_config_funcs = []
                 config_files = []
-                
+
                 # 分配每种类型的机构智能体到当前组
                 curr_start = start_idx
                 for agent_class, count, mem_func, conf_file in institution_params:
@@ -395,30 +419,32 @@ class AgentSimulation:
                         memory_config_funcs.append(mem_func)
                         config_files.append(conf_file)
                     curr_start = max(0, curr_start - count)
-                
-                group_creation_params.append((
-                    agent_classes,
-                    agent_counts, 
-                    memory_config_funcs,
-                    f"InstitutionGroup_{k}",
-                    config_files
-                ))
+
+                group_creation_params.append(
+                    (
+                        agent_classes,
+                        agent_counts,
+                        memory_config_funcs,
+                        f"InstitutionGroup_{k}",
+                        config_files,
+                    )
+                )
 
         # 处理普通智能体组
         if citizen_params:
             total_citizen_count = sum(p[1] for p in citizen_params)
             num_citizen_groups = (total_citizen_count + group_size - 1) // group_size
-            
+
             for k in range(num_citizen_groups):
                 start_idx = k * group_size
                 remaining = total_citizen_count - start_idx
                 number_of_agents = min(remaining, group_size)
-                
+
                 agent_classes = []
                 agent_counts = []
                 memory_config_funcs = []
                 config_files = []
-                
+
                 # 分配每种类型的普通智能体到当前组
                 curr_start = start_idx
                 for agent_class, count, mem_func, conf_file in citizen_params:
@@ -428,14 +454,16 @@ class AgentSimulation:
                         memory_config_funcs.append(mem_func)
                         config_files.append(conf_file)
                     curr_start = max(0, curr_start - count)
-                
-                group_creation_params.append((
-                    agent_classes,
-                    agent_counts,
-                    memory_config_funcs, 
-                    f"CitizenGroup_{k}",
-                    config_files
-                ))
+
+                group_creation_params.append(
+                    (
+                        agent_classes,
+                        agent_counts,
+                        memory_config_funcs,
+                        f"CitizenGroup_{k}",
+                        config_files,
+                    )
+                )
 
         # 初始化mlflow连接
         _mlflow_config = self.config.get("metric_request", {}).get("mlflow")
@@ -463,7 +491,13 @@ class AgentSimulation:
             self._pgsql_writers = _workers = [None for _ in range(_num_workers)]
 
         creation_tasks = []
-        for i, (agent_class, number_of_agents, memory_config_function_group, group_name, config_file) in enumerate(group_creation_params):
+        for i, (
+            agent_class,
+            number_of_agents,
+            memory_config_function_group,
+            group_name,
+            config_file,
+        ) in enumerate(group_creation_params):
             # 直接创建异步任务
             group = AgentGroup.remote(
                 agent_class,
@@ -489,7 +523,9 @@ class AgentSimulation:
             group_agent_uuids = ray.get(group.get_agent_uuids.remote())
             for agent_uuid in group_agent_uuids:
                 self._agent_uuid2group[agent_uuid] = group
-                self._user_chat_topics[agent_uuid] = f"exps/{self.exp_id}/agents/{agent_uuid}/user-chat"
+                self._user_chat_topics[agent_uuid] = (
+                    f"exps/{self.exp_id}/agents/{agent_uuid}/user-chat"
+                )
                 self._user_survey_topics[agent_uuid] = (
                     f"exps/{self.exp_id}/agents/{agent_uuid}/user-survey"
                 )
@@ -511,23 +547,26 @@ class AgentSimulation:
         for group in self._groups.values():
             gather_tasks.append(group.gather.remote(content))
         return await asyncio.gather(*gather_tasks)
-    
-    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]:
         """过滤出指定类型的智能体"""
         if not types and not keys and not values:
             return self._agent_uuids
         group_to_filter = []
-        for t in types:
-            if t in self._type2group:
-                group_to_filter.extend(self._type2group[t])
-            else:
-                raise ValueError(f"type {t} not found in simulation")
+        if types is not None:
+            for t in types:
+                if t in self._type2group:
+                    group_to_filter.extend(self._type2group[t])
+                else:
+                    raise ValueError(f"type {t} not found in simulation")
         filtered_uuids = []
         if keys:
-            if len(keys) != len(values):
+            if values is None or len(keys) != len(values):
                 raise ValueError("the length of key and value does not match")
             for group in group_to_filter:
                 filtered_uuids.extend(await group.filter.remote(types, keys, values))

pycityagent/tools/__init__.py

@@ -0,0 +1,9 @@
+from .tool import ExportMlflowMetrics, GetMap, SencePOI, Tool, UpdateWithSimulator
+
+__all__ = [
+    "SencePOI",
+    "Tool",
+    "ExportMlflowMetrics",
+    "GetMap",
+    "UpdateWithSimulator",
+]

pycityagent/{workflow → tools}/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 ..environment import (LEVEL_ONE_PRE, POI_TYPE_DICT, AoiService,
+                           PersonService)
+from ..utils.decorators import lock_decorator
+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,9 @@ 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]],
@@ -231,6 +233,7 @@ class ExportMlflowMetrics(Tool):
         if clear_cache:
             await self._clear_cache()
 
+    @lock_decorator
     async def _clear_cache(
         self,
     ):

pycityagent/workflow/__init__.py

@@ -7,14 +7,9 @@ This module contains classes for creating blocks and running workflows.
 from .block import (Block, log_and_check, log_and_check_with_memory,
                     trigger_class)
 from .prompt import FormatPrompt
-from .tool import ExportMlflowMetrics, GetMap, SencePOI, Tool
 from .trigger import EventTrigger, MemoryChangeTrigger, TimeTrigger
 
 __all__ = [
-    "SencePOI",
-    "Tool",
-    "ExportMlflowMetrics",
-    "GetMap",
     "MemoryChangeTrigger",
     "TimeTrigger",
     "EventTrigger",

pycityagent/workflow/block.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
+
 import asyncio
 import functools
 import inspect
-from collections.abc import Awaitable, Callable, Coroutine
 import json
-from typing import Any, List, Optional, Union
+from collections.abc import Awaitable, Callable, Coroutine
+from typing import Any, Optional, Union
 
 from pyparsing import Dict
 
@@ -143,7 +144,7 @@ def trigger_class():
 
 # Define a Block, similar to a layer in PyTorch
 class Block:
-    configurable_fields: List[str] = []
+    configurable_fields: list[str] = []
     default_values: dict[str, Any] = {}
 
     def __init__(
@@ -164,22 +165,23 @@ class Block:
             trigger.initialize()  # 立即初始化trigger
         self.trigger = trigger
 
-    def export_config(self) -> Dict[str, Optional[str]]:
+    def export_config(self) -> dict[str, Optional[str]]:
         return {
             field: self.default_values.get(field, "default_value")
             for field in self.configurable_fields
         }
 
     @classmethod
-    def export_class_config(cls) -> Dict[str, str]:
+    def export_class_config(cls) -> dict[str, str]:
         return {
             field: cls.default_values.get(field, "default_value")
             for field in cls.configurable_fields
         }
 
     @classmethod
-    def import_config(cls, config: Dict[str, str]) -> "Block":
+    def import_config(cls, config: dict[str, Union[str, dict]]) -> Block:
         instance = cls(name=config["name"])
+        assert isinstance(config["config"], dict)
         for field, value in config["config"].items():
             if field in cls.configurable_fields:
                 setattr(instance, field, value)
@@ -190,8 +192,8 @@ class Block:
             setattr(instance, child_block.name.lower(), child_block)
 
         return instance
-    
-    def load_from_config(self, config: Dict[str, List[Dict]]) -> None:
+
+    def load_from_config(self, config: dict[str, list[Dict]]) -> None:
         """
         使用配置更新当前Block实例的参数，并递归更新子Block。
         """
@@ -201,8 +203,8 @@ class Block:
                 if config["config"][field] != "default_value":
                     setattr(self, field, config["config"][field])
 
-        def build_or_update_block(block_data: Dict) -> Block:
-            block_name = block_data["name"].lower()
+        def build_or_update_block(block_data: dict) -> Block:
+            block_name = block_data["name"].lower()  # type:ignore
             existing_block = getattr(self, block_name, None)
 
             if existing_block:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a48
+Version: 2.0.0a50
 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
@@ -50,7 +50,6 @@ Requires-Dist: requests>=2.32.3
 Requires-Dist: Shapely>=2.0.6
 Requires-Dist: PyYAML>=6.0.2
 Requires-Dist: zhipuai>=2.1.5.20230904
-Requires-Dist: gradio>=5.7.1
 Requires-Dist: mosstool>=1.3.0
 Requires-Dist: ray>=2.40.0
 Requires-Dist: aiomqtt>=2.3.0