pycityagent 1.0.0__py3-none-any.whl → 2.0.0a2__py3-none-any.whl

pycityagent/simulation/simulation.py

@@ -0,0 +1,352 @@
+import asyncio
+import json
+import logging
+from datetime import datetime
+import random
+from typing import Dict, List, Optional, Callable
+from mosstool.map._map_util.const import AOI_START_ID
+
+from pycityagent.llm.llm import LLM
+from pycityagent.memory.memory import Memory
+
+from ..agent import Agent
+from ..environment import Simulator
+from .interview import InterviewManager
+from .survey import QuestionType, SurveyManager
+from .ui import InterviewUI
+
+logger = logging.getLogger(__name__)
+
+
+class AgentSimulation:
+    """城市智能体模拟器"""
+    def __init__(self, agent_class: type[Agent], simulator: Simulator, llm: LLM, agent_prefix: str = "agent_"):
+        """
+        Args:
+            agent_class: 智能体类
+            simulator: 模拟器
+            llm: 语言模型
+            agent_prefix: 智能体名称前缀
+        """
+        self.agent_class = agent_class
+        self.simulator = simulator
+        self.llm = llm
+        self.agent_prefix = agent_prefix
+        self._agents: Dict[str, Agent] = {}
+        self._interview_manager = InterviewManager()
+        self._interview_lock = asyncio.Lock()
+        self._start_time = datetime.now()
+        self._agent_run_times: Dict[str, datetime] = {}  # 记录每个智能体的运行开始时间
+        self._ui: Optional[InterviewUI] = None
+        self._loop = asyncio.get_event_loop()
+        self._blocked_agents: List[str] = []  # 新增：持续阻塞的智能体列表
+        self._survey_manager = SurveyManager()
+
+    def init_agents(self, agent_count: int, memory_config_func: Callable = None) -> None:
+        """初始化智能体
+        
+        Args:
+            agent_count: 要创建的智能体数量
+            memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组
+        """
+        if memory_config_func is None:
+            memory_config_func = self.default_memory_config_func
+        
+        for i in range(agent_count):
+            agent_name = f"{self.agent_prefix}{i}"
+
+            # 获取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,
+                simulator=self.simulator,
+                llm=self.llm,
+                memory=memory
+            )
+            
+            self.add_agent(agent)
+
+    def default_memory_config_func(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),
+        }
+
+        PROFILE = {
+            "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": random.randint(1000, 10000),
+            "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"]),
+            }
+
+        aois = self.simulator.aois.keys()
+        BASE = {
+            "home": {"aoi_position": {"aoi_id": random.choice(aois)}},
+            "work": {"aoi_position": {"aoi_id": random.choice(aois)}},
+        }
+
+        return EXTRA_ATTRIBUTES, PROFILE, BASE
+
+    def get_agent_runtime(self, agent_name: str) -> str:
+        """获取智能体运行时间"""
+        if agent_name not in self._agent_run_times:
+            return "-"
+        delta = datetime.now() - self._agent_run_times[agent_name]
+        hours = delta.seconds // 3600
+        minutes = (delta.seconds % 3600) // 60
+        seconds = delta.seconds % 60
+        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
+
+    def get_total_runtime(self) -> str:
+        """获取总运行时间"""
+        delta = datetime.now() - self._start_time
+        hours = delta.seconds // 3600
+        minutes = (delta.seconds % 3600) // 60
+        seconds = delta.seconds % 60
+        return f"{hours:02d}:{minutes:02d}:{seconds:02d}"
+
+    def export_chat_history(self, agent_name: str | None) -> str:
+        """导出对话历史
+
+        Args:
+            agent_name: 可选的智能体名称，如果提供则只导出该智能体的对话
+
+        Returns:
+            str: JSON格式的对话历史
+        """
+        history = (
+            self._interview_manager.get_agent_history(agent_name)
+            if agent_name
+            else self._interview_manager.get_recent_history(limit=1000)
+        )
+
+        # 转换为易读格式
+        formatted_history = []
+        for record in history:
+            formatted_history.append(
+                {
+                    "timestamp": record.timestamp.strftime("%Y-%m-%d %H:%M:%S"),
+                    "agent": record.agent_name,
+                    "question": record.question,
+                    "response": record.response,
+                    "blocking": record.blocking,
+                }
+            )
+
+        return json.dumps(formatted_history, ensure_ascii=False, indent=2)
+
+    def toggle_agent_block(self, agent_name: str, blocking: bool) -> str:
+        """切换智能体的阻塞状态
+
+        Args:
+            agent_name: 能体名称
+            blocking: True表示阻塞，False表示取消阻塞
+
+        Returns:
+            str: 状态变更消息
+        """
+        if agent_name not in self._agents:
+            return f"找不到智能体 {agent_name}"
+
+        if blocking and agent_name not in self._blocked_agents:
+            self._blocked_agents.append(agent_name)
+            return f"已阻塞智能体 {agent_name}"
+        elif not blocking and agent_name in self._blocked_agents:
+            self._blocked_agents.remove(agent_name)
+            return f"已取消阻塞智能体 {agent_name}"
+
+        return f"智能体 {agent_name} 状态未变"
+
+    async def interview_agent(self, agent_name: str, question: str) -> str:
+        """采访指定智能体"""
+        agent = self._agents.get(agent_name)
+        if not agent:
+            return "找不到指定的智能体"
+
+        try:
+            response = await agent.generate_response(question)
+            # 记录采访历史
+            self._interview_manager.add_record(
+                agent_name,
+                question,
+                response,
+                blocking=(agent_name in self._blocked_agents),
+            )
+            return response
+
+        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:
+        """向智能体提交问卷
+
+        Args:
+            agent_name: 智能体名称
+            survey_id: 问卷ID
+
+        Returns:
+            str: 处理结果
+        """
+        agent = self._agents.get(agent_name)
+        if not agent:
+            return "找不到指定的智能体"
+
+        survey = self._survey_manager.get_survey(survey_id)
+        if not survey:
+            return "找不到指定的问卷"
+
+        try:
+            # 建问卷提示
+            prompt = f"""请以第一人称回答以下调查问卷:
+
+问卷标题: {survey.title}
+问卷说明: {survey.description}
+
+"""
+            for i, question in enumerate(survey.questions):
+                prompt += f"\n问题{i+1}. {question.content}"
+                if question.type in (
+                    QuestionType.SINGLE_CHOICE,
+                    QuestionType.MULTIPLE_CHOICE,
+                ):
+                    prompt += "\n选项: " + ", ".join(question.options)
+                elif question.type == QuestionType.RATING:
+                    prompt += (
+                        f"\n(请给出{question.min_rating}-{question.max_rating}的评分)"
+                    )
+                elif question.type == QuestionType.LIKERT:
+                    prompt += "\n(1-强烈不同意, 2-不同意, 3-中立, 4-同意, 5-强烈同意)"
+
+            # 生成回答
+            response = await agent.generate_response(prompt)
+
+            # 存储原始回答
+            self._survey_manager.add_response(
+                survey_id, agent_name, {"raw_response": response, "parsed": False}
+            )
+
+            return response
+
+        except Exception as e:
+            logger.error(f"问卷处理出错: {str(e)}")
+            return f"问卷处理出现错误: {str(e)}"
+
+    def create_survey(self, **survey_data: dict) -> None:
+        """创建新问卷
+
+        Args:
+            survey_data: 问卷数据，包含 title, description, questions
+
+        Returns:
+            更新后的问卷列表
+        """
+        self._survey_manager.create_survey(**survey_data)  # type:ignore
+
+    def get_surveys(self) -> list:
+        """获取所有问卷"""
+        return self._survey_manager.get_all_surveys()
+
+    def get_survey_questions(self, survey_id: str) -> dict | None:
+        """获取指定问卷的问题列表
+
+        Args:
+            survey_id: 问卷ID
+
+        Returns:
+            问卷数据，包含 title, description, questions
+        """
+        for _, survey in self._survey_manager._surveys.items():
+            survey_dict = survey.to_dict()
+            if survey_dict["id"] == survey_id:
+                return survey_dict
+        return None
+
+    async def init_ui(
+        self,
+        server_name: str = "127.0.0.1",
+        server_port: int = 7860,
+    ):
+        """初始化UI"""
+        self._interview_lock = asyncio.Lock()
+        # 初始化GradioUI
+        self._ui = InterviewUI(self)
+        interface = self._ui.create_interface()
+        interface.queue().launch(
+            server_name=server_name,
+            server_port=server_port,
+            prevent_thread_lock=True,
+            quiet=True,
+        )
+        print(
+            f"Gradio Frontend is running on http://{server_name}:{server_port}"
+        )
+
+    async def step(self):
+        """运行一步, 即每个智能体执行一次forward"""
+        try:
+            tasks = []
+            for agent in self._agents.values():
+                tasks.append(agent.forward())
+            await asyncio.gather(*tasks)
+        except Exception as e:
+            logger.error(f"运行错误: {str(e)}")
+            raise
+        
+    async def run(
+        self,
+        day: int = 1,
+    ):
+        """运行模拟器
+
+        Args:
+            day: 运行天数,默认为1天
+        """
+        try:
+            # 获取开始时间
+            start_time = self.simulator.GetTime()
+            # 计算结束时间（秒）
+            end_time = start_time + day * 24 * 3600  # 将天数转换为秒
+            
+            while True:
+                current_time = self.simulator.GetTime()
+                if current_time >= end_time:
+                    break
+                
+                tasks = []
+                for agent in self._agents.values():
+                    if agent.name not in self._blocked_agents:
+                        tasks.append(agent.forward())
+                
+                await asyncio.gather(*tasks)
+
+        except Exception as e:
+            logger.error(f"模拟器运行错误: {str(e)}")
+            raise

