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

pycityagent/memory/profile.py

@@ -3,8 +3,7 @@ Agent Profile
 """
 
 from copy import deepcopy
-from typing import (Any, Callable, Dict, List, Optional, Sequence, Tuple,
-                    Union, cast)
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *

pycityagent/memory/self_define.py

@@ -3,8 +3,7 @@ Self Define Data
 """
 
 from copy import deepcopy
-from typing import (Any, Callable, Dict, List, Optional, Sequence, Tuple,
-                    Union, cast)
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *

pycityagent/memory/state.py

@@ -3,8 +3,7 @@ Agent State
 """
 
 from copy import deepcopy
-from typing import (Any, Callable, Dict, List, Optional, Sequence, Tuple,
-                    Union, cast)
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *

pycityagent/message/__init__.py

@@ -1,3 +1,3 @@
 from .messager import Messager
 
-__all__ = ["Messager"]
+__all__ = ["Messager"]

pycityagent/message/messager.py

@@ -4,9 +4,14 @@ import logging
 import math
 from aiomqtt import Client
 
+
 class Messager:
-    def __init__(self, broker, port=1883, timeout=math.inf):
-        self.client = Client(broker, port=port, timeout=timeout)
+    def __init__(
+        self, hostname, port=1883, username=None, password=None, timeout=math.inf
+    ):
+        self.client = Client(
+            hostname, port=port, username=username, password=password, timeout=timeout
+        )
         self.connected = False  # 是否已连接标志
         self.message_queue = asyncio.Queue()  # 用于存储接收到的消息
         self.subscribers = {}  # 订阅者信息，topic -> Agent 映射
@@ -31,7 +36,9 @@ class Messager:
 
     async def subscribe(self, topic, agent):
         if not self.is_connected():
-            logging.error(f"Cannot subscribe to {topic} because not connected to the Broker.")
+            logging.error(
+                f"Cannot subscribe to {topic} because not connected to the Broker."
+            )
             return
         await self.client.subscribe(topic)
         self.subscribers[topic] = agent
@@ -48,7 +55,7 @@ class Messager:
         while not self.message_queue.empty():
             messages.append(await self.message_queue.get())
         return messages
-    
+
     async def send_message(self, topic: str, payload: str, sender_id: int):
         """通过 Messager 发送消息，包含发送者 ID"""
         # 构造消息，payload 中加入 sender_id 以便接收者识别

pycityagent/simulation/__init__.py

@@ -4,4 +4,4 @@
 
 from .simulation import AgentSimulation
 
-__all__ = ["AgentSimulation"] 
+__all__ = ["AgentSimulation"]

pycityagent/simulation/agentgroup.py

@@ -8,15 +8,21 @@ from pycityagent.llm.llm import LLM
 from pycityagent.llm.llmconfig import LLMConfig
 from pycityagent.message import Messager
 
+
 @ray.remote
 class AgentGroup:
     def __init__(self, agents: list[Agent], config: dict, exp_id: str):
         self.agents = agents
         self.config = config
         self.exp_id = exp_id
-        self.messager = Messager(config["simulator_request"]["mqtt"]["server"], config["simulator_request"]["mqtt"]["port"])
+        self.messager = Messager(
+            hostname=config["simulator_request"]["mqtt"]["server"],
+            port=config["simulator_request"]["mqtt"]["port"],
+            username=config["simulator_request"]["mqtt"].get("username", None),
+            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...")
@@ -27,8 +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)
@@ -48,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()
@@ -65,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
@@ -72,14 +97,18 @@ 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("/")
+                payload = payload.decode("utf-8")
+
+            # 提取 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]
@@ -96,15 +125,14 @@ class AgentGroup:
             start_time = await self.simulator.get_time()
             # 计算结束时间（秒）
             end_time = start_time + day * 24 * 3600  # 将天数转换为秒
-            
+
             while True:
                 current_time = await self.simulator.get_time()
                 if current_time >= end_time:
                     break
-                
+
                 await self.step()
 
         except Exception as e:
             logging.error(f"模拟器运行错误: {str(e)}")
             raise
-

pycityagent/simulation/interview.py

@@ -2,20 +2,24 @@ from dataclasses import dataclass
 from datetime import datetime
 from typing import List, Optional
 
+
 @dataclass
 class InterviewRecord:
     """采访记录"""
+
     timestamp: datetime
     agent_name: str
     question: str
     response: str
     blocking: bool
 
+
 class InterviewManager:
     """采访管理器"""
+
     def __init__(self):
         self._history: List[InterviewRecord] = []
-        
+
     def add_record(self, agent_name: str, question: str, response: str, blocking: bool):
         """添加采访记录"""
         record = InterviewRecord(
@@ -23,14 +27,14 @@ class InterviewManager:
             agent_name=agent_name,
             question=question,
             response=response,
-            blocking=blocking
+            blocking=blocking,
         )
         self._history.append(record)
-        
+
     def get_agent_history(self, agent_name: str) -> List[InterviewRecord]:
         """获取指定智能体的采访历史"""
         return [r for r in self._history if r.agent_name == agent_name]
-    
+
     def get_recent_history(self, limit: int = 10) -> List[InterviewRecord]:
         """获取最近的采访记录"""
-        return sorted(self._history, key=lambda x: x.timestamp, reverse=True)[:limit] 
+        return sorted(self._history, key=lambda x: x.timestamp, reverse=True)[:limit]

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
@@ -14,12 +14,19 @@ from .interview import InterviewManager
 from .survey import QuestionType, SurveyManager
 from .ui import InterviewUI
 from .agentgroup import AgentGroup
+
 logger = logging.getLogger(__name__)
 
 
 class AgentSimulation:
     """城市智能体模拟器"""
-    def __init__(self, agent_class: type[Agent], config: dict, agent_prefix: str = "agent_"):
+
+    def __init__(
+        self,
+        agent_class: Union[type[Agent], list[type[Agent]]],
+        config: dict,
+        agent_prefix: str = "agent_",
+    ):
         """
         Args:
             agent_class: 智能体类
