pycityagent 2.0.0a12__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,8 +513,11 @@ class InstitutionAgent(Agent):
             messager,
             simulator,
             memory,
+            avro_file,
         )
-
+        # 添加响应收集器
+        self._gather_responses: Dict[str, asyncio.Future] = {}
+        
     async def bind_to_simulator(self):
         await self._bind_to_economy()
 
@@ -514,18 +604,47 @@ class InstitutionAgent(Agent):
 
     async def handle_gather_message(self, payload: dict):
         """处理收到的消息，识别发送者"""
-        # 从消息中解析发送者 ID 和消息内容
         content = payload["content"]
         sender_id = payload["from"]
-        print(
-            f"Agent {self._uuid} received gather message: '{content}' from Agent {sender_id}"
-        )
-
-    async def gather_messages(self, agent_ids: list[UUID], target: str):
-        """从多个智能体收集消息"""
+        
+        # 将响应存储到对应的Future中
+        response_key = str(sender_id)
+        if response_key in self._gather_responses:
+            self._gather_responses[response_key].set_result({
+                "from": sender_id,
+                "content": content,
+            })
+
+    async def gather_messages(self, agent_ids: list[UUID], target: str) -> List[dict]:
+        """从多个智能体收集消息
+        
+        Args:
+            agent_ids: 目标智能体ID列表
+            target: 要收集的信息类型
+            
+        Returns:
+            List[dict]: 收集到的所有响应
+        """
+        # 为每个agent创建Future
+        futures = {}
+        for agent_id in agent_ids:
+            response_key = str(agent_id)
+            futures[response_key] = asyncio.Future()
+            self._gather_responses[response_key] = futures[response_key]
+            
+        # 发送gather请求
         payload = {
             "from": self._uuid,
             "target": target,
         }
         for agent_id in agent_ids:
             await self._send_message(agent_id, payload, "gather")
+            
+        try:
+            # 等待所有响应
+            responses = await asyncio.gather(*futures.values())
+            return responses
+        finally:
+            # 清理Future
+            for key in futures:
+                self._gather_responses.pop(key, None)

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"],
@@ -71,7 +75,8 @@ class AgentGroup:
                 topic = f"exps/{self.exp_id}/agents/{agent._uuid}/gather"
                 await self.messager.subscribe(topic, agent)
         self.initialized = True
-
+        self.message_dispatch_task = asyncio.create_task(self.message_dispatch())
+        
     async def gather(self, content: str):
         results = {}
         for agent in self.agents:
@@ -82,51 +87,82 @@ class AgentGroup:
         agent = self.id2agent[target_agent_uuid]
         await agent.memory.update(target_key, content)
 
+    async def message_dispatch(self):
+        while True:
+            if not self.messager.is_connected():
+                logging.warning("Messager is not connected. Skipping message processing.")
+
+            # Step 1: 获取消息
+            messages = await self.messager.fetch_messages()
+            logging.info(f"Group {self._uuid} received {len(messages)} messages")
+
+            # Step 2: 分发消息到对应的 Agent
+            for message in messages:
+                topic = message.topic.value
+                payload = message.payload
+
+                # 添加解码步骤，将bytes转换为str
+                if isinstance(payload, bytes):
+                    payload = payload.decode("utf-8")
+                    payload = json.loads(payload)
+
+                # 提取 agent_id（主题格式为 "exps/{exp_id}/agents/{agent_uuid}/{topic_type}"）
+                _, _, _, agent_uuid, topic_type = topic.strip("/").split("/")
+                    
+                if uuid.UUID(agent_uuid) in self.id2agent:
+                    agent = self.id2agent[uuid.UUID(agent_uuid)]
+                    # topic_type: agent-chat, user-chat, user-survey, gather
+                    if topic_type == "agent-chat":
+                        await agent.handle_agent_chat_message(payload)
+                    elif topic_type == "user-chat":
+                        await agent.handle_user_chat_message(payload)
+                    elif topic_type == "user-survey":
+                        await agent.handle_user_survey_message(payload)
+                    elif topic_type == "gather":
+                        await agent.handle_gather_message(payload)
+
+            await asyncio.sleep(0.5)
+
     async def step(self):
         if not self.initialized:
             await self.init_agents()
 
-        # Step 1: 如果 Messager 无法连接，则跳过消息接收
-        if not self.messager.is_connected():
-            logging.warning("Messager is not connected. Skipping message processing.")
-            # 跳过接收和分发消息
-            tasks = [agent.run() for agent in self.agents]
-            await asyncio.gather(*tasks)
-            return
-
-        # Step 2: 从 Messager 获取消息
-        messages = await self.messager.fetch_messages()
-
-        logging.info(f"Group {self._uuid} received {len(messages)} messages")
-
-        # Step 3: 分发消息到对应的 Agent
-        for message in messages:
-            topic = message.topic.value
-            payload = message.payload
-
-            # 添加解码步骤，将bytes转换为str
-            if isinstance(payload, bytes):
-                payload = payload.decode("utf-8")
-                payload = json.loads(payload)
-
-            # 提取 agent_id（主题格式为 "exps/{exp_id}/agents/{agent_uuid}/{topic_type}"）
-            _, _, _, agent_uuid, topic_type = topic.strip("/").split("/")
-                
-            if uuid.UUID(agent_uuid) in self.id2agent:
-                agent = self.id2agent[uuid.UUID(agent_uuid)]
-                # topic_type: agent-chat, user-chat, user-survey, gather
-                if topic_type == "agent-chat":
-                    await agent.handle_agent_chat_message(payload)
-                elif topic_type == "user-chat":
-                    await agent.handle_user_chat_message(payload)
-                elif topic_type == "user-survey":
-                    await agent.handle_user_survey_message(payload)
-                elif topic_type == "gather":
-                    await agent.handle_gather_message(payload)
-
-        # Step 4: 调用每个 Agent 的运行逻辑
         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):
         """运行模拟器