pycityagent 2.0.0a7__py3-none-any.whl → 2.0.0a8__py3-none-any.whl

pycityagent/agent.py

@@ -287,11 +287,13 @@ class Agent(ABC):
             f"Agent {self._agent_id} received message: '{message}' from Agent {sender_id}"
         )
 
-    async def send_message(self, to_agent_id: int, message: str):
+    async def send_message(
+        self, to_agent_id: int, message: str, sub_topic: str = "chat"
+    ):
         """通过 Messager 发送消息，附带发送者的 ID"""
         if self._messager is None:
             raise RuntimeError("Messager is not set")
-        topic = f"/exps/{self._exp_id}/agents/{to_agent_id}/chat"
+        topic = f"/exps/{self._exp_id}/agents/{to_agent_id}/{sub_topic}"
         await self._messager.send_message(topic, message, self._agent_id)
 
     @abstractmethod
@@ -332,6 +334,13 @@ class CitizenAgent(Agent):
             memory,
         )
 
+    async def handle_gather_message(self, payload: str):
+        """处理收到的消息，识别发送者"""
+        # 从消息中解析发送者 ID 和消息内容
+        target, sender_id = payload.split("|from:")
+        content = await self.memory.get(f"{target}")
+        await self.send_message(int(sender_id), content, "gather")
+
 
 class InstitutionAgent(Agent):
     """
@@ -356,3 +365,16 @@ class InstitutionAgent(Agent):
             simulator,
             memory,
         )
+
+    async def handle_gather_message(self, payload: str):
+        """处理收到的消息，识别发送者"""
+        # 从消息中解析发送者 ID 和消息内容
+        content, sender_id = payload.split("|from:")
+        print(
+            f"Agent {self._agent_id} received gather message: '{content}' from Agent {sender_id}"
+        )
+
+    async def gather_messages(self, agent_ids: list[int], content: str):
+        """从多个智能体收集消息"""
+        for agent_id in agent_ids:
+            await self.send_message(agent_id, content, "gather")

pycityagent/simulation/agentgroup.py

@@ -22,7 +22,7 @@ class AgentGroup:
             password=config["simulator_request"]["mqtt"].get("password", None),
         )
         self.initialized = False
-
+        self.id2agent = {}
         # Step:1 prepare LLM client
         llmConfig = LLMConfig(config["llm_request"])
         logging.info("-----Creating LLM client in remote...")
@@ -33,10 +33,13 @@ class AgentGroup:
         self.simulator = Simulator(config["simulator_request"])
 
         # Step:3 prepare Economy client
-        logging.info("-----Creating Economy client in remote...")
-        self.economy_client = EconomyClient(
-            config["simulator_request"]["economy"]["server"]
-        )
+        if "economy" in config["simulator_request"]:
+            logging.info("-----Creating Economy client in remote...")
+            self.economy_client = EconomyClient(
+                config["simulator_request"]["economy"]["server"]
+            )
+        else:
+            self.economy_client = None
 
         for agent in self.agents:
             agent.set_exp_id(self.exp_id)
@@ -56,8 +59,20 @@ class AgentGroup:
                 agent.set_messager(self.messager)
                 topic = f"/exps/{self.exp_id}/agents/{agent._agent_id}/chat"
                 await self.messager.subscribe(topic, agent)
+                topic = f"/exps/{self.exp_id}/agents/{agent._agent_id}/gather"
+                await self.messager.subscribe(topic, agent)
         self.initialized = True
 
+    async def gather(self, content: str):
+        results = {}
+        for agent in self.agents:
+            results[agent._agent_id] = await agent.memory.get(content)
+        return results
+
+    async def update(self, target_agent_id: str, target_key: str, content: any):
+        agent = self.id2agent[target_agent_id]
+        await agent.memory.update(target_key, content)
+
     async def step(self):
         if not self.initialized:
             await self.init_agents()
@@ -73,6 +88,8 @@ class AgentGroup:
         # Step 2: 从 Messager 获取消息
         messages = await self.messager.fetch_messages()
 
+        print(f"Received {len(messages)} messages")
+
         # Step 3: 分发消息到对应的 Agent
         for message in messages:
             topic = message.topic.value
@@ -81,13 +98,17 @@ class AgentGroup:
             # 添加解码步骤，将bytes转换为str
             if isinstance(payload, bytes):
                 payload = payload.decode("utf-8")
-            # 提取 agent_id（主题格式为 "/exps/{exp_id}/agents/{agent_id}/chat"）
-            _, _, _, agent_id, _ = topic.strip("/").split("/")
+
+            # 提取 agent_id（主题格式为 "/exps/{exp_id}/agents/{agent_id}/chat" 或 "/exps/{exp_id}/agents/{agent_id}/gather"）
+            _, _, _, agent_id, topic_type = topic.strip("/").split("/")
             agent_id = int(agent_id)
 
             if agent_id in self.id2agent:
                 agent = self.id2agent[agent_id]
-                await agent.handle_message(payload)
+                if topic_type == "chat":
+                    await agent.handle_message(payload)
+                elif topic_type == "gather":
+                    await agent.handle_gather_message(payload)
 
         # Step 4: 调用每个 Agent 的运行逻辑
         tasks = [agent.run() for agent in self.agents]

pycityagent/simulation/simulation.py

@@ -4,7 +4,7 @@ import logging
 import uuid
 from datetime import datetime
 import random
-from typing import Dict, List, Optional, Callable
+from typing import Dict, List, Optional, Callable, Union
 from mosstool.map._map_util.const import AOI_START_ID
 
 from pycityagent.memory.memory import Memory
@@ -22,7 +22,10 @@ class AgentSimulation:
     """城市智能体模拟器"""
 
     def __init__(
-        self, agent_class: type[Agent], config: dict, agent_prefix: str = "agent_"
+        self,
+        agent_class: Union[type[Agent], list[type[Agent]]],
+        config: dict,
+        agent_prefix: str = "agent_",
     ):
         """
         Args:
