pycityagent 2.0.0a57__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a59__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/cityagent/societyagent.py

@@ -161,6 +161,7 @@ class MindBlock(Block):
         await self.cognitionBlock.forward()
 
 class SocietyAgent(CitizenAgent):
+    update_with_sim = UpdateWithSimulator()
     mindBlock: MindBlock
     planAndActionBlock: PlanAndActionBlock
     update_with_sim = UpdateWithSimulator()

pycityagent/environment/simulator.py

@@ -65,6 +65,7 @@ class Simulator:
                     sim_addr=config["simulator"].get("server", None),
                 )
                 self.server_addr = sim_env.sim_addr
+                config["simulator"]["server"] = self.server_addr
                 # using local client
                 self._client = CityClient(sim_env.sim_addr, secure=False)
                 """
@@ -72,10 +73,8 @@ class Simulator:
                 - grpc client of simulator
                 """
             else:
-                raise ValueError(
-                    f"Passing Traffic Simulation address is not supported!"
-                )
                 self._client = CityClient(config["simulator"]["server"], secure=False)
+                self.server_addr = config["simulator"]["server"]
         else:
             self.server_addr = None
             logger.warning(

pycityagent/simulation/agentgroup.py

@@ -34,10 +34,7 @@ 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: dict[type[Agent], Callable],
         config: dict,
         exp_name: str,
         exp_id: str | UUID,
@@ -49,14 +46,12 @@ class AgentGroup:
         mlflow_run_id: str,
         embedding_model: Embeddings,
         logging_level: int,
-        agent_config_file: Optional[Union[str, list[str]]] = None,
+        agent_config_file: Optional[dict[type[Agent], str]] = None,
     ):
         logger.setLevel(logging_level)
         self._uuid = str(uuid.uuid4())
         if not isinstance(agent_class, list):
             agent_class = [agent_class]
-        if not isinstance(memory_config_function_group, list):
-            memory_config_function_group = [memory_config_function_group]
         if not isinstance(number_of_agents, list):
             number_of_agents = [number_of_agents]
         self.agent_class = agent_class
@@ -123,13 +118,10 @@ class AgentGroup:
         self.projector = pyproj.Proj(self.simulator.map.header["projection"])
 
         # prepare Economy client
-        if "economy" in config["simulator_request"]:
-            logger.info(f"-----Creating Economy client in AgentGroup {self._uuid} ...")
-            self.economy_client = EconomyClient(
-                config["simulator_request"]["economy"]["server"]
-            )
-        else:
-            self.economy_client = None
+        logger.info(f"-----Creating Economy client in AgentGroup {self._uuid} ...")
+        self.economy_client = EconomyClient(
+            config["simulator_request"]["simulator"]["server"]
+        )
 
         # set FaissQuery
         if self.embedding_model is not None:
@@ -142,7 +134,7 @@ class AgentGroup:
             agent_class_i = agent_class[i]
             number_of_agents_i = number_of_agents[i]
             for j in range(number_of_agents_i):
