pycityagent 2.0.0a13__py3-none-any.whl → 2.0.0a15__py3-none-any.whl

pycityagent/simulation/simulation.py

@@ -1,21 +1,25 @@
 import asyncio
 import json
 import logging
+import os
+from pathlib import Path
 import uuid
-from datetime import datetime
+from datetime import datetime, timezone
 import random
 from typing import Dict, List, Optional, Callable, Union,Any
 from mosstool.map._map_util.const import AOI_START_ID
 import pycityproto.city.economy.v2.economy_pb2 as economyv2
+from pycityagent.environment.simulator import Simulator
 from pycityagent.memory.memory import Memory
+from pycityagent.message.messager import Messager
+from pycityagent.survey import Survey
+import yaml
+from concurrent.futures import ThreadPoolExecutor
 
 from ..agent import Agent, InstitutionAgent
-from .interview import InterviewManager
-from .survey import QuestionType, SurveyManager
-from .ui import InterviewUI
 from .agentgroup import AgentGroup
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("pycityagent")
 
 
 class AgentSimulation:
@@ -26,28 +30,72 @@ class AgentSimulation:
         agent_class: Union[type[Agent], list[type[Agent]]],
         config: dict,
         agent_prefix: str = "agent_",
+        exp_name: str = "default_experiment",
+        logging_level: int = logging.WARNING
     ):
         """
         Args:
             agent_class: 智能体类
             config: 配置
             agent_prefix: 智能体名称前缀
+            exp_name: 实验名称
         """
-        self.exp_id = uuid.uuid4()
+        self.exp_id = str(uuid.uuid4())
         if isinstance(agent_class, list):
             self.agent_class = agent_class
         else:
             self.agent_class = [agent_class]
+        self.logging_level = logging_level
         self.config = config
+        self._simulator = Simulator(config["simulator_request"])
         self.agent_prefix = agent_prefix
         self._agents: Dict[uuid.UUID, Agent] = {}
         self._groups: Dict[str, AgentGroup] = {}
         self._agent_uuid2group: Dict[uuid.UUID, AgentGroup] = {}
         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._loop = asyncio.get_event_loop()
-        self._interview_manager = InterviewManager()
-        self._survey_manager = SurveyManager()
+
+        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),
+        )
+        asyncio.create_task(self._messager.connect())
+
+        self._enable_avro = config["storage"]["avro"]["enabled"]
+        if not self._enable_avro:
+            logger.warning("AVRO is not enabled, NO AVRO LOCAL STORAGE")
+        self._avro_path = Path(config["storage"]["avro"]["path"]) / f"{self.exp_id}"
+        self._avro_path.mkdir(parents=True, exist_ok=True)
+
+        self._enable_pgsql = config["storage"]["pgsql"]["enabled"]
+        self._pgsql_host = config["storage"]["pgsql"]["host"]
+        self._pgsql_port = config["storage"]["pgsql"]["port"]
+        self._pgsql_database = config["storage"]["pgsql"]["database"]
+        self._pgsql_user = config["storage"]["pgsql"]["user"]
+        self._pgsql_password = config["storage"]["pgsql"]["password"]
+
+        # 添加实验信息相关的属性
+        self._exp_info = {
+            "id": self.exp_id,
+            "name": exp_name,
+            "num_day": 0,  # 将在 run 方法中更新
+            "status": 0,
+            "cur_day": 0,
+            "cur_t": 0.0,
+            "config": json.dumps(config),
+            "error": "",
+            "created_at": datetime.now(timezone.utc).isoformat()
+        }
+        
+        # 创建异步任务保存实验信息
+        self._exp_info_file = self._avro_path / "experiment_info.yaml"
+        with open(self._exp_info_file, 'w') as f:
+            yaml.dump(self._exp_info, f)
 
     @property
     def agents(self):
@@ -64,6 +112,11 @@ class AgentSimulation:
     @property
     def agent_uuid2group(self):
         return self._agent_uuid2group