@@ -27,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] = {}
@@ -40,88 +50,200 @@ 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, group_size: int = 1000, memory_config_func: Callable = None) -> None:
+    async def init_agents(
+        self,
+        agent_count: Union[int, list[int]],
+        group_size: int = 1000,
+        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()
-            )
-            
-            # 创建智能体时传入Memory配置
-            agent = self.agent_class(
-                name=agent_name,
-                memory=memory,
+        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]
 
-            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
+        if len(memory_config_func) != len(agent_count):
+            logging.warning(
+                "memory_config_func和agent_count的长度不一致，使用默认的memory_config_func"
+            )
+            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配置函数"""
         EXTRA_ATTRIBUTES = {
             # 需求信息
-            "needs": (dict, {
-                'hungry': random.random(),  # 饥饿感
-                'tired': random.random(),   # 疲劳感
-                'safe': random.random(),    # 安全需
-                'social': random.random(),  # 社会需求
-            }, True),
+            "needs": (
+                dict,
+                {
+                    "hungry": random.random(),  # 饥饿感
+                    "tired": random.random(),  # 疲劳感
+                    "safe": random.random(),  # 安全需
+                    "social": random.random(),  # 社会需求
+                },
+                True,
+            ),
             "current_need": (str, "none", True),
             "current_plan": (list, [], True),
             "current_step": (dict, {"intention": "", "type": ""}, True),
-            "execution_context" : (dict, {}, 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 = {
             "gender": random.choice(["male", "female"]),
-            "education": random.choice(["Doctor", "Master", "Bachelor", "College", "High School"]),
+            "education": random.choice(
+                ["Doctor", "Master", "Bachelor", "College", "High School"]
+            ),
             "consumption": random.choice(["sightly low", "low", "medium", "high"]),
-            "occupation": random.choice(["Student", "Teacher", "Doctor", "Engineer", "Manager", "Businessman", "Artist", "Athlete", "Other"]),
+            "occupation": random.choice(
+                [
+                    "Student",
+                    "Teacher",
+                    "Doctor",
+                    "Engineer",
+                    "Manager",
+                    "Businessman",
+                    "Artist",
+                    "Athlete",
+                    "Other",
+                ]
+            ),
             "age": random.randint(18, 65),
-            "skill": random.choice(["Good at problem-solving", "Good at communication", "Good at creativity", "Good at teamwork", "Other"]),
+            "skill": random.choice(
+                [
+                    "Good at problem-solving",
+                    "Good at communication",
+                    "Good at creativity",
+                    "Good at teamwork",
+                    "Other",
+                ]
+            ),
             "family_consumption": random.choice(["low", "medium", "high"]),
-            "personality": random.choice(["outgoint", "introvert", "ambivert", "extrovert"]),
+            "personality": random.choice(
+                ["outgoint", "introvert", "ambivert", "extrovert"]
+            ),
             "income": random.randint(1000, 10000),
             "currency": random.randint(10000, 100000),
             "residence": random.choice(["city", "suburb", "rural"]),
-            "race": random.choice(["Chinese", "American", "British", "French", "German", "Japanese", "Korean", "Russian", "Other"]),
-            "religion": random.choice(["none", "Christian", "Muslim", "Buddhist", "Hindu", "Other"]),
-            "marital_status": random.choice(["not married", "married", "divorced", "widowed"]),
-            }
+            "race": random.choice(
+                [
+                    "Chinese",
+                    "American",
+                    "British",
+                    "French",
+                    "German",
+                    "Japanese",
+                    "Korean",
+                    "Russian",
+                    "Other",
+                ]
+            ),
+            "religion": random.choice(
+                ["none", "Christian", "Muslim", "Buddhist", "Hindu", "Other"]
+            ),
+            "marital_status": random.choice(
+                ["not married", "married", "divorced", "widowed"]
+            ),
+        }
 
         BASE = {
-            "home": {"aoi_position": {"aoi_id": AOI_START_ID + random.randint(1, 50000)}},
-            "work": {"aoi_position": {"aoi_id": AOI_START_ID + random.randint(1, 50000)}},
+            "home": {
+                "aoi_position": {"aoi_id": AOI_START_ID + random.randint(1, 50000)}
+            },
+            "work": {
+                "aoi_position": {"aoi_id": AOI_START_ID + random.randint(1, 50000)}
+            },
         }
 
         return EXTRA_ATTRIBUTES, PROFILE, BASE
@@ -218,7 +340,7 @@ class AgentSimulation:
         except Exception as e:
             logger.error(f"采访过程出错: {str(e)}")
             return f"采访过程出现错误: {str(e)}"
-        
+
     async def submit_survey(self, agent_name: str, survey_id: str) -> str:
         """向智能体提交问卷
 