@@ -31,7 +34,10 @@ class AgentSimulation:
             agent_prefix: 智能体名称前缀
         """
         self.exp_id = uuid.uuid4()
-        self.agent_class = agent_class
+        if isinstance(agent_class, list):
+            self.agent_class = agent_class
+        else:
+            self.agent_class = [agent_class]
         self.config = config
         self.agent_prefix = agent_prefix
         self._agents: Dict[str, Agent] = {}
@@ -44,53 +50,105 @@ class AgentSimulation:
         self._loop = asyncio.get_event_loop()
         self._blocked_agents: List[str] = []  # 新增：持续阻塞的智能体列表
         self._survey_manager = SurveyManager()
+        self._agentid2group: Dict[str, AgentGroup] = {}
+        self._agent_ids: List[str] = []
 
     async def init_agents(
         self,
-        agent_count: int,
+        agent_count: Union[int, list[int]],
         group_size: int = 1000,
-        memory_config_func: Callable = None,
+        memory_config_func: Union[Callable, list[Callable]] = None,
     ) -> None:
         """初始化智能体
 
         Args:
-            agent_count: 要创建的总智能体数量
+            agent_count: 要创建的总智能体数量, 如果为列表，则每个元素表示一个智能体类创建的智能体数量
             group_size: 每个组的智能体数量，每一个组为一个独立的ray actor