pycityagent/simulation/survey/__init__.py

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

pycityagent/simulation/survey/manager.py

@@ -0,0 +1,67 @@
+from typing import List, Dict, Optional
+from datetime import datetime
+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:
+        """创建新问卷"""
+        survey_id = str(uuid.uuid4())
+        
+        # 转换问题数据
+        survey_questions = []
+        for q in questions:
+            question = Question(
+                content=q["content"],
+                type=QuestionType(q["type"]),
+                required=q.get("required", True),
+                options=q.get("options", []),
+                min_rating=q.get("min_rating", 1),
+                max_rating=q.get("max_rating", 5)
+            )
+            survey_questions.append(question)
+            
+        survey = Survey(
+            id=survey_id,
+            title=title,
+            description=description,
+            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
+        }
+        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) 

pycityagent/simulation/survey/models.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass, field
+from typing import List, Dict, Optional
+from datetime import datetime
+from enum import Enum
+import uuid
+
+class QuestionType(Enum):
+    TEXT = "文本"
+    SINGLE_CHOICE = "单选"
+    MULTIPLE_CHOICE = "多选"
+    RATING = "评分"
+    LIKERT = "李克特量表"
+
+@dataclass
+class Question:
+    content: str
+    type: QuestionType
+    required: bool = True
+    options: List[str] = field(default_factory=list)
+    min_rating: int = 1
+    max_rating: int = 5
+    
+    def to_dict(self) -> dict:
+        return {
+            "content": self.content,
+            "type": self.type.value,
+            "required": self.required,
+            "options": self.options,
+            "min_rating": self.min_rating,
+            "max_rating": self.max_rating
+        }
+
+@dataclass
+class Survey:
+    id: str
+    title: str
+    description: str
+    questions: List[Question]
+    responses: Dict[str, dict] = field(default_factory=dict)
+    created_at: datetime = field(default_factory=datetime.now)
+    
+    def to_dict(self) -> dict:
+        return {
+            "id": self.id,
+            "title": self.title,
+            "description": self.description,
+            "questions": [q.to_dict() for q in self.questions],
+            "response_count": len(self.responses)
+        } 

pycityagent/simulation/ui/__init__.py

@@ -0,0 +1,3 @@
+from .interface import InterviewUI
+
+__all__ = ["InterviewUI"]