@@ -319,9 +441,7 @@ class AgentSimulation:
             prevent_thread_lock=True,
             quiet=True,
         )
-        logger.info(
-            f"Gradio Frontend is running on http://{server_name}:{server_port}"
-        )
+        logger.info(f"Gradio Frontend is running on http://{server_name}:{server_port}")
 
     async def step(self):
         """运行一步, 即每个智能体执行一次forward"""
@@ -333,7 +453,7 @@ class AgentSimulation:
         except Exception as e:
             logger.error(f"运行错误: {str(e)}")
             raise
-        
+
     async def run(
         self,
         day: int = 1,
@@ -348,7 +468,7 @@ class AgentSimulation:
             tasks = []
             for group in self._groups.values():
                 tasks.append(group.run.remote(day))
-            
+
             await asyncio.gather(*tasks)
 
         except Exception as e:

pycityagent/simulation/survey/__init__.py

@@ -1,9 +1,4 @@
 from .models import QuestionType, Question, Survey
 from .manager import SurveyManager
 
-__all__ = [
-    'QuestionType',
-    'Question', 
-    'Survey',
-    'SurveyManager'
-]
+__all__ = ["QuestionType", "Question", "Survey", "SurveyManager"]

pycityagent/simulation/survey/manager.py

@@ -4,14 +4,17 @@ import uuid
 import json
 from .models import Survey, Question, QuestionType
 
+
 class SurveyManager:
     def __init__(self):
         self._surveys: Dict[str, Survey] = {}
-        
-    def create_survey(self, title: str, description: str, questions: List[dict]) -> Survey:
+
+    def create_survey(
+        self, title: str, description: str, questions: List[dict]
+    ) -> Survey:
         """创建新问卷"""
         survey_id = str(uuid.uuid4())
-        
+
         # 转换问题数据
         survey_questions = []
         for q in questions:
@@ -21,47 +24,45 @@ class SurveyManager:
                 required=q.get("required", True),
                 options=q.get("options", []),
                 min_rating=q.get("min_rating", 1),
-                max_rating=q.get("max_rating", 5)
+                max_rating=q.get("max_rating", 5),
             )
             survey_questions.append(question)
-            
+
         survey = Survey(
             id=survey_id,
             title=title,
             description=description,
-            questions=survey_questions
+            questions=survey_questions,
         )
-        
+
         self._surveys[survey_id] = survey
         return survey
-        
+
     def get_survey(self, survey_id: str) -> Optional[Survey]:
         """获取指定问卷"""
         return self._surveys.get(survey_id)
-        
+
     def get_all_surveys(self) -> List[Survey]:
         """获取所有问卷"""
         return list(self._surveys.values())
-        
+
     def add_response(self, survey_id: str, agent_name: str, response: dict) -> bool:
         """添加问卷回答"""
         survey = self.get_survey(survey_id)
         if not survey:
             return False
-            
-        survey.responses[agent_name] = {
-            "timestamp": datetime.now(),
-            **response
-        }
+
+        survey.responses[agent_name] = {"timestamp": datetime.now(), **response}
         return True
-        
+
     def export_results(self, survey_id: str) -> str:
         """导出问卷结果"""
         survey = self.get_survey(survey_id)
         if not survey:
             return json.dumps({"error": "问卷不存在"})
-            
-        return json.dumps({
-            "survey": survey.to_dict(),
-            "responses": survey.responses
-        }, ensure_ascii=False, indent=2) 
+
+        return json.dumps(
+            {"survey": survey.to_dict(), "responses": survey.responses},
+            ensure_ascii=False,
+            indent=2,
+        )