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

pycityagent/agent.py

@@ -13,11 +13,15 @@ import random
 import uuid
 from typing import Dict, List, Optional
 
+import fastavro
+
 from pycityagent.environment.sim.person_service import PersonService
 from mosstool.util.format_converter import dict2pb
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
+from pycityagent.utils import process_survey_for_llm
 
 from pycityagent.message.messager import Messager
+from pycityagent.utils import SURVEY_SCHEMA, DIALOG_SCHEMA
 
 from .economy import EconomyClient
 from .environment import Simulator
@@ -52,6 +56,7 @@ class Agent(ABC):
         messager: Optional[Messager] = None,
         simulator: Optional[Simulator] = None,
         memory: Optional[Memory] = None,
+        avro_file: Optional[Dict[str, str]] = None,
     ) -> None:
         """
         Initialize the Agent.
@@ -64,6 +69,7 @@ class Agent(ABC):
             messager (Messager, optional): The messager object. Defaults to None.
             simulator (Simulator, optional): The simulator object. Defaults to None.
             memory (Memory, optional): The memory of the agent. Defaults to None.
+            avro_file (Dict[str, str], optional): The avro file of the agent. Defaults to None.
         """
         self._name = name
         self._type = type
@@ -79,6 +85,7 @@ class Agent(ABC):
         self._blocked = False
         self._interview_history: List[Dict] = []  # 存储采访历史
         self._person_template = PersonService.default_dict_person()
+        self._avro_file = avro_file
 
     def __getstate__(self):
         state = self.__dict__.copy()
@@ -168,11 +175,59 @@ class Agent(ABC):
             )
         return self._simulator
 
-    async def generate_response(self, question: str) -> str:
-        """生成回答
+    async def generate_user_survey_response(self, survey: dict) -> str:
+        """生成回答 —— 可重写
+        基于智能体的记忆和当前状态，生成对问卷调查的回答。
+        Args:
+            survey: 需要回答的问卷 dict
+        Returns:
+            str: 智能体的回答
+        """
+        survey_prompt = process_survey_for_llm(survey)
+        dialog = []
 
-        基于智能体的记忆和当前状态，生成对问题的回答。
+        # 添加系统提示
+        system_prompt = "Please answer the survey question in first person. Follow the format requirements strictly and provide clear and specific answers."
+        dialog.append({"role": "system", "content": system_prompt})
+
+        # 添加记忆上下文
+        if self._memory:
+            relevant_memories = await self._memory.search(survey_prompt)
+            if relevant_memories:
+                dialog.append(
+                    {
+                        "role": "system",
+                        "content": f"Answer based on these memories:\n{relevant_memories}",
+                    }
+                )
+
+        # 添加问卷问题
+        dialog.append({"role": "user", "content": survey_prompt})
+
+        # 使用LLM生成回答
+        if not self._llm_client:
+            return "Sorry, I cannot answer survey questions right now."
+
+        response = await self._llm_client.atext_request(dialog)  # type:ignore
 
+        return response  # type:ignore
+    
+    async def _process_survey(self, survey: dict):
+        survey_response = await self.generate_user_survey_response(survey)
+        response_to_avro = [{
+            "id": str(self._uuid),
+            "day": await self._simulator.get_simulator_day(),
+            "t": await self._simulator.get_simulator_second_from_start_of_day(),
+            "survey_id": survey["id"],
+            "result": survey_response,
+            "created_at": int(datetime.now().timestamp() * 1000),
+        }]
+        with open(self._avro_file["survey"], "a+b") as f:
+            fastavro.writer(f, SURVEY_SCHEMA, response_to_avro, codec="snappy")
+
+    async def generate_user_chat_response(self, question: str) -> str:
+        """生成回答 —— 可重写
+        基于智能体的记忆和当前状态，生成对问题的回答。
         Args:
             question: 需要回答的问题
 
@@ -182,7 +237,7 @@ class Agent(ABC):
         dialog = []
 
         # 添加系统提示
-        system_prompt = f"请以第一人称的方式回答问题,保持回答简洁明了。"
+        system_prompt = "Please answer the question in first person and keep the response concise and clear."
         dialog.append({"role": "system", "content": system_prompt})
 
         # 添加记忆上下文
@@ -192,7 +247,7 @@ class Agent(ABC):
                 dialog.append(
                     {
                         "role": "system",
-                        "content": f"基于以下记忆回答问题:\n{relevant_memories}",
+                        "content": f"Answer based on these memories:\n{relevant_memories}",
                     }
                 )
 
