pycityagent 2.0.0a58__cp310-cp310-macosx_11_0_arm64.whl → 2.0.0a60__cp310-cp310-macosx_11_0_arm64.whl

pycityagent/agent/agent.py

@@ -94,24 +94,19 @@ class CitizenAgent(Agent):
             logger.warning("Economy client is not set")
             return
         if not self._has_bound_to_economy:
-            if self._has_bound_to_simulator:
-                try:
-                    await self._economy_client.remove_agents([self._agent_id])
-                except:
-                    pass
-                person_id = await self.status.get("id")
-                currency = await self.status.get("currency")
-                await self._economy_client.add_agents(
-                    {
-                        "id": person_id,
-                        "currency": currency,
-                    }
-                )
-                self._has_bound_to_economy = True
-            else:
-                logger.debug(
-                    f"Binding to Economy before binding to Simulator, skip binding to Economy Simulator"
-                )
+            try:
+                await self._economy_client.remove_agents([self._agent_id])
+            except:
+                pass
+            person_id = await self.status.get("id")
+            currency = await self.status.get("currency")
+            await self._economy_client.add_agents(
+                {
+                    "id": person_id,
+                    "currency": currency,
+                }
+            )
+            self._has_bound_to_economy = True
 
     async def handle_gather_message(self, payload: dict):
         """处理收到的消息，识别发送者"""
@@ -176,7 +171,6 @@ class InstitutionAgent(Agent):
         await self._bind_to_economy()
 
     async def _bind_to_economy(self):
-        print("Debug:", self._economy_client, self._has_bound_to_economy)
         if self._economy_client is None:
             logger.debug("Economy client is not set")
             return

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/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,27 +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
-        self._simulator = Simulator(config["simulator_request"])
         _simulator_config = config["simulator_request"].get("simulator", {})
         if "server" in _simulator_config:
             raise ValueError(f"Passing Traffic Simulation address is not supported!")
-        if enable_economy:
+        self._simulator = Simulator(config["simulator_request"])
+        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(
@@ -205,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]
@@ -220,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"
@@ -243,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),
         )
@@ -275,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):
@@ -305,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):
@@ -409,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:
         """初始化智能体
 
@@ -418,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]
@@ -429,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 = []
 
@@ -451,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)
@@ -481,8 +481,8 @@ class AgentSimulation:
 
                 agent_classes = []
                 agent_counts = []
-                memory_config_funcs = []
-                config_files = []
+                memory_config_funcs = {}
+                config_files = {}
 
                 # 分配每种类型的机构智能体到当前组
                 curr_start = start_idx
@@ -490,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(
@@ -516,8 +516,8 @@ class AgentSimulation:
 
                 agent_classes = []
                 agent_counts = []
-                memory_config_funcs = []
-                config_files = []
+                memory_config_funcs = {}
+                config_files = {}
 
                 # 分配每种类型的普通智能体到当前组
                 curr_start = start_idx
@@ -525,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(
@@ -688,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.0a58.dist-info → pycityagent-2.0.0a60.dist-info}/METADATA

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

@@ -1,10 +1,10 @@
-pycityagent-2.0.0a58.dist-info/RECORD,,
-pycityagent-2.0.0a58.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
-pycityagent-2.0.0a58.dist-info/WHEEL,sha256=ezfKMaDztqf77C8lvQ0NCnZxkTaOaKLprqJ8q932MhU,109
-pycityagent-2.0.0a58.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
-pycityagent-2.0.0a58.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
-pycityagent-2.0.0a58.dist-info/METADATA,sha256=AUgPOiWxGpZ5A9LMIruA_h9B0S41Zu_Z_eCY1_Z23iE,9110
-pycityagent/pycityagent-sim,sha256=1Nu-QYC0AuZyVWciGNa2XkYsUntbwAS15Bo7l-y6eok,35449490
+pycityagent-2.0.0a60.dist-info/RECORD,,
+pycityagent-2.0.0a60.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
+pycityagent-2.0.0a60.dist-info/WHEEL,sha256=ezfKMaDztqf77C8lvQ0NCnZxkTaOaKLprqJ8q932MhU,109
+pycityagent-2.0.0a60.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
+pycityagent-2.0.0a60.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
+pycityagent-2.0.0a60.dist-info/METADATA,sha256=oyO0iwWkmdYb_ImPGko3_7auKYsx4ve45hSRRTEmYFs,9110
+pycityagent/pycityagent-sim,sha256=vskCJGHJEh0B2dUfmYlVyrcy3sDZ3kBNwjqcYUZpmO8,35449490
 pycityagent/__init__.py,sha256=PUKWTXc-xdMG7px8oTNclodsILUgypANj2Z647sY63k,808
 pycityagent/pycityagent-ui,sha256=cHZjqtrQ4Fh4qtRahFNCNbT2DNHLmUexiDAa-72Z3RQ,40333378
 pycityagent/metrics/mlflow_client.py,sha256=bTrqYsRfNfjJd6l91hU8vN7JXzJC7DHiM2XbkhMuXIU,4613
@@ -28,9 +28,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=RZvDC6QBHjbm4d4WVdBJICG-X7TOSiQ04elZQBJWbeo,33646
+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
@@ -46,7 +46,7 @@ pycityagent/utils/parsers/parser_base.py,sha256=KBKO4zLZPNdGjPAGqIus8LseZ8W3Tlt2
 pycityagent/utils/parsers/json_parser.py,sha256=tjwyPluYfkWgsvLi0hzfJwFhO3L6yQfZMKza20HaGrY,2911
 pycityagent/agent/agent_base.py,sha256=vC1nMWahwi_nVBUXVdC1CnZpqimKZApixD2m2Iqu_bg,24196
 pycityagent/agent/__init__.py,sha256=U20yKu9QwSqAx_PHk5JwipfODkDfxONtumVfnsKjWFg,180
-pycityagent/agent/agent.py,sha256=FNon_OPvGrh976AEHpaAQyLQN7erbGOX-BxGnCVZza8,12038
+pycityagent/agent/agent.py,sha256=H5AXh92LZZzrGHH8dkwAVTKQIfI8a7xH-83iv50Qejs,11696
 pycityagent/cli/wrapper.py,sha256=9aG6D2cnmh6aIhiVXKPGKuMjVcdAa-TwzANudb2q_FU,1147
 pycityagent/workflow/__init__.py,sha256=H08Ko3eliZvuuCMajbEri-IP4-SeswYU6UjHBNA4Ze0,490
 pycityagent/workflow/prompt.py,sha256=6jI0Rq54JLv3-IXqZLYug62vse10wTI83xvf4ZX42nk,2929
@@ -87,7 +87,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