-            memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组
+            memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组, 如果为列表，则每个元素表示一个智能体类创建的Memory配置函数
         """
-        if memory_config_func is None:
-            memory_config_func = self.default_memory_config_func
+        if not isinstance(agent_count, list):
+            agent_count = [agent_count]
 
-        for i in range(agent_count):
-            agent_name = f"{self.agent_prefix}{i}"
+        if len(self.agent_class) != len(agent_count):
+            raise ValueError("agent_class和agent_count的长度不一致")
 
-            # 获取Memory配置
-            extra_attributes, profile, base = memory_config_func()
-            memory = Memory(
-                config=extra_attributes, profile=profile.copy(), base=base.copy()
+        if memory_config_func is None:
+            logging.warning(
+                "memory_config_func is None, using default memory config function"
             )
+            memory_config_func = [self.default_memory_config_func]
+        elif not isinstance(memory_config_func, list):
+            memory_config_func = [memory_config_func]
 
-            # 创建智能体时传入Memory配置
-            agent = self.agent_class(
-                name=agent_name,
-                memory=memory,
+        if len(memory_config_func) != len(agent_count):
+            logging.warning(
+                "memory_config_func和agent_count的长度不一致，使用默认的memory_config_func"
             )
-
-            self._agents[agent_name] = agent
-
-        # 计算需要的组数,向上取整以处理不足一组的情况
-        num_group = (agent_count + group_size - 1) // group_size
-
-        for i in range(num_group):
-            # 计算当前组的起始和结束索引
-            start_idx = i * group_size
-            end_idx = min((i + 1) * group_size, agent_count)
-
-            # 获取当前组的agents
-            agents = list(self._agents.values())[start_idx:end_idx]
-            group_name = f"{self.agent_prefix}_group_{i}"
-            group = AgentGroup.remote(agents, self.config, self.exp_id)
-            self._groups[group_name] = group
+            memory_config_func = [self.default_memory_config_func] * len(agent_count)
+
+        class_init_index = 0
+        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]
+            for j in range(agent_count_i):
+                agent_name = f"{self.agent_prefix}_{i}_{j}"
+
+                # 获取Memory配置
+                extra_attributes, profile, base = memory_config_func_i()
+                memory = Memory(
+                    config=extra_attributes, profile=profile.copy(), base=base.copy()
+                )
+
+                # 创建智能体时传入Memory配置
+                agent = agent_class(
+                    name=agent_name,
+                    memory=memory,
+                )
+
+                self._agents[agent_name] = agent
+
+            # 计算需要的组数,向上取整以处理不足一组的情况
+            num_group = (agent_count_i + group_size - 1) // group_size
+
+            for k in range(num_group):
+                # 计算当前组的起始和结束索引
+                start_idx = class_init_index + k * group_size
+                end_idx = min(
+                    class_init_index + start_idx + group_size,
+                    class_init_index + agent_count_i,
+                )
+
+                # 获取当前组的agents
+                agents = list(self._agents.values())[start_idx:end_idx]
+                group_name = f"{self.agent_prefix}_{i}_group_{k}"
+                group = AgentGroup.remote(agents, self.config, self.exp_id)
+                self._groups[group_name] = group
+
+            class_init_index += agent_count_i  # 更新类初始索引
+
+        init_tasks = []
+        for group in self._groups.values():
+            init_tasks.append(group.init_agents.remote())
+        await asyncio.gather(*init_tasks)
+
+        for group in self._groups.values():
+            agent_ids = await group.gather.remote("id")
+            for agent_id in agent_ids:
+                self._agent_ids.append(agent_id)
+                self._agentid2group[agent_id] = group
+
+    async def gather(self, content: str):
+        """收集所有智能体的ID"""
+        gather_tasks = []
+        for group in self._groups.values():
+            gather_tasks.append(group.gather.remote(content))
+        return await asyncio.gather(*gather_tasks)
+
+    async def update(self, target_agent_id: str, target_key: str, content: any):
+        """更新指定智能体的记忆"""
+        group = self._agentid2group[target_agent_id]
+        await group.update.remote(target_agent_id, target_key, content)
 
     def default_memory_config_func(self):
         """默认的Memory配置函数"""
@@ -111,6 +169,15 @@ class AgentSimulation:
             "current_step": (dict, {"intention": "", "type": ""}, True),
             "execution_context": (dict, {}, True),
             "plan_history": (list, [], True),
+            # cognition
+            "fulfillment": (int, 5, True),
+            "emotion": (int, 5, True),
+            "attitude": (int, 5, True),
+            "thought": (str, "Currently nothing good or bad is happening", True),
+            "emotion_types": (str, "Relief", True),
+            "incident": (list, [], True),
+            # social
+            "friends": (list, [], True),
         }
 
         PROFILE = {

{pycityagent-2.0.0a7.dist-info → pycityagent-2.0.0a8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a7
+Version: 2.0.0a8
 Summary: LLM-based城市环境agent构建库
 License: MIT
 Author: Yuwei Yan

{pycityagent-2.0.0a7.dist-info → pycityagent-2.0.0a8.dist-info}/RECORD

@@ -1,5 +1,5 @@
 pycityagent/__init__.py,sha256=n56bWkAUEcvjDsb7LcJpaGjlrriSKPnR0yBhwRfEYBA,212
-pycityagent/agent.py,sha256=s7FYOkOwdOU6ufp6HzGx0GNt3LZc4qMxRRv8hZxJI7k,11281
+pycityagent/agent.py,sha256=cn8g9Sn8g3L5LaCb-zN831bW_sZWuUZbRydS3pegHdQ,12232
 pycityagent/economy/__init__.py,sha256=aonY4WHnx-6EGJ4WKrx4S-2jAkYNLtqUA04jp6q8B7w,75
 pycityagent/economy/econ_client.py,sha256=qQb_kZneEXGBRaS_y5Jdoi95I8GyjKEsDSC4s6V6R7w,10829
 pycityagent/environment/__init__.py,sha256=awHxlOud-btWbk0FCS4RmGJ13W84oVCkbGfcrhKqihA,240
@@ -46,9 +46,9 @@ pycityagent/memory/utils.py,sha256=97lkenn-36wgt7uWb3Z39BXdJ5zlEQTQnQBFpoND1gg,8
 pycityagent/message/__init__.py,sha256=TCjazxqb5DVwbTu1fF0sNvaH_EPXVuj2XQ0p6W-QCLU,55
 pycityagent/message/messager.py,sha256=2hUIuX8Jqad2p6KM8PyfOWMFl1NjqF9B5yE-FnsYx3c,2598
 pycityagent/simulation/__init__.py,sha256=jYaqaNpzM5M_e_ykISS_M-mIyYdzJXJWhgpfBpA6l5k,111
-pycityagent/simulation/agentgroup.py,sha256=bVHXjiYtvqvvYOMsQPzufOFG8KZZDRUDTYjwpIjlQgE,4308
+pycityagent/simulation/agentgroup.py,sha256=Hu5IXFYZlaZ22m0xlDPBu89rRr1bJdSDeWWJIBxv92k,5233
 pycityagent/simulation/interview.py,sha256=S2uv8MFCB4-u_4Q202VFoPJOIleqpKK9Piju0BDSb_0,1158
-pycityagent/simulation/simulation.py,sha256=5DRRG4iTWvTTa10G4_qsoUy9wa2dmpDpfPWcK2Rnfvs,13860
+pycityagent/simulation/simulation.py,sha256=7ONLou5qoA0dcVp3-mr6AFop6q3JqTyLudxfnbf9trc,16910
 pycityagent/simulation/survey/__init__.py,sha256=rxwou8U9KeFSP7rMzXtmtp2fVFZxK4Trzi-psx9LPIs,153
 pycityagent/simulation/survey/manager.py,sha256=Tqini-4-uBZDsChVVL4ezlgXatnrxAfvceAZZRP8E48,2062
 pycityagent/simulation/survey/models.py,sha256=sY4OrrG1h9iBnjBsyDage4T3mUFPBHHZQe-ORtwSjKc,1305
@@ -65,6 +65,6 @@ pycityagent/workflow/block.py,sha256=6EmiRMLdOZC1wMlmLMIjfrp9TuiI7Gw4s3nnXVMbrnw
 pycityagent/workflow/prompt.py,sha256=tY69nDO8fgYfF_dOA-iceR8pAhkYmCqoox8uRPqEuGY,2956
 pycityagent/workflow/tool.py,sha256=wD9WZ5rma6HCKugtHTwbShNE0f-Rjlwvn_1be3fCAsk,6682
 pycityagent/workflow/trigger.py,sha256=t5X_i0WtL32bipZSsq_E3UUyYYudYLxQUpvxbgClp2s,5683
-pycityagent-2.0.0a7.dist-info/METADATA,sha256=jl2Q3mkfZidXJQGQuTTNMXCaGD6z7CMYy4R9jk3I0vA,7614
-pycityagent-2.0.0a7.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-pycityagent-2.0.0a7.dist-info/RECORD,,
+pycityagent-2.0.0a8.dist-info/METADATA,sha256=IR7pOLENcFN0uxU3lshslneGOmYHrRjuuyJ13x5BWSc,7614
+pycityagent-2.0.0a8.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+pycityagent-2.0.0a8.dist-info/RECORD,,