@@ -201,46 +256,76 @@ class Agent(ABC):
 
         # 使用LLM生成回答
         if not self._llm_client:
-            return "抱歉，我现在无法回答问题。"
+            return "Sorry, I cannot answer questions right now."
 
         response = await self._llm_client.atext_request(dialog)  # type:ignore
 
-        # 记录采访历史
-        self._interview_history.append(
-            {
-                "timestamp": datetime.now().isoformat(),
-                "question": question,
-                "response": response,
-            }
-        )
-
         return response  # type:ignore
-
-    def get_interview_history(self) -> List[Dict]:
-        """获取采访历史记录"""
-        return self._interview_history
-
+    
+    async def _process_interview(self, payload: dict):
+        auros = [{
+            "id": str(self._uuid),
+            "day": await self._simulator.get_simulator_day(),
+            "t": await self._simulator.get_simulator_second_from_start_of_day(),
+            "type": 2,
+            "speaker": "user",
+            "content": payload["content"],
+            "created_at": int(datetime.now().timestamp() * 1000),
+        }]
+        question = payload["content"]
+        response = await self.generate_user_chat_response(question)
+        auros.append({
+            "id": str(self._uuid),
+            "day": await self._simulator.get_simulator_day(),
+            "t": await self._simulator.get_simulator_second_from_start_of_day(),
+            "type": 2,
+            "speaker": "",
+            "content": response,
+            "created_at": int(datetime.now().timestamp() * 1000),
+        })
+        with open(self._avro_file["dialog"], "a+b") as f:
+            fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+
+    async def process_agent_chat_response(self, payload: dict) -> str:
+        logging.info(f"Agent {self._uuid} received agent chat response: {payload}")
+
+    async def _process_agent_chat(self, payload: dict):
+        auros = [{
+            "id": str(self._uuid),
+            "day": payload["day"],
+            "t": payload["t"],
+            "type": 1,
+            "speaker": payload["from"],
+            "content": payload["content"],
+            "created_at": int(datetime.now().timestamp() * 1000),
+        }]
+        asyncio.create_task(self.process_agent_chat_response(payload))
+        with open(self._avro_file["dialog"], "a+b") as f:
+            fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+
+    # Callback functions for MQTT message
     async def handle_agent_chat_message(self, payload: dict):
         """处理收到的消息，识别发送者"""
         # 从消息中解析发送者 ID 和消息内容
-        print(
-            f"Agent {self._uuid} received agent chat message: '{payload['content']}' from Agent {payload['from']}"
-        )
+        logging.info(f"Agent {self._uuid} received agent chat message: {payload}")
+        asyncio.create_task(self._process_agent_chat(payload))
 
     async def handle_user_chat_message(self, payload: dict):
         """处理收到的消息，识别发送者"""
         # 从消息中解析发送者 ID 和消息内容
-        print(
-            f"Agent {self._uuid} received user chat message: '{payload['content']}' from User"
-        )
+        logging.info(f"Agent {self._uuid} received user chat message: {payload}")
+        asyncio.create_task(self._process_interview(payload))
 
     async def handle_user_survey_message(self, payload: dict):
         """处理收到的消息，识别发送者"""
         # 从消息中解析发送者 ID 和消息内容
-        print(
-            f"Agent {self._uuid} received user survey message: '{payload['content']}' from User"
-        )
+        logging.info(f"Agent {self._uuid} received user survey message: {payload}")
+        asyncio.create_task(self._process_survey(payload["data"]))
 
+    async def handle_gather_message(self, payload: str):
+        raise NotImplementedError
+
+    # MQTT send message
     async def _send_message(
         self, to_agent_uuid: UUID, payload: dict, sub_topic: str
     ):
@@ -251,7 +336,7 @@ class Agent(ABC):
         await self._messager.send_message(topic, payload)
 
     async def send_message_to_agent(
-        self, to_agent_uuid: UUID, content: dict
+        self, to_agent_uuid: UUID, content: str
     ):
         """通过 Messager 发送消息"""
         if self._messager is None:
@@ -259,29 +344,28 @@ class Agent(ABC):
         payload = {
             "from": self._uuid,
             "content": content,
-            "timestamp": int(time.time()),
+            "timestamp": int(datetime.now().timestamp() * 1000),
             "day": await self._simulator.get_simulator_day(),
             "t": await self._simulator.get_simulator_second_from_start_of_day(),
         }
         await self._send_message(to_agent_uuid, payload, "agent-chat")
