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

pycityagent/message/messager.py

@@ -5,6 +5,7 @@ import logging
 import math
 from aiomqtt import Client
 
+logger = logging.getLogger("pycityagent")
 
 class Messager:
     def __init__(
@@ -21,15 +22,15 @@ class Messager:
         try:
             await self.client.__aenter__()
             self.connected = True
-            logging.info("Connected to MQTT Broker")
+            logger.info("Connected to MQTT Broker")
         except Exception as e:
             self.connected = False
-            logging.error(f"Failed to connect to MQTT Broker: {e}")
+            logger.error(f"Failed to connect to MQTT Broker: {e}")
 
     async def disconnect(self):
         await self.client.__aexit__(None, None, None)
         self.connected = False
-        logging.info("Disconnected from MQTT Broker")
+        logger.info("Disconnected from MQTT Broker")
 
     def is_connected(self):
         """检查是否成功连接到 Broker"""
@@ -37,13 +38,13 @@ class Messager:
 
     async def subscribe(self, topic, agent):
         if not self.is_connected():
-            logging.error(
+            logger.error(
                 f"Cannot subscribe to {topic} because not connected to the Broker."
             )
             return
         await self.client.subscribe(topic)
         self.subscribers[topic] = agent
-        logging.info(f"Subscribed to {topic} for Agent {agent._uuid}")
+        logger.info(f"Subscribed to {topic} for Agent {agent._uuid}")
 
     async def receive_messages(self):
         """监听并将消息存入队列"""
@@ -61,11 +62,11 @@ class Messager:
         """通过 Messager 发送消息"""
         message = json.dumps(payload, default=str)
         await self.client.publish(topic, message)
-        logging.info(f"Message sent to {topic}: {message}")
+        logger.info(f"Message sent to {topic}: {message}")
 
     async def start_listening(self):
         """启动消息监听任务"""
         if self.is_connected():
             asyncio.create_task(self.receive_messages())
         else:
-            logging.error("Cannot start listening because not connected to the Broker.")
+            logger.error("Cannot start listening because not connected to the Broker.")

pycityagent/simulation/agentgroup.py

@@ -1,24 +1,42 @@
 import asyncio
+from datetime import datetime
 import json
 import logging
+from pathlib import Path
 import uuid
+import fastavro
 import ray
 from uuid import UUID
-from pycityagent.agent import Agent
+from pycityagent.agent import Agent, CitizenAgent, InstitutionAgent
 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, PROFILE_SCHEMA, DIALOG_SCHEMA, SURVEY_SCHEMA, INSTITUTION_STATUS_SCHEMA
 from typing import Any
 
+logger = logging.getLogger("pycityagent")
+
 @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, enable_avro: bool, avro_path: Path, logging_level: int = logging.WARNING):
+        logger.setLevel(logging_level)
+        self._uuid = str(uuid.uuid4())
         self.agents = agents
         self.config = config
         self.exp_id = exp_id
-        self._uuid = uuid.uuid4()
+        self.enable_avro = enable_avro
+        self.avro_path = avro_path / f"{self._uuid}"
+        if enable_avro:
+            self.avro_path.mkdir(parents=True, exist_ok=True)
+            self.avro_file = {
+                "profile": self.avro_path / f"profile.avro",
+                "dialog": self.avro_path / f"dialog.avro",
+                "status": self.avro_path / f"status.avro",
+                "survey": self.avro_path / f"survey.avro",
+            }
+            
         self.messager = Messager(
             hostname=config["simulator_request"]["mqtt"]["server"],
             port=config["simulator_request"]["mqtt"]["port"],
@@ -29,16 +47,16 @@ class AgentGroup:
         self.id2agent = {}
         # Step:1 prepare LLM client
         llmConfig = LLMConfig(config["llm_request"])
-        logging.info("-----Creating LLM client in remote...")
+        logger.info(f"-----Creating LLM client in AgentGroup {self._uuid} ...")
         self.llm = LLM(llmConfig)
 
         # Step:2 prepare Simulator
-        logging.info("-----Creating Simulator in remote...")
+        logger.info(f"-----Creating Simulator in AgentGroup {self._uuid} ...")
         self.simulator = Simulator(config["simulator_request"])
 
         # Step:3 prepare Economy client
         if "economy" in config["simulator_request"]:
-            logging.info("-----Creating Economy client in remote...")
+            logger.info(f"-----Creating Economy client in AgentGroup {self._uuid} ...")
             self.economy_client = EconomyClient(
                 config["simulator_request"]["economy"]["server"]
             )
@@ -52,11 +70,16 @@ class AgentGroup:
             if self.economy_client is not None:
                 agent.set_economy_client(self.economy_client)
             agent.set_messager(self.messager)
+            if self.enable_avro:
+                agent.set_avro_file(self.avro_file)
 
     async def init_agents(self):
+        logger.debug(f"-----Initializing Agents in AgentGroup {self._uuid} ...")
+        logger.debug(f"-----Binding Agents to Simulator in AgentGroup {self._uuid} ...")
         for agent in self.agents:
             await agent.bind_to_simulator()
         self.id2agent = {agent._uuid: agent for agent in self.agents}
+        logger.debug(f"-----Binding Agents to Messager in AgentGroup {self._uuid} ...")
         await self.messager.connect()
         if self.messager.is_connected():
             await self.messager.start_listening()
@@ -70,27 +93,65 @@ class AgentGroup:
                 await self.messager.subscribe(topic, agent)
                 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())
-        
+        if self.enable_avro:
+            logger.debug(f"-----Creating Avro files in AgentGroup {self._uuid} ...")
+            # profile
+            if not issubclass(type(self.agents[0]), InstitutionAgent):
+                filename = self.avro_file["profile"]
+                with open(filename, "wb") as f:
+                    profiles = []
+                    for agent in self.agents:
+                        profile = await agent.memory._profile.export()
+                        profile = profile[0]
+                        profile['id'] = 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 = []
+                if not issubclass(type(self.agents[0]), InstitutionAgent):
+                    fastavro.writer(f, STATUS_SCHEMA, statuses)
+                else:
+                    fastavro.writer(f, INSTITUTION_STATUS_SCHEMA, statuses)
+
+            # survey
+            filename = self.avro_file["survey"]
+            with open(filename, "wb") as f:
+                surveys = []
+                fastavro.writer(f, SURVEY_SCHEMA, surveys)
+        self.initialized = True
+        logger.debug(f"-----AgentGroup {self._uuid} initialized")
+
     async def gather(self, content: str):
+        logger.debug(f"-----Gathering {content} from all agents in group {self._uuid}")
         results = {}
         for agent in self.agents:
             results[agent._uuid] = await agent.memory.get(content)
         return results
 
     async def update(self, target_agent_uuid: str, target_key: str, content: Any):
+        logger.debug(f"-----Updating {target_key} for agent {target_agent_uuid} in group {self._uuid}")
         agent = self.id2agent[target_agent_uuid]
         await agent.memory.update(target_key, content)
 
     async def message_dispatch(self):
+        logger.debug(f"-----Starting message dispatch for group {self._uuid}")
         while True:
             if not self.messager.is_connected():
-                logging.warning("Messager is not connected. Skipping message processing.")
+                logger.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")
+            logger.info(f"Group {self._uuid} received {len(messages)} messages")
 
             # Step 2: 分发消息到对应的 Agent
             for message in messages:
@@ -105,8 +166,8 @@ class AgentGroup:
                 # 提取 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)]
+                if agent_uuid in self.id2agent:
+                    agent = self.id2agent[agent_uuid]
                     # topic_type: agent-chat, user-chat, user-survey, gather
                     if topic_type == "agent-chat":
                         await agent.handle_agent_chat_message(payload)
@@ -117,7 +178,67 @@ class AgentGroup:
                     elif topic_type == "gather":
                         await agent.handle_gather_message(payload)
 
-            await asyncio.sleep(1)
+            await asyncio.sleep(0.5)
+
+    async def save_status(self):
+        if self.enable_avro:
+            logger.debug(f"-----Saving status for group {self._uuid}")
+            avros = []
+            if not issubclass(type(self.agents[0]), InstitutionAgent):
+                for agent in self.agents:
+                    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": agent._uuid,
+                        "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")
+            else:
+                for agent in self.agents:
+                    avro = {
+                        "id": agent._uuid,
+                        "day": await self.simulator.get_simulator_day(),
+                        "t": await self.simulator.get_simulator_second_from_start_of_day(),
+                        "type": await agent.memory.get("type"),
+                        "nominal_gdp": await agent.memory.get("nominal_gdp"),
+                        "real_gdp": await agent.memory.get("real_gdp"),
+                        "unemployment": await agent.memory.get("unemployment"),
+                        "wages": await agent.memory.get("wages"),
+                        "prices": await agent.memory.get("prices"),
+                        "inventory": await agent.memory.get("inventory"),
+                        "price": await agent.memory.get("price"),
+                        "interest_rate": await agent.memory.get("interest_rate"),
+                        "bracket_cutoffs": await agent.memory.get("bracket_cutoffs"),
+                        "bracket_rates": await agent.memory.get("bracket_rates"),
+                        "employees": await agent.memory.get("employees"),
+                        "customers": await agent.memory.get("customers"),
+                    }
+                    avros.append(avro)
+                with open(self.avro_file["status"], "a+b") as f:
+                    fastavro.writer(f, INSTITUTION_STATUS_SCHEMA, avros, codec="snappy")
 
     async def step(self):
         if not self.initialized:
@@ -125,6 +246,7 @@ class AgentGroup:
 
         tasks = [agent.run() for agent in self.agents]
         await asyncio.gather(*tasks)
+        await self.save_status()
 
     async def run(self, day: int = 1):
         """运行模拟器
@@ -147,5 +269,5 @@ class AgentGroup:
                 await self.step()
 
         except Exception as e:
-            logging.error(f"模拟器运行错误: {str(e)}")
+            logger.error(f"模拟器运行错误: {str(e)}")
             raise