pycityagent 2.0.0a60__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a61__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/cityagent/message_intercept.py

@@ -55,6 +55,10 @@ class EdgeMessageBlock(MessageBlockBase):
 
 
 class PointMessageBlock(MessageBlockBase):
+    def __init__(self, name: str = "", max_violation_time: int = 3) -> None:
+        super().__init__(name)
+        self.max_violation_time = max_violation_time
+
     async def forward(  # type:ignore
         self,
         from_uuid: str,
@@ -74,7 +78,10 @@ class PointMessageBlock(MessageBlockBase):
                 llm_client=self.llm,
                 content=msg,
             )
-            if not is_valid and violation_counts[from_uuid] >= 3 - 1:
+            if (
+                not is_valid
+                and violation_counts[from_uuid] >= self.max_violation_time - 1
+            ):
                 # 直接添加即可 在框架内部的异步锁保证不会冲突
                 black_list.append((from_uuid, to_uuid))
             return is_valid

pycityagent/simulation/simulation.py

@@ -13,13 +13,16 @@ 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 ..cityagent.message_intercept import (EdgeMessageBlock,
+                                           MessageBlockListener,
+                                           PointMessageBlock)
 from ..economy.econ_client import EconomyClient
+from ..environment import Simulator
 from ..llm import SimpleEmbedding
 from ..message import (MessageBlockBase, MessageBlockListenerBase,
                        MessageInterceptor, Messager)
@@ -75,7 +78,9 @@ class AgentSimulation:
                 }
             else:
                 self.agent_class = [SocietyAgent]