+        auros = [{
+            "id": str(self._uuid),
+            "day": await self._simulator.get_simulator_day(),
+            "t": await self._simulator.get_simulator_second_from_start_of_day(),
+            "type": 1,
+            "speaker": str(self._uuid),
+            "content": content,
+            "created_at": int(datetime.now().timestamp() * 1000),
+        }]
+        with open(self._avro_file["dialog"], "a+b") as f:
+            fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
 
-    async def send_message_to_user(
-        self, content: dict
-    ):
-        pass
-
-    async def send_message_to_survey(
-        self, content: dict
-    ):
-        pass
-
+    # Agent logic
     @abstractmethod
     async def forward(self) -> None:
         """智能体行为逻辑"""
         raise NotImplementedError
-    
-    async def handle_gather_message(self, payload: str):
-        raise NotImplementedError
 
     async def run(self) -> None:
         """
@@ -305,6 +389,7 @@ class CitizenAgent(Agent):
         memory: Optional[Memory] = None,
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,
+        avro_file: Optional[dict] = None,
     ) -> None:
         super().__init__(
             name,
@@ -314,6 +399,7 @@ class CitizenAgent(Agent):
             messager,
             simulator,
             memory,
+            avro_file,
         )
 
     async def bind_to_simulator(self):
@@ -417,6 +503,7 @@ class InstitutionAgent(Agent):
         memory: Optional[Memory] = None,
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,
+        avro_file: Optional[dict] = None,
     ) -> None:
         super().__init__(
             name,
@@ -426,6 +513,7 @@ class InstitutionAgent(Agent):
             messager,
             simulator,
             memory,
+            avro_file,
         )
         # 添加响应收集器
         self._gather_responses: Dict[str, asyncio.Future] = {}

pycityagent/economy/econ_client.py

@@ -8,6 +8,8 @@ import pycityproto.city.economy.v2.org_service_pb2 as org_service
 import pycityproto.city.economy.v2.org_service_pb2_grpc as org_grpc
 from google.protobuf import descriptor
 
+economyv2.ORG_TYPE_BANK
+
 __all__ = [
     "EconomyClient",
 ]

pycityagent/memory/const.py

@@ -1,6 +1,7 @@
 from pycityproto.city.person.v2.motion_pb2 import Status
 
 PROFILE_ATTRIBUTES = {
+    "name": str(),
     "gender": str(),
     "age": float(),
     "education": str(),

pycityagent/simulation/agentgroup.py

@@ -1,23 +1,27 @@
 import asyncio
+from datetime import datetime
 import json
 import logging
 import uuid
+import fastavro
 import ray
 from uuid import UUID
-from pycityagent.agent import Agent
+from pycityagent.agent import Agent, CitizenAgent
 from pycityagent.economy.econ_client import EconomyClient
 from pycityagent.environment.simulator import Simulator
 from pycityagent.llm.llm import LLM
 from pycityagent.llm.llmconfig import LLMConfig
 from pycityagent.message import Messager
+from pycityagent.utils import STATUS_SCHEMA
 from typing import Any
 
 @ray.remote
 class AgentGroup:
-    def __init__(self, agents: list[Agent], config: dict, exp_id: str|UUID):
+    def __init__(self, agents: list[Agent], config: dict, exp_id: str|UUID, avro_file: dict):
         self.agents = agents
         self.config = config
         self.exp_id = exp_id
+        self.avro_file = avro_file
         self._uuid = uuid.uuid4()
         self.messager = Messager(
             hostname=config["simulator_request"]["mqtt"]["server"],
@@ -117,7 +121,7 @@ class AgentGroup:
                     elif topic_type == "gather":
                         await agent.handle_gather_message(payload)
 
-            await asyncio.sleep(1)
+            await asyncio.sleep(0.5)
 
     async def step(self):
         if not self.initialized:
@@ -125,6 +129,40 @@ class AgentGroup:
 
         tasks = [agent.run() for agent in self.agents]
         await asyncio.gather(*tasks)
+        avros = []
+        for agent in self.agents:
+            if not issubclass(type(agent), CitizenAgent):
+                continue
+            position = await agent.memory.get("position")
+            lng = position["longlat_position"]["longitude"]
+            lat = position["longlat_position"]["latitude"]
+            if "aoi_position" in position:
+                parent_id = position["aoi_position"]["aoi_id"]
+            elif "lane_position" in position:
+                parent_id = position["lane_position"]["lane_id"]
+            else:
+                # BUG: 需要处理
+                parent_id = -1
+            needs = await agent.memory.get("needs")
+            action = await agent.memory.get("current_step")
+            action = action["intention"]
+            avro = {
+                "id": str(agent._uuid),  # uuid as string
+                "day": await self.simulator.get_simulator_day(),
+                "t": await self.simulator.get_simulator_second_from_start_of_day(),
+                "lng": lng,
+                "lat": lat,
+                "parent_id": parent_id,
+                "action": action,
+                "hungry": needs["hungry"],
+                "tired": needs["tired"],
+                "safe": needs["safe"],
+                "social": needs["social"],
+                "created_at": int(datetime.now().timestamp() * 1000),
+            }
+            avros.append(avro)
+        with open(self.avro_file["status"], "a+b") as f:
+            fastavro.writer(f, STATUS_SCHEMA, avros, codec="snappy")
 
     async def run(self, day: int = 1):
         """运行模拟器