+    
+    def create_remote_group(self, group_name: str, agents: list[Agent], config: dict, exp_id: str, enable_avro: bool, avro_path: Path, logging_level: int = logging.WARNING):
+        """创建远程组"""
+        group = AgentGroup.remote(agents, config, exp_id, enable_avro, avro_path, logging_level)
+        return group_name, group, agents
 
     async def init_agents(
         self,
@@ -85,7 +138,7 @@ class AgentSimulation:
             raise ValueError("agent_class和agent_count的长度不一致")
 
         if memory_config_func is None:
-            logging.warning(
+            logger.warning(
                 "memory_config_func is None, using default memory config function"
             )
             memory_config_func = []
@@ -98,17 +151,21 @@ class AgentSimulation:
             memory_config_func = [memory_config_func]
 
         if len(memory_config_func) != len(agent_count):
-            logging.warning(
+            logger.warning(
                 "memory_config_func和agent_count的长度不一致，使用默认的memory_config"
             )
             memory_config_func = []
             for agent_class in self.agent_class:
-                if agent_class == InstitutionAgent:
+                if issubclass(agent_class, InstitutionAgent):
                     memory_config_func.append(self.default_memory_config_institution)
                 else:
                     memory_config_func.append(self.default_memory_config_citizen)
 
+        # 使用线程池并行创建 AgentGroup
+        group_creation_params = []
         class_init_index = 0
+        
+        # 首先收集所有需要创建的组的参数
         for i in range(len(self.agent_class)):
             agent_class = self.agent_class[i]
             agent_count_i = agent_count[i]
@@ -133,30 +190,51 @@ class AgentSimulation:
 
             # 计算需要的组数,向上取整以处理不足一组的情况
             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,
+                    class_init_index + (k + 1) * group_size,  # 修正了索引计算
+                    class_init_index + agent_count_i
                 )
-
-                # 获取当前组的agents
+                
                 agents = list(self._agents.values())[start_idx:end_idx]
                 group_name = f"AgentType_{i}_Group_{k}"
-                group = AgentGroup.remote(agents, self.config, self.exp_id)
-                self._groups[group_name] = group
-                for agent in agents:
-                    self._agent_uuid2group[agent._uuid] = group
-
-            class_init_index += agent_count_i  # 更新类初始索引
-
+                
+                # 收集创建参数
+                group_creation_params.append((
+                    group_name, 
+                    agents
+                ))
+            
+            class_init_index += agent_count_i
+
+        # 收集所有创建组的参数
+        creation_tasks = []
+        for group_name, agents in group_creation_params:
+            # 直接创建异步任务
+            group = AgentGroup.remote(agents, self.config, self.exp_id, 
+                                    self._enable_avro, self._avro_path, 
+                                    self.logging_level)
+            creation_tasks.append((group_name, group, agents))
+        
+        # 更新数据结构
+        for group_name, group, agents in creation_tasks:
+            self._groups[group_name] = group
+            for agent in agents:
+                self._agent_uuid2group[agent._uuid] = 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"
+
     async def gather(self, content: str):
         """收集智能体的特定信息"""
         gather_tasks = []
@@ -164,10 +242,10 @@ class AgentSimulation:
             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):
+    async def update(self, target_agent_uuid: uuid.UUID, target_key: str, content: Any):
         """更新指定智能体的记忆"""
-        group = self._agent_uuid2group[target_agent_id]
-        await group.update.remote(target_agent_id, target_key, content)
+        group = self._agent_uuid2group[target_agent_uuid]
+        await group.update.remote(target_agent_uuid, target_key, content)
 
     def default_memory_config_institution(self):
         """默认的Memory配置函数"""
@@ -219,6 +297,7 @@ class AgentSimulation:
         }
 
         PROFILE = {
+            "name": "unknown",
             "gender": random.choice(["male", "female"]),
             "education": random.choice(
                 ["Doctor", "Master", "Bachelor", "College", "High School"]
@@ -251,7 +330,7 @@ class AgentSimulation:
             "personality": random.choice(
                 ["outgoint", "introvert", "ambivert", "extrovert"]
             ),
-            "income": random.randint(1000, 10000),
+            "income": str(random.randint(1000, 10000)),
             "currency": random.randint(10000, 100000),
             "residence": random.choice(["city", "suburb", "rural"]),
             "race": random.choice(
@@ -285,6 +364,32 @@ class AgentSimulation:
         }
 
         return EXTRA_ATTRIBUTES, PROFILE, BASE
+    
+    async def send_survey(self, survey: Survey, agent_uuids: Optional[List[uuid.UUID]] = None):
+        """发送问卷"""
+        survey = survey.to_dict()
+        if agent_uuids is None:
+            agent_uuids = self._agent_uuids
+        payload = {
+            "from": "none",
+            "survey_id": survey["id"],
+            "timestamp": int(datetime.now().timestamp() * 1000),
+            "data": survey,
+        }
+        for uuid in agent_uuids:
+            topic = self._user_survey_topics[uuid]
+            await self._messager.send_message(topic, payload)
+
+    async def send_interview_message(self, content: str, agent_uuids: Union[uuid.UUID, List[uuid.UUID]]):
+        """发送面试消息"""
+        payload = {
+            "from": "none",
+            "content": content,
+            "timestamp": int(datetime.now().timestamp() * 1000),
+        }
+        for uuid in agent_uuids:
+            topic = self._user_chat_topics[uuid]
+            await self._messager.send_message(topic, payload)
 
     async def step(self):
         """运行一步, 即每个智能体执行一次forward"""
@@ -297,23 +402,88 @@ class AgentSimulation:
             logger.error(f"运行错误: {str(e)}")
             raise
 
-    async def run(
-        self,
-        day: int = 1,
-    ):
-        """运行模拟器
+    async def _save_exp_info(self) -> None:
+        """异步保存实验信息到YAML文件"""
+        try:
+            with open(self._exp_info_file, 'w') as f:
+                yaml.dump(self._exp_info, f)
+        except Exception as e:
+            logger.error(f"保存实验信息失败: {str(e)}")
+
+    async def _update_exp_status(self, status: int, error: str = "") -> None:
+        """更新实验状态并保存"""
+        self._exp_info["status"] = status
+        self._exp_info["error"] = error
+        await self._save_exp_info()
 
+    async def _monitor_exp_status(self, stop_event: asyncio.Event):
+        """监控实验状态并更新
+        
         Args:
-            day: 运行天数,默认为1天
+            stop_event: 用于通知监控任务停止的事件
         """
         try:
-            # 获取开始时间
-            tasks = []
-            for group in self._groups.values():
-                tasks.append(group.run.remote(day))
+            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.CancelError:
+            # 正常取消，不需要特殊处理
+            pass
+        except Exception as e:
+            logger.error(f"监控实验状态时发生错误: {str(e)}")
+            raise
 
-            await asyncio.gather(*tasks)
+    async def run(
+        self,
+        day: int = 1,
+    ):
+        """运行模拟器"""
+        try:
+            self._exp_info["num_day"] += day
+            await self._update_exp_status(1)  # 更新状态为运行中
+            
+            # 创建停止事件
+            stop_event = asyncio.Event()
+            # 创建监控任务
+            monitor_task = asyncio.create_task(self._monitor_exp_status(stop_event))
+            
+            try:
+                tasks = []
+                for group in self._groups.values():
+                    tasks.append(group.run.remote())
+
+                # 等待所有group运行完成
+                await asyncio.gather(*tasks)
+                
+            finally:
+                # 设置停止事件
+                stop_event.set()
+                # 等待监控任务结束
+                await monitor_task
+            
+            # 运行成功后更新状态
+            await self._update_exp_status(2)
 
         except Exception as e:
-            logger.error(f"模拟器运行错误: {str(e)}")
-            raise
+            error_msg = f"模拟器运行错误: {str(e)}"
+            logger.error(error_msg)
+            await self._update_exp_status(3, error_msg)
+            raise 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)

pycityagent/survey/manager.py

@@ -0,0 +1,58 @@
+from typing import List, Dict, Optional
+from datetime import datetime
+import uuid
+import json
+from .models import Survey, Question, QuestionType, Page
+
+
+class SurveyManager:
+    def __init__(self):
+        self._surveys: Dict[str, Survey] = {}
+
+    def create_survey(
+        self, title: str, description: str, pages: List[dict]
+    ) -> Survey:
+        """创建新问卷"""
+        survey_id = uuid.uuid4()
+
+        # 转换页面和问题数据
+        survey_pages = []
+        for page_data in pages:
+            questions = []
+            for q in page_data["elements"]:
+                question = Question(
+                    name=q["name"],
+                    title=q["title"],
+                    type=QuestionType(q["type"]),
+                    required=q.get("required", True),
+                    choices=q.get("choices", []),
+                    columns=q.get("columns", []),
+                    rows=q.get("rows", []),
+                    min_rating=q.get("min_rating", 1),
+                    max_rating=q.get("max_rating", 5),
+                )
+                questions.append(question)
+            
+            page = Page(
+                name=page_data["name"],
+                elements=questions
+            )
+            survey_pages.append(page)
+
+        survey = Survey(
+            id=survey_id,
+            title=title,
+            description=description,
+            pages=survey_pages,
+        )
+
+        self._surveys[str(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())

pycityagent/survey/models.py

@@ -0,0 +1,120 @@
+from dataclasses import dataclass, field
+from typing import List, Dict, Optional
+from datetime import datetime
+from enum import Enum
+import uuid
+import json
+
+
+class QuestionType(Enum):
+    TEXT = "text"
+    RADIO = "radiogroup"
+    CHECKBOX = "checkbox"
+    BOOLEAN = "boolean"
+    RATING = "rating"
+    MATRIX = "matrix"
+
+
+@dataclass
+class Question:
+    name: str
+    title: str
+    type: QuestionType
+    choices: List[str] = field(default_factory=list)
+    columns: List[str] = field(default_factory=list)
+    rows: List[str] = field(default_factory=list)
+    min_rating: int = 1
+    max_rating: int = 5
+
+    def to_dict(self) -> dict:
+        base_dict = {
+            "type": self.type.value,
+            "name": self.name,
+            "title": self.title,
+        }
+        
+        if self.type in [QuestionType.RADIO, QuestionType.CHECKBOX]:
+            base_dict["choices"] = self.choices
+        elif self.type == QuestionType.MATRIX:
+            base_dict["columns"] = self.columns
+            base_dict["rows"] = self.rows
+        elif self.type == QuestionType.RATING:
+            base_dict["min_rating"] = self.min_rating
+            base_dict["max_rating"] = self.max_rating
+            
+        return base_dict
+
+
+@dataclass
+class Page:
+    name: str
+    elements: List[Question]
+
+    def to_dict(self) -> dict:
+        return {
+            "name": self.name,
+            "elements": [q.to_dict() for q in self.elements]
+        }
+
+
+@dataclass
+class Survey:
+    id: uuid.UUID
+    title: str
+    description: str
+    pages: List[Page]
+    responses: Dict[str, dict] = field(default_factory=dict)
+    created_at: datetime = field(default_factory=datetime.now)
+
+    def to_dict(self) -> dict:
+        return {
+            "id": str(self.id),
+            "title": self.title,
+            "description": self.description,
+            "pages": [p.to_dict() for p in self.pages],
+            "response_count": len(self.responses),
+        }
+
+    def to_json(self) -> str:
+        """Convert the survey to a JSON string for MQTT transmission"""
+        survey_dict = {
+            "id": str(self.id),
+            "title": self.title,
+            "description": self.description,
+            "pages": [p.to_dict() for p in self.pages],
+            "responses": self.responses,
+            "created_at": self.created_at.isoformat()
+        }
+        return json.dumps(survey_dict)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> 'Survey':
+        """Create a Survey instance from a JSON string"""
+        data = json.loads(json_str)
+        pages = [
+            Page(
+                name=p["name"],
+                elements=[
+                    Question(
+                        name=q["name"],
+                        title=q["title"],
+                        type=QuestionType(q["type"]),
+                        required=q.get("required", True),
+                        choices=q.get("choices", []),
+                        columns=q.get("columns", []),
+                        rows=q.get("rows", []),
+                        min_rating=q.get("min_rating", 1),
+                        max_rating=q.get("max_rating", 5)
+                    ) for q in p["elements"]
+                ]
+            ) for p in data["pages"]
+        ]
+        
+        return cls(
+            id=uuid.UUID(data["id"]),
+            title=data["title"],
+            description=data["description"],
+            pages=pages,
+            responses=data.get("responses", {}),
+            created_at=datetime.fromisoformat(data["created_at"])
+        )

pycityagent/utils/__init__.py

@@ -0,0 +1,7 @@
+from .avro_schema import PROFILE_SCHEMA, DIALOG_SCHEMA, STATUS_SCHEMA, SURVEY_SCHEMA, INSTITUTION_STATUS_SCHEMA
+from .survey_util import process_survey_for_llm
+
+__all__ = [
+    "PROFILE_SCHEMA", "DIALOG_SCHEMA", "STATUS_SCHEMA", "SURVEY_SCHEMA", "INSTITUTION_STATUS_SCHEMA",
+    "process_survey_for_llm"
+]

pycityagent/utils/avro_schema.py

@@ -0,0 +1,110 @@
+PROFILE_SCHEMA = {
+    "doc": "Agent属性",
+    "name": "AgentProfile",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "name", "type": "string"},
+        {"name": "gender", "type": "string"},
+        {"name": "age", "type": "float"},
+        {"name": "education", "type": "string"},
+        {"name": "skill", "type": "string"},
+        {"name": "occupation", "type": "string"},
+        {"name": "family_consumption", "type": "string"},
+        {"name": "consumption", "type": "string"},
+        {"name": "personality", "type": "string"},
+        {"name": "income", "type": "string"},
+        {"name": "currency", "type": "float"},
+        {"name": "residence", "type": "string"},
+        {"name": "race", "type": "string"},
+        {"name": "religion", "type": "string"},
+        {"name": "marital_status", "type": "string"},
+    ],
+}
+
+DIALOG_SCHEMA = {
+    "doc": "Agent对话",
+    "name": "AgentDialog",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "day", "type": "int"},
+        {"name": "t", "type": "float"},
+        {"name": "type", "type": "int"},
+        {"name": "speaker", "type": "string"},
+        {"name": "content", "type": "string"},
+        {
+            "name": "created_at",
+            "type": {"type": "long", "logicalType": "timestamp-millis"},
+        },
+    ],
+}
+
+STATUS_SCHEMA = {
+    "doc": "Agent状态",
+    "name": "AgentStatus",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "day", "type": "int"},
+        {"name": "t", "type": "float"},
+        {"name": "lng", "type": "double"},
+        {"name": "lat", "type": "double"},
+        {"name": "parent_id", "type": "int"},
+        {"name": "action", "type": "string"},
+        {"name": "hungry", "type": "float"},
+        {"name": "tired", "type": "float"},
+        {"name": "safe", "type": "float"},
+        {"name": "social", "type": "float"},
+        {
+            "name": "created_at",
+            "type": {"type": "long", "logicalType": "timestamp-millis"},
+        },
+    ],
+}
+
+INSTITUTION_STATUS_SCHEMA = {
+    "doc": "Institution状态",
+    "name": "InstitutionStatus",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "day", "type": "int"},
+        {"name": "t", "type": "float"},
+        {"name": "type", "type": "int"},
+        {"name": "nominal_gdp", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "real_gdp", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "unemployment", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "wages", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "prices", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "inventory", "type": ["int", "null"]},
+        {"name": "price", "type": ["float", "null"]},
+        {"name": "interest_rate", "type": ["float", "null"]},
+        {"name": "bracket_cutoffs", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "bracket_rates", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "employees", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "customers", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+    ],
+}
+
+SURVEY_SCHEMA = {
+    "doc": "Agent问卷",
+    "name": "AgentSurvey",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "day", "type": "int"},
+        {"name": "t", "type": "float"},
+        {"name": "survey_id", "type": "string"},
+        {"name": "result", "type": "string"},
+        {
+            "name": "created_at",
+            "type": {"type": "long", "logicalType": "timestamp-millis"},
+        },
+    ],
+}