-                self.default_memory_config_func = {SocietyAgent: 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
@@ -229,6 +234,9 @@ class AgentSimulation:
                 - 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"
+        - message_intercept
+            - mode: "point"|"edge"
+            - max_violation_time: Optional[int], default to 3. The maximum time for someone to send bad message before banned. Used only in `point` mode.
         - logging_level: Optional[int]
         - exp_name: Optional[str]
         """
@@ -260,7 +268,31 @@ class AgentSimulation:
         agent_count.append(config["agent_config"]["number_of_government"])
         agent_count.append(config["agent_config"]["number_of_bank"])
         agent_count.append(config["agent_config"]["number_of_nbs"])
-        # TODO(yanjunbo): support MessageInterceptor
+        # support MessageInterceptor
+        if "message_intercept" in config:
+            _intercept_config = config["message_intercept"]
+            _mode = _intercept_config.get("mode", "point")
+            if _mode == "point":
+                _kwargs = {
+                    k: v
+                    for k, v in _intercept_config.items()
+                    if k
+                    in {
+                        "max_violation_time",
+                    }
+                }
+                _interceptor_blocks = [PointMessageBlock(**_kwargs)]
+            elif _mode == "edge":
+                _kwargs = {}
+                _interceptor_blocks = [EdgeMessageBlock(**_kwargs)]
+            else:
+                raise ValueError(f"Unsupported interception mode `{_mode}!`")
+            _message_intercept_kwargs = {
+                "message_interceptor_blocks": _interceptor_blocks,
+                "message_listener": MessageBlockListener(),
+            }
+        else:
+            _message_intercept_kwargs = {}
         await simulation.init_agents(
             agent_count=agent_count,
             group_size=config["agent_config"].get("group_size", 10000),
@@ -268,6 +300,7 @@ class AgentSimulation:
                 "embedding_model", SimpleEmbedding()
             ),
             memory_config_func=config["agent_config"].get("memory_config_func", None),
+            **_message_intercept_kwargs,
         )
         logger.info("Running Init Functions...")
         for init_func in config["agent_config"].get(
@@ -312,7 +345,7 @@ class AgentSimulation:
         self,
     ) -> Path:
         return self._avro_path  # type:ignore
-    
+
     @property
     def economy_client(self):
         return self._economy_client
@@ -438,7 +471,7 @@ class AgentSimulation:
             raise ValueError("agent_class和agent_count的长度不一致")
 
         if memory_config_func is None:
-            memory_config_func = self.default_memory_config_func
+            memory_config_func = self.default_memory_config_func  # type:ignore
 
         # 使用线程池并行创建 AgentGroup
         group_creation_params = []
@@ -451,7 +484,10 @@ 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.get(agent_class, self.default_memory_config_func[agent_class])
+            assert memory_config_func is not None
+            memory_config_func_i = memory_config_func.get(
+                agent_class, self.default_memory_config_func[agent_class]  # type:ignore
+            )
 
             if self.agent_config_file is not None:
                 config_file = self.agent_config_file.get(agent_class, None)
@@ -564,6 +600,7 @@ class AgentSimulation:
             _num_workers = 1
             self._pgsql_writers = _workers = [None for _ in range(_num_workers)]
         # message interceptor
+        self.message_listener = message_listener
         if message_listener is not None:
             self._message_abort_listening_queue = _queue = ray.util.queue.Queue()  # type: ignore
             await message_listener.set_queue(_queue)
@@ -688,18 +725,19 @@ 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"):
+    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
+        await self.economy_client.update(
+            id=target_agent_id, key=target_key, value=content, mode=mode
         )
 
-    async def send_survey(
-        self, survey: Survey, agent_uuids: list[str] = None
-    ):
+    async def send_survey(self, survey: Survey, agent_uuids: list[str] = []):
         """发送问卷"""
         survey_dict = survey.to_dict()
         _date_time = datetime.now(timezone.utc)

{pycityagent-2.0.0a60.dist-info → pycityagent-2.0.0a61.dist-info}/METADATA

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

@@ -1,9 +1,3 @@
-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=VujM3ypTCyUW6hcTDdK2ej0ARVMxlU1Djlh_zWnDgqk,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
@@ -28,7 +22,7 @@ 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=sbrX-e8gV2Wwi8k6ir7TVgQhzejpUZ85SfuiXO4E39I,33983
+pycityagent/simulation/simulation.py,sha256=yDQBzIpxAEvr3oEt2PChgAhTnPgy2t4n-nbgCiqiE6s,35578
 pycityagent/simulation/__init__.py,sha256=P5czbcg2d8S0nbbnsQXFIhwzO4CennAhZM8OmKvAeYw,194
 pycityagent/simulation/agentgroup.py,sha256=V-ChLJtJUWJZnmTGllXqGk4XW4iDz9h3V89RgW0e4gg,31376
 pycityagent/simulation/storage/pg.py,sha256=xRshSOGttW-p0re0fNBOjOpb-nQ5msIE2LsdT79_E_Y,8425
@@ -88,7 +82,7 @@ pycityagent/cityagent/firmagent.py,sha256=UVlNN0lpa4cC4PZVqYzQhbc5VJ2oGsA1731mhb
 pycityagent/cityagent/nbsagent.py,sha256=WIXW__6dZ5IrqBqDCjvGbrCshpXzuFRV3Ww6gkYw7p4,4387
 pycityagent/cityagent/initial.py,sha256=7hgCt_tGdnVTXGfEQOn1GTW5dAs1b-ru_FwXxRLI6tM,4549
 pycityagent/cityagent/societyagent.py,sha256=DKkPtCPXTpHs83hi3bvXxkQYxKqZI0YldYCyFu4yjsk,24811
-pycityagent/cityagent/message_intercept.py,sha256=HiNOtBXKq40LCmvsw2KMLxMsbmXSVcbRytAO-eI-1sU,3279
+pycityagent/cityagent/message_intercept.py,sha256=2RI4--5SVQy9aLXulxxldzmwOynwAKFqtmBaFUnEtsk,3511
 pycityagent/cityagent/governmentagent.py,sha256=HJLuhvEmllu_1KnFEJsYCIasaBJT0BV9Cn_4Y2QGPqg,2791
 pycityagent/cityagent/blocks/dispatcher.py,sha256=mEa1r3tRS3KI1BMZR_w_sbUGzOj6aUJuiUrsHv1n2n0,2943
 pycityagent/cityagent/blocks/needs_block.py,sha256=s8LikgtKORfo_Sw9SQ5_3biNPTof15QuUs4cDynXCyM,15332
@@ -103,3 +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.0a61.dist-info/RECORD,,
+pycityagent-2.0.0a61.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
+pycityagent-2.0.0a61.dist-info/WHEEL,sha256=VujM3ypTCyUW6hcTDdK2ej0ARVMxlU1Djlh_zWnDgqk,109
+pycityagent-2.0.0a61.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
+pycityagent-2.0.0a61.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
+pycityagent-2.0.0a61.dist-info/METADATA,sha256=hvkHrK65VERp64_mKCo6cwULjsAl_GTwsui2NCwOek4,9110