pycityagent/simulation/simulation.py

@@ -1,18 +1,21 @@
 import asyncio
 import json
 import logging
+import os
+from pathlib import Path
 import uuid
 from datetime import datetime
 import random
 from typing import Dict, List, Optional, Callable, Union,Any
+import fastavro
 from mosstool.map._map_util.const import AOI_START_ID
 import pycityproto.city.economy.v2.economy_pb2 as economyv2
 from pycityagent.memory.memory import Memory
+from pycityagent.message.messager import Messager
+from pycityagent.survey import Survey
+from pycityagent.utils.avro_schema import PROFILE_SCHEMA, DIALOG_SCHEMA, STATUS_SCHEMA, SURVEY_SCHEMA
 
 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__)
@@ -44,10 +47,34 @@ class AgentSimulation:
         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"]
+        self._avro_path = Path(config["storage"]["avro"]["path"])
+        self._avro_file = {
+            "profile": self._avro_path / f"{self.exp_id}_profile.avro",
+            "dialog": self._avro_path / f"{self.exp_id}_dialog.avro",
+            "status": self._avro_path / f"{self.exp_id}_status.avro",
+            "survey": self._avro_path / f"{self.exp_id}_survey.avro",
+        }
+
+        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"]
 
     @property
     def agents(self):
@@ -126,6 +153,7 @@ class AgentSimulation:
                 agent = agent_class(
                     name=agent_name,
                     memory=memory,
+                    avro_file=self._avro_file,
                 )
 
                 self._agents[agent._uuid] = agent
@@ -145,7 +173,7 @@ class AgentSimulation:
                 # 获取当前组的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)
+                group = AgentGroup.remote(agents, self.config, self.exp_id, self._avro_file)
                 self._groups[group_name] = group
                 for agent in agents:
                     self._agent_uuid2group[agent._uuid] = group
@@ -156,6 +184,42 @@ class AgentSimulation:
         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"
+
+        # save profile
+        if self._enable_avro:
+            self._avro_path.mkdir(parents=True, exist_ok=True)
+            # profile
+            filename = self._avro_file["profile"]
+            with open(filename, "wb") as f:
+                profiles = []
+                for agent in self._agents.values():
+                    profile = await agent.memory._profile.export()
+                    profile = profile[0]
+                    profile['id'] = str(agent._uuid)
+                    profiles.append(profile)
+                fastavro.writer(f, PROFILE_SCHEMA, profiles)
+
+            # dialog
+            filename = self._avro_file["dialog"]
+            with open(filename, "wb") as f:
+                dialogs = []
+                fastavro.writer(f, DIALOG_SCHEMA, dialogs)
+
+            # status
+            filename = self._avro_file["status"]
+            with open(filename, "wb") as f:
+                statuses = []
+                fastavro.writer(f, STATUS_SCHEMA, statuses)
+
+            # survey
+            filename = self._avro_file["survey"]
+            with open(filename, "wb") as f:
+                surveys = []
+                fastavro.writer(f, SURVEY_SCHEMA, surveys)
+
 
     async def gather(self, content: str):
         """收集智能体的特定信息"""
@@ -219,6 +283,7 @@ class AgentSimulation:
         }
 
         PROFILE = {
+            "name": "unknown",
             "gender": random.choice(["male", "female"]),
             "education": random.choice(
                 ["Doctor", "Master", "Bachelor", "College", "High School"]
@@ -251,7 +316,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 +350,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"""

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())