-                memory_config_function_group_i = memory_config_function_group[i]
+                memory_config_function_group_i = memory_config_function_group[agent_class_i]
                 extra_attributes, profile, base = memory_config_function_group_i()
                 memory = Memory(config=extra_attributes, profile=profile, base=base)
                 agent = agent_class_i(
@@ -161,8 +153,8 @@ class AgentGroup:
                     agent.set_avro_file(self.avro_file)  # type: ignore
                 if self.enable_pgsql:
                     agent.set_pgsql_writer(self._pgsql_writer)
-                if self.agent_config_file is not None and self.agent_config_file[i]:
-                    agent.load_from_file(self.agent_config_file[i])
+                if self.agent_config_file is not None and self.agent_config_file[agent_class_i]:
+                    agent.load_from_file(self.agent_config_file[agent_class_i])
                 if self._message_interceptor is not None:
                     agent.set_message_interceptor(self._message_interceptor)
                 self.agents.append(agent)

pycityagent/simulation/simulation.py

@@ -3,24 +3,24 @@ import json
 import logging
 import time
 import uuid
-from collections.abc import Callable, Sequence
+from collections.abc import Callable
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Any, Optional, Type, Union
+from typing import Any, Literal, Optional, Type, Union
 
 import ray
 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,
+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 import Simulator
-from ..llm import LLM, LLMConfig, SimpleEmbedding
-from ..memory import Memory
+from ..economy.econ_client import EconomyClient
+from ..llm import SimpleEmbedding
 from ..message import (MessageBlockBase, MessageBlockListenerBase,
                        MessageInterceptor, Messager)
 from ..metrics import init_mlflow_connection
@@ -42,7 +42,7 @@ class AgentSimulation:
         agent_class: Union[None, type[Agent], list[type[Agent]]] = None,
         agent_config_file: Optional[dict] = None,
         metric_extractor: Optional[list[tuple[int, Callable]]] = None,
-        enable_economy: bool = True,
+        enable_institution: bool = True,
         agent_prefix: str = "agent_",
         exp_name: str = "default_experiment",
         logging_level: int = logging.WARNING,
@@ -58,7 +58,7 @@ class AgentSimulation:
         if isinstance(agent_class, list):
             self.agent_class = agent_class
         elif agent_class is None:
-            if enable_economy:
+            if enable_institution:
                 self.agent_class = [
                     SocietyAgent,
                     FirmAgent,
@@ -66,24 +66,30 @@ class AgentSimulation:
                     NBSAgent,
                     GovernmentAgent,
                 ]
-                self.default_memory_config_func = [
-                    memory_config_societyagent,
-                    memory_config_firm,
-                    memory_config_bank,
-                    memory_config_nbs,
-                    memory_config_government,
-                ]
+                self.default_memory_config_func = {
+                    SocietyAgent: memory_config_societyagent,
+                    FirmAgent: memory_config_firm,
+                    BankAgent: memory_config_bank,
+                    NBSAgent: memory_config_nbs,
+                    GovernmentAgent: memory_config_government,
+                }
             else:
                 self.agent_class = [SocietyAgent]
-                self.default_memory_config_func = [memory_config_societyagent]
+                self.default_memory_config_func = {SocietyAgent: memory_config_societyagent}
         else:
             self.agent_class = [agent_class]
         self.agent_config_file = agent_config_file
         self.logging_level = logging_level
         self.config = config
         self.exp_name = exp_name
+        _simulator_config = config["simulator_request"].get("simulator", {})
+        if "server" in _simulator_config:
+            raise ValueError(f"Passing Traffic Simulation address is not supported!")
         self._simulator = Simulator(config["simulator_request"])
-        if enable_economy:
+        self._economy_client = EconomyClient(
+            config["simulator_request"]["simulator"]["server"]
+        )
+        if enable_institution:
             self._economy_addr = economy_addr = self._simulator.server_addr
             if economy_addr is None:
                 raise ValueError(
@@ -202,9 +208,11 @@ class AgentSimulation:
         """Directly run from config file
         Basic config file should contain:
         - simulation_config: file_path
+        - enable_institution: bool, default is True
         - agent_config:
-            - agent_config_file: Optional[dict]
-            - memory_config_func: Optional[Union[Callable, list[Callable]]]
+            - agent_config_file: Optional[dict[type[Agent], str]]
+            - memory_config_func: Optional[dict[type[Agent], Callable]]
+            - metric_extractor: Optional[list[tuple[int, Callable]]]
             - init_func: Optional[list[Callable[AgentSimulation, None]]]
             - group_size: Optional[int]
             - embedding_model: Optional[EmbeddingModel]
@@ -217,7 +225,7 @@ class AgentSimulation:
             - list[Step]
             - Step:
                 - type: str, "step", "run", "interview", "survey", "intervene", "pause", "resume"
-                - day: int if type is "run", else None
+                - days: int if type is "run", else None
                 - times: int if type is "step", else None
                 - description: Optional[str], description of the step
                 - func: Optional[Callable[AgentSimulation, None]], only used when type is "interview", "survey" and "intervene"
@@ -240,6 +248,8 @@ class AgentSimulation:
         simulation = cls(
             config=simulation_config,
             agent_config_file=config["agent_config"].get("agent_config_file", None),
+            metric_extractor=config["agent_config"].get("metric_extractor", None),
+            enable_institution=config.get("enable_institution", True),
             exp_name=config.get("exp_name", "default_experiment"),
             logging_level=config.get("logging_level", logging.WARNING),
         )
@@ -272,7 +282,7 @@ class AgentSimulation:
             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))
+                await simulation.run(step.get("days", 1))
             elif step["type"] == "step":
                 times = step.get("times", 1)
                 for _ in range(times):
@@ -302,6 +312,10 @@ class AgentSimulation:
         self,
     ) -> Path:
         return self._avro_path  # type:ignore
+    
+    @property
+    def economy_client(self):
+        return self._economy_client
 
     @property
     def groups(self):
@@ -406,7 +420,7 @@ class AgentSimulation:
         social_black_list: Optional[list[tuple[str, str]]] = None,
         message_listener: Optional[MessageBlockListenerBase] = None,
         embedding_model: Embeddings = SimpleEmbedding(),
-        memory_config_func: Optional[Union[Callable, list[Callable]]] = None,
+        memory_config_func: Optional[dict[type[Agent], Callable]] = None,
     ) -> None:
         """初始化智能体
 
@@ -415,7 +429,7 @@ class AgentSimulation:
             group_size: 每个组的智能体数量，每一个组为一个独立的ray actor
             pg_sql_writers: 独立的PgSQL writer数量
             message_interceptors: message拦截器数量
-            memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组, 如果为列表，则每个元素表示一个智能体类创建的Memory配置函数
+            memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组, 每个元素表示一个智能体类创建的Memory配置函数
         """
         if not isinstance(agent_count, list):
             agent_count = [agent_count]
@@ -426,17 +440,6 @@ class AgentSimulation:
         if memory_config_func is None:
             memory_config_func = self.default_memory_config_func
 
-        elif not isinstance(memory_config_func, list):
-            logger.warning(
-                "memory_config_func is not a list, using specific memory config function"
-            )
-            memory_config_func = [memory_config_func]
-
-        if len(memory_config_func) != len(agent_count):
-            logger.warning(
-                "The length of memory_config_func and agent_count does not match, using default memory_config"
-            )
-            memory_config_func = self.default_memory_config_func
         # 使用线程池并行创建 AgentGroup
         group_creation_params = []
 
@@ -448,7 +451,7 @@ class AgentSimulation:
         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]
+            memory_config_func_i = memory_config_func.get(agent_class, self.default_memory_config_func[agent_class])
 
             if self.agent_config_file is not None:
                 config_file = self.agent_config_file.get(agent_class, None)
@@ -478,8 +481,8 @@ class AgentSimulation:
 
                 agent_classes = []
                 agent_counts = []
-                memory_config_funcs = []
-                config_files = []
+                memory_config_funcs = {}
+                config_files = {}
 
                 # 分配每种类型的机构智能体到当前组
                 curr_start = start_idx
@@ -487,8 +490,8 @@ class AgentSimulation:
                     if curr_start < count:
                         agent_classes.append(agent_class)
                         agent_counts.append(min(count - curr_start, number_of_agents))
-                        memory_config_funcs.append(mem_func)
-                        config_files.append(conf_file)
+                        memory_config_funcs[agent_class] = mem_func
+                        config_files[agent_class] = conf_file
                     curr_start = max(0, curr_start - count)
 
                 group_creation_params.append(
@@ -513,8 +516,8 @@ class AgentSimulation:
 
                 agent_classes = []
                 agent_counts = []
-                memory_config_funcs = []
-                config_files = []
+                memory_config_funcs = {}
+                config_files = {}
 
                 # 分配每种类型的普通智能体到当前组
                 curr_start = start_idx
@@ -522,8 +525,8 @@ class AgentSimulation:
                     if curr_start < count:
                         agent_classes.append(agent_class)
                         agent_counts.append(min(count - curr_start, number_of_agents))
-                        memory_config_funcs.append(mem_func)
-                        config_files.append(conf_file)
+                        memory_config_funcs[agent_class] = mem_func
+                        config_files[agent_class] = conf_file
                     curr_start = max(0, curr_start - count)
 
                 group_creation_params.append(
@@ -685,14 +688,20 @@ class AgentSimulation:
         group = self._agent_uuid2group[target_agent_uuid]
         await group.update.remote(target_agent_uuid, target_key, content)
 
+    async def economy_update(self, target_agent_id: int, target_key: str, content: Any, mode: Literal["replace", "merge"] = "replace"):
+        """更新指定智能体的经济数据"""
+        self.economy_client.update(
+            id = target_agent_id, 
+            key = target_key, 
+            value=content,
+            mode=mode
+        )
+
     async def send_survey(
-        self, survey: Survey, agent_uuids: Optional[list[str]] = None
+        self, survey: Survey, agent_uuids: list[str] = None
     ):
         """发送问卷"""
-        await self.messager.connect.remote()  # type:ignore
         survey_dict = survey.to_dict()
-        if agent_uuids is None:
-            agent_uuids = self._agent_uuids
         _date_time = datetime.now(timezone.utc)
         payload = {
             "from": SURVEY_SENDER_UUID,

{pycityagent-2.0.0a57.dist-info → pycityagent-2.0.0a59.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pycityagent
-Version: 2.0.0a57
+Version: 2.0.0a59
 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.0a57.dist-info → pycityagent-2.0.0a59.dist-info}/RECORD

@@ -22,9 +22,9 @@ pycityagent/memory/utils.py,sha256=oJWLdPeJy_jcdKcDTo9JAH9kDZhqjoQhhv_zT9qWC0w,8
 pycityagent/memory/const.py,sha256=6zpJPJXWoH9-yf4RARYYff586agCoud9BRn7sPERB1g,932
 pycityagent/memory/faiss_query.py,sha256=V3rIw6d1_xcpNqZBbAYz3qfjVNE7NfJ7xOS5SibPtVU,13180
 pycityagent/memory/state.py,sha256=TYItiyDtehMEQaSBN7PpNrnNxdDM5jGppr9R9Ufv3kA,5134
-pycityagent/simulation/simulation.py,sha256=XOhHYx4zvDv6zgPBagY9Ao109wSZPSbdCH-amsw2pHo,33441
+pycityagent/simulation/simulation.py,sha256=sbrX-e8gV2Wwi8k6ir7TVgQhzejpUZ85SfuiXO4E39I,33983
 pycityagent/simulation/__init__.py,sha256=P5czbcg2d8S0nbbnsQXFIhwzO4CennAhZM8OmKvAeYw,194
-pycityagent/simulation/agentgroup.py,sha256=UnWUQ3Ev4qDQCOasb-EN08aUQi4SWHmM6aNSdBc850Q,31693
+pycityagent/simulation/agentgroup.py,sha256=V-ChLJtJUWJZnmTGllXqGk4XW4iDz9h3V89RgW0e4gg,31376
 pycityagent/simulation/storage/pg.py,sha256=xRshSOGttW-p0re0fNBOjOpb-nQ5msIE2LsdT79_E_Y,8425
 pycityagent/message/message_interceptor.py,sha256=w8XTyZStQtMjILpeAX3VMhAWcYAuaxCgSMwXQU1OryM,8951
 pycityagent/message/__init__.py,sha256=f5QH7DKPqEAMyfSlBMnl3uouOKlsoel909STlIe7nUk,276
@@ -47,7 +47,7 @@ pycityagent/workflow/prompt.py,sha256=6jI0Rq54JLv3-IXqZLYug62vse10wTI83xvf4ZX42n
 pycityagent/workflow/block.py,sha256=4QufS8XnyP6SYp8g1gDODW-H0nAHA7lvivrPGUq1p-w,9922
 pycityagent/workflow/trigger.py,sha256=Df-MOBEDWBbM-v0dFLQLXteLsipymT4n8vqexmK2GiQ,5643
 pycityagent/environment/__init__.py,sha256=MyZBwsweDIHOKSX2iSZs748foNtaiyEcyg6sc747T2g,263
-pycityagent/environment/simulator.py,sha256=d2meQZewFFSvUZa7BkgCHgOkhFf6G_TQyIz_QBFgDI0,12437
+pycityagent/environment/simulator.py,sha256=8imXFHhNxVhritO7YEZqM59x4VkKS8rc0Dwu6SHaREo,12439
 pycityagent/environment/message/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycityagent/environment/utils/port.py,sha256=3OM6kSUt3PxvDUOlgyiendBtETaWU8Mzk_8H0TzTmYg,295
 pycityagent/environment/utils/grpc.py,sha256=6EJwKXXktIWb1NcUiJzIRmfrY0S03QAXXGcCDHqAT00,1998
@@ -81,7 +81,7 @@ pycityagent/cityagent/__init__.py,sha256=gcBQ-a50XegFtjigQ7xDXRBZrywBKqifiQFSRnE
 pycityagent/cityagent/firmagent.py,sha256=UVlNN0lpa4cC4PZVqYzQhbc5VJ2oGsA1731mhbCjnR8,4109
 pycityagent/cityagent/nbsagent.py,sha256=WIXW__6dZ5IrqBqDCjvGbrCshpXzuFRV3Ww6gkYw7p4,4387
 pycityagent/cityagent/initial.py,sha256=7hgCt_tGdnVTXGfEQOn1GTW5dAs1b-ru_FwXxRLI6tM,4549
-pycityagent/cityagent/societyagent.py,sha256=vvFGAr8wlZ82f4qHvpawa1VsHungDegjKn1E51bDUnc,24767
+pycityagent/cityagent/societyagent.py,sha256=DKkPtCPXTpHs83hi3bvXxkQYxKqZI0YldYCyFu4yjsk,24811
 pycityagent/cityagent/message_intercept.py,sha256=HiNOtBXKq40LCmvsw2KMLxMsbmXSVcbRytAO-eI-1sU,3279
 pycityagent/cityagent/governmentagent.py,sha256=HJLuhvEmllu_1KnFEJsYCIasaBJT0BV9Cn_4Y2QGPqg,2791
 pycityagent/cityagent/blocks/dispatcher.py,sha256=mEa1r3tRS3KI1BMZR_w_sbUGzOj6aUJuiUrsHv1n2n0,2943
@@ -97,9 +97,9 @@ pycityagent/cityagent/blocks/mobility_block.py,sha256=xWbARMfJ3-6fddrW3VOUKJrXfM
 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.0a57.dist-info/RECORD,,
-pycityagent-2.0.0a57.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
-pycityagent-2.0.0a57.dist-info/WHEEL,sha256=NW1RskY9zow1Y68W-gXg0oZyBRAugI1JHywIzAIai5o,109
-pycityagent-2.0.0a57.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
-pycityagent-2.0.0a57.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
-pycityagent-2.0.0a57.dist-info/METADATA,sha256=yn1syqalcmvcSdrXMOmcXa7Cj3-IHDE4AIL20UvDaEs,9110
+pycityagent-2.0.0a59.dist-info/RECORD,,
+pycityagent-2.0.0a59.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
+pycityagent-2.0.0a59.dist-info/WHEEL,sha256=NW1RskY9zow1Y68W-gXg0oZyBRAugI1JHywIzAIai5o,109
+pycityagent-2.0.0a59.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
+pycityagent-2.0.0a59.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
+pycityagent-2.0.0a59.dist-info/METADATA,sha256=RMs-AgvrIOTGBqaoqMei0bF2eJRRqIwQn-r3UGBzwr0,9110