pycityagent 2.0.0a47__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a49__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/simulation/simulation.py

@@ -1,21 +1,16 @@
 import asyncio
 import json
 import logging
-import os
-import random
 import time
 import uuid
 from collections.abc import Callable, Sequence
-from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Optional, Type, Union
 
-import pycityproto.city.economy.v2.economy_pb2 as economyv2
 import ray
 import yaml
 from langchain_core.embeddings import Embeddings
-from mosstool.map._map_util.const import AOI_START_ID
 
 from ..agent import Agent, InstitutionAgent
 from ..environment.simulator import Simulator
@@ -27,17 +22,20 @@ from ..survey import Survey
 from ..utils import TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
 from .agentgroup import AgentGroup
 from .storage.pg import PgWriter, create_pg_tables
+from ..cityagent import SocietyAgent, FirmAgent, BankAgent, NBSAgent, GovernmentAgent, memory_config_societyagent, memory_config_government, memory_config_firm, memory_config_bank, memory_config_nbs
+from ..cityagent.initial import bind_agent_info, initialize_social_network
 
 logger = logging.getLogger("pycityagent")
 
-
 class AgentSimulation:
     """城市智能体模拟器"""
 
     def __init__(
         self,
-        agent_class: Union[type[Agent], list[type[Agent]]],
         config: dict,
+        agent_class: Union[None, type[Agent], list[type[Agent]]] = None,
+        agent_config_file: Optional[dict] = None,
+        enable_economy: bool = True,
         agent_prefix: str = "agent_",
         exp_name: str = "default_experiment",
         logging_level: int = logging.WARNING,
@@ -52,20 +50,34 @@ class AgentSimulation:
         self.exp_id = str(uuid.uuid4())
         if isinstance(agent_class, list):
             self.agent_class = agent_class
+        elif agent_class is None:
+            if enable_economy:
+                self.agent_class = [SocietyAgent, FirmAgent, BankAgent, NBSAgent, GovernmentAgent]
+                self.default_memory_config_func = [
+                    memory_config_societyagent,
+                    memory_config_firm,
+                    memory_config_bank,
+                    memory_config_nbs,
+                    memory_config_government,
+                ]
+            else:
+                self.agent_class = [SocietyAgent]
+                self.default_memory_config_func = [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"])
         self.agent_prefix = agent_prefix
-        self._agents: dict[uuid.UUID, Agent] = {}
         self._groups: dict[str, AgentGroup] = {}  # type:ignore
-        self._agent_uuid2group: dict[uuid.UUID, AgentGroup] = {}  # type:ignore
-        self._agent_uuids: list[uuid.UUID] = []
-        self._user_chat_topics: dict[uuid.UUID, str] = {}
-        self._user_survey_topics: dict[uuid.UUID, str] = {}
-        self._user_interview_topics: dict[uuid.UUID, str] = {}
+        self._agent_uuid2group: dict[str, AgentGroup] = {}  # type:ignore
+        self._agent_uuids: list[str] = []
+        self._type2group: dict[Type[Agent], AgentGroup] = {}
+        self._user_chat_topics: dict[str, str] = {}
+        self._user_survey_topics: dict[str, str] = {}
+        self._user_interview_topics: dict[str, str] = {}
         self._loop = asyncio.get_event_loop()
         # self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
 
@@ -126,6 +138,80 @@ class AgentSimulation:
             with open(self._exp_info_file, "w") as f:
                 yaml.dump(self._exp_info, f)
 
+    @classmethod
+    async def run_from_config(cls, config: dict):
+        """Directly run from config file
+        Basic config file should contain:
+        - simulation_config: file_path
+        - agent_config: 
+            - agent_config_file: Optional[dict]
+            - memory_config_func: Optional[Union[Callable, list[Callable]]]
+            - init_func: Optional[list[Callable[AgentSimulation, None]]]
+            - group_size: Optional[int]
+            - embedding_model: Optional[EmbeddingModel]
+            - number_of_citizen: required, int
+            - number_of_firm: required, int
+            - number_of_government: required, int
+            - number_of_bank: required, int
+            - number_of_nbs: required, int
+        - workflow:
+            - list[Step]
+            - Step:
+                - type: str, "step", "run", "interview", "survey", "intervene"
+                - day: int if type is "run", else None
+                - time: int if type is "step", else None
+                - description: Optional[str], description of the step
+                - step_func: Optional[Callable[AgentSimulation, None]], only used when type is "interview", "survey" and "intervene"
+        - logging_level: Optional[int]
+        - exp_name: Optional[str]
+        """
+        # required key check
+        if "simulation_config" not in config:
+            raise ValueError("simulation_config is required")
+        if "agent_config" not in config:
+            raise ValueError("agent_config is required")
+        if "workflow" not in config:
+            raise ValueError("workflow is required")
+        import yaml
+        logger.info("Loading config file...")
+        with open(config["simulation_config"], "r") as f:
+            simulation_config = yaml.safe_load(f)
+        logger.info("Creating AgentSimulation Task...")
+        simulation = cls(
+            config=simulation_config, 
+            agent_config_file=config["agent_config"].get("agent_config_file", None), 
+            exp_name=config.get("exp_name", "default_experiment"),
+            logging_level=config.get("logging_level", logging.WARNING),
+        )
+        logger.info("Initializing Agents...")
+        agent_count = []
+        agent_count.append(config["agent_config"]["number_of_citizen"])
+        agent_count.append(config["agent_config"]["number_of_firm"])
+        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"])
+        await simulation.init_agents(
+            agent_count=agent_count,
+            group_size=config["agent_config"].get("group_size", 10000),
+            embedding_model=config["agent_config"].get("embedding_model", SimpleEmbedding()),
+            memory_config_func=config["agent_config"].get("memory_config_func", None),
+        )
+        logger.info("Running Init Functions...")
+        for init_func in config["agent_config"].get("init_func", [bind_agent_info, initialize_social_network]):
+            await init_func(simulation)
+        logger.info("Starting Simulation...")
+        for step in config["workflow"]:
+            logger.info(f"Running step: type: {step['type']} - description: {step.get('description', 'no description')}")
+            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))
+            elif step["type"] == "step":
+                await simulation.step(step.get("time", 1))
+            else:
+                await step["step_func"](simulation)
+        logger.info("Simulation finished")
+
     @property
     def enable_avro(
         self,
@@ -138,10 +224,6 @@ class AgentSimulation:
     ) -> bool:
         return self._enable_pgsql
 
-    @property
-    def agents(self) -> dict[uuid.UUID, Agent]:
-        return self._agents
-
     @property
     def avro_path(
         self,
@@ -159,42 +241,82 @@ class AgentSimulation:
     @property
     def agent_uuid2group(self):
         return self._agent_uuid2group
+    
+    @property
+    def messager(self):
+        return self._messager
+    
+    async def _save_exp_info(self) -> None:
+        """异步保存实验信息到YAML文件"""
+        try:
+            if self.enable_avro:
+                with open(self._exp_info_file, "w") as f:
+                    yaml.dump(self._exp_info, f)
+        except Exception as e:
+            logger.error(f"Avro保存实验信息失败: {str(e)}")
+        try:
+            if self.enable_pgsql:
+                worker: ray.ObjectRef = self._pgsql_writers[0]  # type:ignore
+                pg_exp_info = {
+                    k: self._exp_info[k] for (k, _) in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                }
+                pg_exp_info["created_at"] = self._exp_created_time
+                pg_exp_info["updated_at"] = self._exp_updated_time
+                await worker.async_update_exp_info.remote(  # type:ignore
+                    pg_exp_info
+                )
+        except Exception as e:
+            logger.error(f"PostgreSQL保存实验信息失败: {str(e)}")
 
-    def create_remote_group(
-        self,
-        group_name: str,
-        agents: list[Agent],
-        config: dict,
-        exp_id: str,
-        exp_name: str,
-        enable_avro: bool,
-        avro_path: Path,
-        enable_pgsql: bool,
-        pgsql_writer: ray.ObjectRef,
-        mlflow_run_id: str = None,  # type: ignore
-        embedding_model: Embeddings = None,  # type: ignore
-        logging_level: int = logging.WARNING,
-    ):
-        """创建远程组"""
-        group = AgentGroup.remote(
-            agents,
-            config,
-            exp_id,
-            exp_name,
-            enable_avro,
-            avro_path,
-            enable_pgsql,
-            pgsql_writer,
-            mlflow_run_id,
-            embedding_model,
-            logging_level,
-        )
-        return group_name, group, agents
+    async def _update_exp_status(self, status: int, error: str = "") -> None:
+        self._exp_updated_time = datetime.now(timezone.utc)
+        """更新实验状态并保存"""
+        self._exp_info["status"] = status
+        self._exp_info["error"] = error
+        self._exp_info["updated_at"] = self._exp_updated_time.isoformat()
+        await self._save_exp_info()
+
+    async def _monitor_exp_status(self, stop_event: asyncio.Event):
+        """监控实验状态并更新
+
+        Args:
+            stop_event: 用于通知监控任务停止的事件
+        """
+        try:
+            while not stop_event.is_set():
+                # 更新实验状态
+                # 假设所有group的cur_day和cur_t是同步的，取第一个即可
+                self._exp_info["cur_day"] = await self._simulator.get_simulator_day()
+                self._exp_info["cur_t"] = (
+                    await self._simulator.get_simulator_second_from_start_of_day()
+                )
+                await self._save_exp_info()
+
+                await asyncio.sleep(1)  # 避免过于频繁的更新
+        except asyncio.CancelledError:
+            # 正常取消，不需要特殊处理
+            pass
+        except Exception as e:
+            logger.error(f"监控实验状态时发生错误: {str(e)}")
+            raise
+
+    async def __aenter__(self):
+        """异步上下文管理器入口"""
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """异步上下文管理器出口"""
+        if exc_type is not None:
+            # 如果发生异常，更新状态为错误
+            await self._update_exp_status(3, str(exc_val))
+        elif self._exp_info["status"] != 3:
+            # 如果没有发生异常且状态不是错误，则更新为完成
+            await self._update_exp_status(2)
 
     async def init_agents(
         self,
         agent_count: Union[int, list[int]],
-        group_size: int = 1000,
+        group_size: int = 10000,
         pg_sql_writers: int = 32,
         embedding_model: Embeddings = SimpleEmbedding(),
         memory_config_func: Optional[Union[Callable, list[Callable]]] = None,
@@ -206,7 +328,6 @@ class AgentSimulation:
             group_size: 每个组的智能体数量，每一个组为一个独立的ray actor
             memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组, 如果为列表，则每个元素表示一个智能体类创建的Memory配置函数
         """
-        await self._messager.connect.remote()
         if not isinstance(agent_count, list):
             agent_count = [agent_count]
 
@@ -217,67 +338,104 @@ class AgentSimulation:
             logger.warning(
                 "memory_config_func is None, using default memory config function"
             )
-            memory_config_func = []
-            for agent_class in self.agent_class:
-                if issubclass(agent_class, InstitutionAgent):
-                    memory_config_func.append(self.default_memory_config_institution)
-                else:
-                    memory_config_func.append(self.default_memory_config_citizen)
+            memory_config_func = self.default_memory_config_func
+
         elif not isinstance(memory_config_func, list):
             memory_config_func = [memory_config_func]
 
         if len(memory_config_func) != len(agent_count):
             logger.warning(
-                "memory_config_func和agent_count的长度不一致，使用默认的memory_config"
+                "The length of memory_config_func and agent_count does not match, using default memory_config"
             )
-            memory_config_func = []
-            for agent_class in self.agent_class:
-                if issubclass(agent_class, InstitutionAgent):
-                    memory_config_func.append(self.default_memory_config_institution)
-                else:
-                    memory_config_func.append(self.default_memory_config_citizen)
+            memory_config_func = self.default_memory_config_func
         # 使用线程池并行创建 AgentGroup
         group_creation_params = []
-        class_init_index = 0
 
-        # 首先收集所有需要创建的组的参数
+        # 分别处理机构智能体和普通智能体
+        institution_params = []
+        citizen_params = []
+        
+        # 收集所有参数
         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, base=base)
-
-                # 创建智能体时传入Memory配置
-                agent = agent_class(
-                    name=agent_name,
-                    memory=memory,
-                )
-
-                self._agents[agent._uuid] = agent  # type:ignore
-                self._agent_uuids.append(agent._uuid)  # type:ignore
-
-            # 计算需要的组数,向上取整以处理不足一组的情况
-            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 + (k + 1) * group_size,  # 修正了索引计算
-                    class_init_index + agent_count_i,
-                )
-
-                agents = list(self._agents.values())[start_idx:end_idx]
-                group_name = f"AgentType_{i}_Group_{k}"
-
-                # 收集创建参数
-                group_creation_params.append((group_name, agents))
-
-            class_init_index += agent_count_i
+            
+            if self.agent_config_file is not None:
+                config_file = self.agent_config_file.get(agent_class, None) 
+            else:
+                config_file = None
+                
+            if issubclass(agent_class, InstitutionAgent):
+                institution_params.append((agent_class, agent_count_i, memory_config_func_i, config_file))
+            else:
+                citizen_params.append((agent_class, agent_count_i, memory_config_func_i, config_file))
+
+        # 处理机构智能体组
+        if institution_params:
+            total_institution_count = sum(p[1] for p in institution_params)
+            num_institution_groups = (total_institution_count + group_size - 1) // group_size
+            
+            for k in range(num_institution_groups):
+                start_idx = k * group_size
+                remaining = total_institution_count - start_idx
+                number_of_agents = min(remaining, group_size)
+                
+                agent_classes = []
+                agent_counts = []
+                memory_config_funcs = []
+                config_files = []
+                
+                # 分配每种类型的机构智能体到当前组
+                curr_start = start_idx
+                for agent_class, count, mem_func, conf_file in institution_params:
+                    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)
+                    curr_start = max(0, curr_start - count)
+                
+                group_creation_params.append((
+                    agent_classes,
+                    agent_counts, 
+                    memory_config_funcs,
+                    f"InstitutionGroup_{k}",
+                    config_files
+                ))
+
+        # 处理普通智能体组
+        if citizen_params:
+            total_citizen_count = sum(p[1] for p in citizen_params)
+            num_citizen_groups = (total_citizen_count + group_size - 1) // group_size
+            
+            for k in range(num_citizen_groups):
+                start_idx = k * group_size
+                remaining = total_citizen_count - start_idx
+                number_of_agents = min(remaining, group_size)
+                
+                agent_classes = []
+                agent_counts = []
+                memory_config_funcs = []
+                config_files = []
+                
+                # 分配每种类型的普通智能体到当前组
+                curr_start = start_idx
+                for agent_class, count, mem_func, conf_file in citizen_params:
+                    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)
+                    curr_start = max(0, curr_start - count)
+                
+                group_creation_params.append((
+                    agent_classes,
+                    agent_counts,
+                    memory_config_funcs, 
+                    f"CitizenGroup_{k}",
+                    config_files
+                ))
 
         # 初始化mlflow连接
         _mlflow_config = self.config.get("metric_request", {}).get("mlflow")
@@ -303,12 +461,14 @@ class AgentSimulation:
         else:
             _num_workers = 1
             self._pgsql_writers = _workers = [None for _ in range(_num_workers)]
-        # 收集所有创建组的参数
+
         creation_tasks = []
-        for i, (group_name, agents) in enumerate(group_creation_params):
+        for i, (agent_class, number_of_agents, memory_config_function_group, group_name, config_file) in enumerate(group_creation_params):
             # 直接创建异步任务
             group = AgentGroup.remote(
-                agents,
+                agent_class,
+                number_of_agents,
+                memory_config_function_group,
                 self.config,
                 self.exp_id,
                 self.exp_name,
@@ -319,27 +479,31 @@ class AgentSimulation:
                 mlflow_run_id,  # type:ignore
                 embedding_model,
                 self.logging_level,
+                config_file,
             )
-            creation_tasks.append((group_name, group, agents))
+            creation_tasks.append((group_name, group))
 
         # 更新数据结构
-        for group_name, group, agents in creation_tasks:
+        for group_name, group in creation_tasks:
             self._groups[group_name] = group
-            for agent in agents:
-                self._agent_uuid2group[agent._uuid] = group
+            group_agent_uuids = ray.get(group.get_agent_uuids.remote())
+            for agent_uuid in group_agent_uuids:
+                self._agent_uuid2group[agent_uuid] = group
+                self._user_chat_topics[agent_uuid] = f"exps/{self.exp_id}/agents/{agent_uuid}/user-chat"
+                self._user_survey_topics[agent_uuid] = (
+                    f"exps/{self.exp_id}/agents/{agent_uuid}/user-survey"
+                )
+            group_agent_type = ray.get(group.get_agent_type.remote())
+            for agent_type in group_agent_type:
+                if agent_type not in self._type2group:
+                    self._type2group[agent_type] = []
+                self._type2group[agent_type].append(group)
 
         # 并行初始化所有组的agents
         init_tasks = []
         for group in self._groups.values():
             init_tasks.append(group.init_agents.remote())
-        await asyncio.gather(*init_tasks)
-
-        # 设置用户主题
-        for uuid, agent in self._agents.items():
-            self._user_chat_topics[uuid] = f"exps/{self.exp_id}/agents/{uuid}/user-chat"
-            self._user_survey_topics[uuid] = (
-                f"exps/{self.exp_id}/agents/{uuid}/user-survey"
-            )
+        ray.get(init_tasks)
 
     async def gather(self, content: str):
         """收集智能体的特定信息"""
@@ -347,145 +511,42 @@ class AgentSimulation:
         for group in self._groups.values():
             gather_tasks.append(group.gather.remote(content))
         return await asyncio.gather(*gather_tasks)
+    
+    async def filter(self, 
+                     types: Optional[list[Type[Agent]]] = None, 
+                     keys: Optional[list[str]] = None, 
+                     values: Optional[list[Any]] = None) -> list[str]:
+        """过滤出指定类型的智能体"""
+        if not types and not keys and not values:
+            return self._agent_uuids
+        group_to_filter = []
+        for t in types:
+            if t in self._type2group:
+                group_to_filter.extend(self._type2group[t])
+            else:
+                raise ValueError(f"type {t} not found in simulation")
+        filtered_uuids = []
+        if keys:
+            if len(keys) != len(values):
+                raise ValueError("the length of key and value does not match")
+            for group in group_to_filter:
+                filtered_uuids.extend(await group.filter.remote(types, keys, values))
+            return filtered_uuids
+        else:
+            for group in group_to_filter:
+                filtered_uuids.extend(await group.filter.remote(types))
+            return filtered_uuids
 
-    async def update(self, target_agent_uuid: uuid.UUID, target_key: str, content: Any):
+    async def update(self, target_agent_uuid: str, target_key: str, content: Any):
         """更新指定智能体的记忆"""
         group = self._agent_uuid2group[target_agent_uuid]
         await group.update.remote(target_agent_uuid, target_key, content)
 
-    def default_memory_config_institution(self):
-        """默认的Memory配置函数"""
-        EXTRA_ATTRIBUTES = {
-            "type": (
-                int,
-                random.choice(
-                    [
-                        economyv2.ORG_TYPE_BANK,
-                        economyv2.ORG_TYPE_GOVERNMENT,
-                        economyv2.ORG_TYPE_FIRM,
-                        economyv2.ORG_TYPE_NBS,
-                        economyv2.ORG_TYPE_UNSPECIFIED,
-                    ]
-                ),
-            ),
-            "nominal_gdp": (list, [], True),
-            "real_gdp": (list, [], True),
-            "unemployment": (list, [], True),
-            "wages": (list, [], True),
-            "prices": (list, [], True),
-            "inventory": (int, 0, True),
-            "price": (float, 0.0, True),
-            "interest_rate": (float, 0.0, True),
-            "bracket_cutoffs": (list, [], True),
-            "bracket_rates": (list, [], True),
-            "employees": (list, [], True),
-            "customers": (list, [], True),
-        }
-        return EXTRA_ATTRIBUTES, None, None
-
-    def default_memory_config_citizen(self):
-        """默认的Memory配置函数"""
-        EXTRA_ATTRIBUTES = {
-            # 需求信息
-            "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),
-            "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 = {
-            "name": "unknown",
-            "gender": random.choice(["male", "female"]),
-            "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",
-                ]
-            ),
-            "age": random.randint(18, 65),
-            "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"]
-            ),
-            "income": str(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"]
-            ),
-        }
-
-        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)}
-            },
-        }
-
-        return EXTRA_ATTRIBUTES, PROFILE, BASE
-
     async def send_survey(
-        self, survey: Survey, agent_uuids: Optional[list[uuid.UUID]] = None
+        self, survey: Survey, agent_uuids: Optional[list[str]] = None
     ):
         """发送问卷"""
+        await self.messager.connect()
         survey_dict = survey.to_dict()
         if agent_uuids is None:
             agent_uuids = self._agent_uuids
@@ -499,12 +560,13 @@ class AgentSimulation:
         }
         for uuid in agent_uuids:
             topic = self._user_survey_topics[uuid]
-            await self._messager.send_message.remote(topic, payload)
+            await self.messager.send_message(topic, payload)
 
     async def send_interview_message(
-        self, content: str, agent_uuids: Union[uuid.UUID, list[uuid.UUID]]
+        self, content: str, agent_uuids: Union[str, list[str]]
     ):
-        """发送面试消息"""
+        """发送采访消息"""
+        await self.messager.connect()
         _date_time = datetime.now(timezone.utc)
         payload = {
             "from": "none",
@@ -512,11 +574,11 @@ class AgentSimulation:
             "timestamp": int(_date_time.timestamp() * 1000),
             "_date_time": _date_time,
         }
-        if not isinstance(agent_uuids, Sequence):
+        if not isinstance(agent_uuids, list):
             agent_uuids = [agent_uuids]
         for uuid in agent_uuids:
             topic = self._user_chat_topics[uuid]
-            await self._messager.send_message.remote(topic, payload)
+            await self.messager.send_message(topic, payload)
 
     async def step(self):
         """运行一步, 即每个智能体执行一次forward"""
@@ -529,60 +591,6 @@ class AgentSimulation:
             logger.error(f"运行错误: {str(e)}")
             raise
 
-    async def _save_exp_info(self) -> None:
-        """异步保存实验信息到YAML文件"""
-        try:
-            if self.enable_avro:
-                with open(self._exp_info_file, "w") as f:
-                    yaml.dump(self._exp_info, f)
-        except Exception as e:
-            logger.error(f"Avro保存实验信息失败: {str(e)}")
-        try:
-            if self.enable_pgsql:
-                worker: ray.ObjectRef = self._pgsql_writers[0]  # type:ignore
-                pg_exp_info = {
-                    k: self._exp_info[k] for (k, _) in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
-                }
-                pg_exp_info["created_at"] = self._exp_created_time
-                pg_exp_info["updated_at"] = self._exp_updated_time
-                await worker.async_update_exp_info.remote(  # type:ignore
-                    pg_exp_info
-                )
-        except Exception as e:
-            logger.error(f"PostgreSQL保存实验信息失败: {str(e)}")
-
-    async def _update_exp_status(self, status: int, error: str = "") -> None:
-        self._exp_updated_time = datetime.now(timezone.utc)
-        """更新实验状态并保存"""
-        self._exp_info["status"] = status
-        self._exp_info["error"] = error
-        self._exp_info["updated_at"] = self._exp_updated_time.isoformat()
-        await self._save_exp_info()
-
-    async def _monitor_exp_status(self, stop_event: asyncio.Event):
-        """监控实验状态并更新
-
-        Args:
-            stop_event: 用于通知监控任务停止的事件
-        """
-        try:
-            while not stop_event.is_set():
-                # 更新实验状态
-                # 假设所有group的cur_day和cur_t是同步的，取第一个即可
-                self._exp_info["cur_day"] = await self._simulator.get_simulator_day()
-                self._exp_info["cur_t"] = (
-                    await self._simulator.get_simulator_second_from_start_of_day()
-                )
-                await self._save_exp_info()
-
-                await asyncio.sleep(1)  # 避免过于频繁的更新
-        except asyncio.CancelledError:
-            # 正常取消，不需要特殊处理
-            pass
-        except Exception as e:
-            logger.error(f"监控实验状态时发生错误: {str(e)}")
-            raise
-
     async def run(
         self,
         day: int = 1,
@@ -619,16 +627,3 @@ class AgentSimulation:
             logger.error(error_msg)
             await self._update_exp_status(3, error_msg)
             raise RuntimeError(error_msg) from e
-
-    async def __aenter__(self):
-        """异步上下文管理器入口"""
-        return self
-
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        """异步上下文管理器出口"""
-        if exc_type is not None:
-            # 如果发生异常，更新状态为错误
-            await self._update_exp_status(3, str(exc_val))
-        elif self._exp_info["status"] != 3:
-            # 如果没有发生异常且状态不是错误，则更新为完成
-            await self._update_exp_status(2)