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

pycityagent/__init__.py

@@ -4,5 +4,19 @@ Pycityagent: 城市智能体构建框架
 
 from .agent import Agent, CitizenAgent, InstitutionAgent
 from .environment import Simulator
+import logging
+
+# 创建一个 pycityagent 记录器
+logger = logging.getLogger("pycityagent")
+logger.setLevel(logging.WARNING)  # 默认级别
+
+# 如果没有处理器，则添加一个
+if not logger.hasHandlers():
+    handler = logging.StreamHandler()
+    formatter = logging.Formatter(
+        "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
 
 __all__ = ["Agent", "Simulator", "CitizenAgent", "InstitutionAgent"]

pycityagent/agent.py

@@ -2,16 +2,14 @@
 
 from abc import ABC, abstractmethod
 import asyncio
-import json
 from uuid import UUID
 from copy import deepcopy
 from datetime import datetime
-import time
 from enum import Enum
 import logging
 import random
 import uuid
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional,Any
 
 import fastavro
 
@@ -28,6 +26,8 @@ from .environment import Simulator
 from .llm import LLM
 from .memory import Memory
 
+logger = logging.getLogger("pycityagent")
+
 
 class AgentType(Enum):
     """
@@ -73,13 +73,14 @@ class Agent(ABC):
         """
         self._name = name
         self._type = type
-        self._uuid = uuid.uuid4()
+        self._uuid = str(uuid.uuid4())
         self._llm_client = llm_client
         self._economy_client = economy_client
         self._messager = messager
         self._simulator = simulator
         self._memory = memory
         self._exp_id = -1
+        self._agent_id = -1
         self._has_bound_to_simulator = False
         self._has_bound_to_economy = False
         self._blocked = False
@@ -123,12 +124,18 @@ class Agent(ABC):
         """
         self._memory = memory
 
-    def set_exp_id(self, exp_id: str|UUID):
+    def set_exp_id(self, exp_id: str):
         """
         Set the exp_id of the agent.
         """
         self._exp_id = exp_id
 
+    def set_avro_file(self, avro_file: Dict[str, str]):
+        """
+        Set the avro file of the agent.
+        """
+        self._avro_file = avro_file
+
     @property
     def uuid(self):
         """The Agent's UUID"""
@@ -214,10 +221,12 @@ class Agent(ABC):
     
     async def _process_survey(self, survey: dict):
         survey_response = await self.generate_user_survey_response(survey)
+        if self._avro_file is None:
+            return
         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(),
+            "id": 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),
@@ -264,9 +273,9 @@ class Agent(ABC):
     
     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(),
+            "id": 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"],
@@ -275,23 +284,27 @@ class Agent(ABC):
         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(),
+            "id": 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),
         })
+        if self._avro_file is None:
+            return
         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}")
+        resp = f"Agent {self._uuid} received agent chat response: {payload}"
+        logger.info(resp)
+        return resp
 
     async def _process_agent_chat(self, payload: dict):
         auros = [{
-            "id": str(self._uuid),
+            "id": self._uuid,
             "day": payload["day"],
             "t": payload["t"],
             "type": 1,
@@ -300,6 +313,8 @@ class Agent(ABC):
             "created_at": int(datetime.now().timestamp() * 1000),
         }]
         asyncio.create_task(self.process_agent_chat_response(payload))
+        if self._avro_file is None:
+            return
         with open(self._avro_file["dialog"], "a+b") as f:
             fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
 
@@ -307,27 +322,27 @@ class Agent(ABC):
     async def handle_agent_chat_message(self, payload: dict):
         """处理收到的消息，识别发送者"""
         # 从消息中解析发送者 ID 和消息内容
-        logging.info(f"Agent {self._uuid} received agent chat message: {payload}")
+        logger.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 和消息内容
-        logging.info(f"Agent {self._uuid} received user chat message: {payload}")
+        logger.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 和消息内容
-        logging.info(f"Agent {self._uuid} received user survey message: {payload}")
+        logger.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):
+    async def handle_gather_message(self, payload: Any):
         raise NotImplementedError
 
     # MQTT send message
     async def _send_message(
-        self, to_agent_uuid: UUID, payload: dict, sub_topic: str
+        self, to_agent_uuid: str, payload: dict, sub_topic: str
     ):
         """通过 Messager 发送消息"""
         if self._messager is None:
@@ -336,7 +351,7 @@ class Agent(ABC):
         await self._messager.send_message(topic, payload)
 
     async def send_message_to_agent(
-        self, to_agent_uuid: UUID, content: str
+        self, to_agent_uuid: str, content: str
     ):
         """通过 Messager 发送消息"""
         if self._messager is None:
@@ -345,19 +360,21 @@ class Agent(ABC):
             "from": self._uuid,
             "content": content,
             "timestamp": int(datetime.now().timestamp() * 1000),
-            "day": await self._simulator.get_simulator_day(),
-            "t": await self._simulator.get_simulator_second_from_start_of_day(),
+            "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(),
+            "id": 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),
+            "speaker": self._uuid,
             "content": content,
             "created_at": int(datetime.now().timestamp() * 1000),
         }]
+        if self._avro_file is None:
+            return
         with open(self._avro_file["dialog"], "a+b") as f:
             fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
 
@@ -414,7 +431,7 @@ class CitizenAgent(Agent):
             person_template (dict, optional): The person template in dict format. Defaults to PersonService.default_dict_person().
         """
         if self._simulator is None:
-            logging.warning("Simulator is not set")
+            logger.warning("Simulator is not set")
             return
         if not self._has_bound_to_simulator:
             FROM_MEMORY_KEYS = {
@@ -426,13 +443,13 @@ class CitizenAgent(Agent):
                 "pedestrian_attribute",
                 "bike_attribute",
             }
-            simulator = self._simulator
-            memory = self._memory
+            simulator = self.simulator
+            memory = self.memory
             person_id = await memory.get("id")
             # ATTENTION:模拟器分配的id从0开始
             if person_id >= 0:
                 await simulator.get_person(person_id)
-                logging.debug(f"Binding to Person `{person_id}` already in Simulator")
+                logger.debug(f"Binding to Person `{person_id}` already in Simulator")
             else:
                 dict_person = deepcopy(self._person_template)
                 for _key in FROM_MEMORY_KEYS:
@@ -447,7 +464,7 @@ class CitizenAgent(Agent):
                 )
                 person_id = resp["person_id"]
                 await memory.update("id", person_id, protect_llm_read_only_fields=False)
-                logging.debug(
+                logger.debug(
                     f"Binding to Person `{person_id}` just added to Simulator"
                 )
                 # 防止模拟器还没有到prepare阶段导致get_person出错
@@ -456,7 +473,7 @@ class CitizenAgent(Agent):
 
     async def _bind_to_economy(self):
         if self._economy_client is None:
-            logging.warning("Economy client is not set")
+            logger.warning("Economy client is not set")
             return
         if not self._has_bound_to_economy:
             if self._has_bound_to_simulator:
@@ -464,16 +481,16 @@ class CitizenAgent(Agent):
                     await self._economy_client.remove_agents([self._agent_id])
                 except:
                     pass
-                person_id = await self._memory.get("id")
+                person_id = await self.memory.get("id")
                 await self._economy_client.add_agents(
                     {
                         "id": person_id,
-                        "currency": await self._memory.get("currency"),
+                        "currency": await self.memory.get("currency"),
                     }
                 )
                 self._has_bound_to_economy = True
             else:
-                logging.debug(
+                logger.debug(
                     f"Binding to Economy before binding to Simulator, skip binding to Economy Simulator"
                 )
 
@@ -523,68 +540,69 @@ class InstitutionAgent(Agent):
 
     async def _bind_to_economy(self):
         if self._economy_client is None:
-            logging.debug("Economy client is not set")
+            logger.debug("Economy client is not set")
             return
         if not self._has_bound_to_economy:
             # TODO: More general id generation
             _id = random.randint(100000, 999999)
             self._agent_id = _id
-            await self._memory.update("id", _id, protect_llm_read_only_fields=False)
+            await self.memory.update("id", _id, protect_llm_read_only_fields=False)
             try:
                 await self._economy_client.remove_orgs([self._agent_id])
             except:
                 pass
             try:
-                id = await self._memory.get("id")
-                type = await self._memory.get("type")
+                _memory = self.memory
+                _id = await _memory.get("id")
+                _type = await _memory.get("type")
                 try:
-                    nominal_gdp = await self._memory.get("nominal_gdp")
+                    nominal_gdp = await _memory.get("nominal_gdp")
                 except:
                     nominal_gdp = []
                 try:
-                    real_gdp = await self._memory.get("real_gdp")
+                    real_gdp = await _memory.get("real_gdp")
                 except:
                     real_gdp = []
                 try:
-                    unemployment = await self._memory.get("unemployment")
+                    unemployment = await _memory.get("unemployment")
                 except:
                     unemployment = []
                 try:
-                    wages = await self._memory.get("wages")
+                    wages = await _memory.get("wages")
                 except:
                     wages = []
                 try:
-                    prices = await self._memory.get("prices")
+                    prices = await _memory.get("prices")
                 except:
                     prices = []
                 try:
-                    inventory = await self._memory.get("inventory")
+                    inventory = await _memory.get("inventory")
                 except:
                     inventory = 0
                 try:
-                    price = await self._memory.get("price")
+                    price = await _memory.get("price")
                 except:
                     price = 0
                 try:
-                    currency = await self._memory.get("currency")
+                    currency = await _memory.get("currency")
                 except:
                     currency = 0.0
                 try:
-                    interest_rate = await self._memory.get("interest_rate")
+                    interest_rate = await _memory.get("interest_rate")
                 except:
                     interest_rate = 0.0
                 try:
-                    bracket_cutoffs = await self._memory.get("bracket_cutoffs")
+                    bracket_cutoffs = await _memory.get("bracket_cutoffs")
                 except:
                     bracket_cutoffs = []
                 try:
-                    bracket_rates = await self._memory.get("bracket_rates")
+                    bracket_rates = await _memory.get("bracket_rates")
                 except:
                     bracket_rates = []
                 await self._economy_client.add_orgs(
                     {
-                        "id": id,
-                        "type": type,
+                        "id": _id,
+                        "type": _type,
                         "nominal_gdp": nominal_gdp,
                         "real_gdp": real_gdp,
                         "unemployment": unemployment,
@@ -599,7 +617,7 @@ class InstitutionAgent(Agent):
                     }
                 )
             except Exception as e:
-                logging.error(f"Failed to bind to Economy: {e}")
+                logger.error(f"Failed to bind to Economy: {e}")
             self._has_bound_to_economy = True
 
     async def handle_gather_message(self, payload: dict):
@@ -615,11 +633,11 @@ class InstitutionAgent(Agent):
                 "content": content,
             })
 
-    async def gather_messages(self, agent_ids: list[UUID], target: str) -> List[dict]:
+    async def gather_messages(self, agent_uuids: list[str], target: str) -> List[dict]:
         """从多个智能体收集消息
         
         Args:
-            agent_ids: 目标智能体ID列表
+            agent_uuids: 目标智能体UUID列表
             target: 要收集的信息类型
             
         Returns:
@@ -627,18 +645,17 @@ class InstitutionAgent(Agent):
         """
         # 为每个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]
+        for agent_uuid in agent_uuids:
+            futures[agent_uuid] = asyncio.Future()
+            self._gather_responses[agent_uuid] = futures[agent_uuid]
             
         # 发送gather请求
         payload = {
             "from": self._uuid,
             "target": target,
         }
-        for agent_id in agent_ids:
-            await self._send_message(agent_id, payload, "gather")
+        for agent_uuid in agent_uuids:
+            await self._send_message(agent_uuid, payload, "gather")
             
         try:
             # 等待所有响应

pycityagent/environment/simulator.py

@@ -20,6 +20,7 @@ from shapely.strtree import STRtree
 from .sim import CityClient, ControlSimEnv
 from .utils.const import *
 
+logger = logging.getLogger("pycityagent")
 
 class Simulator:
     """
@@ -72,7 +73,7 @@ class Simulator:
             else:
                 self._client = CityClient(config["simulator"]["server"], secure=False)
         else:
-            logging.warning(
+            logger.warning(
                 "No simulator config found, no simulator client will be used"
             )
         self.map = SimMap(
@@ -285,7 +286,7 @@ class Simulator:
             reset_position["aoi_position"] = {"aoi_id": aoi_id}
             if poi_id is not None:
                 reset_position["aoi_position"]["poi_id"] = poi_id
-            logging.debug(
+            logger.debug(
                 f"Setting person {person_id} pos to AoiPosition {reset_position}"
             )
             await self._client.person_service.ResetPersonPosition(
@@ -298,14 +299,14 @@ class Simulator:
             }
             if s is not None:
                 reset_position["lane_position"]["s"] = s
-            logging.debug(
+            logger.debug(
                 f"Setting person {person_id} pos to LanePosition {reset_position}"
             )
             await self._client.person_service.ResetPersonPosition(
                 {"person_id": person_id, "position": reset_position}
             )
         else:
-            logging.debug(
+            logger.debug(
                 f"Neither aoi or lane pos provided for person {person_id} position reset!!"
             )
 

pycityagent/memory/memory.py

@@ -13,6 +13,7 @@ from .profile import ProfileMemory
 from .self_define import DynamicMemory
 from .state import StateMemory
 
+logger = logging.getLogger("pycityagent")
 
 class Memory:
     """
@@ -83,7 +84,7 @@ class Memory:
                                 _type.extend(_value)
                                 _value = deepcopy(_type)
                             else:
-                                logging.warning(f"type `{_type}` is not supported!")
+                                logger.warning(f"type `{_type}` is not supported!")
                                 pass
                     except TypeError as e:
                         pass
@@ -99,7 +100,7 @@ class Memory:
                     or k in STATE_ATTRIBUTES
                     or k == TIME_STAMP_KEY
                 ):
-                    logging.warning(f"key `{k}` already declared in memory!")
+                    logger.warning(f"key `{k}` already declared in memory!")
                     continue
 
                 _dynamic_config[k] = deepcopy(_value)
@@ -112,19 +113,19 @@ class Memory:
         if profile is not None:
             for k, v in profile.items():
                 if k not in PROFILE_ATTRIBUTES:
-                    logging.warning(f"key `{k}` is not a correct `profile` field!")
+                    logger.warning(f"key `{k}` is not a correct `profile` field!")
                     continue
                 _profile_config[k] = v
         if motion is not None:
             for k, v in motion.items():
                 if k not in STATE_ATTRIBUTES:
-                    logging.warning(f"key `{k}` is not a correct `motion` field!")
+                    logger.warning(f"key `{k}` is not a correct `motion` field!")
                     continue
                 _state_config[k] = v
         if base is not None:
             for k, v in base.items():
                 if k not in STATE_ATTRIBUTES:
-                    logging.warning(f"key `{k}` is not a correct `base` field!")
+                    logger.warning(f"key `{k}` is not a correct `base` field!")
                     continue
                 _state_config[k] = v
         self._state = StateMemory(
@@ -182,7 +183,7 @@ class Memory:
         """更新记忆值并在必要时更新embedding"""
         if protect_llm_read_only_fields:
             if any(key in _attrs for _attrs in [STATE_ATTRIBUTES]):
-                logging.warning(f"Trying to write protected key `{key}`!")
+                logger.warning(f"Trying to write protected key `{key}`!")
                 return
         for _mem in [self._state, self._profile, self._dynamic]:
             try:
@@ -208,7 +209,7 @@ class Memory:
                     elif isinstance(original_value, deque):
                         original_value.extend(deque(value))
                     else:
-                        logging.debug(
+                        logger.debug(
                             f"Type of {type(original_value)} does not support mode `merge`, using `replace` instead!"
                         )
                         await _mem.update(key, value, store_snapshot)

pycityagent/memory/memory_base.py

@@ -10,6 +10,8 @@ from typing import Any, Callable, Dict, Optional, Sequence, Union
 
 from .const import *
 
+logger = logging.getLogger("pycityagent")
+
 
 class MemoryUnit:
     def __init__(
@@ -57,7 +59,7 @@ class MemoryUnit:
                 orig_v = self._content[k]
                 orig_type, new_type = type(orig_v), type(v)
                 if not orig_type == new_type:
-                    logging.debug(
+                    logger.debug(
                         f"Type warning: The type of the value for key '{k}' is changing from `{orig_type.__name__}` to `{new_type.__name__}`!"
                     )
         self._content.update(content)
@@ -82,7 +84,7 @@ class MemoryUnit:
         await self._lock.acquire()
         values = self._content[key]
         if not isinstance(values, Sequence):
-            logging.warning(
+            logger.warning(
                 f"the value stored in key `{key}` is not `sequence`, return value `{values}` instead!"
             )
             return values
@@ -93,7 +95,7 @@ class MemoryUnit:
             )
             top_k = len(values) if top_k is None else top_k
             if len(_sorted_values_with_idx) < top_k:
-                logging.debug(
+                logger.debug(
                     f"Length of values {len(_sorted_values_with_idx)} is less than top_k {top_k}, returning all values."
                 )
             self._lock.release()
@@ -149,7 +151,7 @@ class MemoryBase(ABC):
         if recent_n is None:
             return _list_units
         if len(_memories) < recent_n:
-            logging.debug(
+            logger.debug(
                 f"Length of memory {len(_memories)} is less than recent_n {recent_n}, returning all available memories."
             )
         return _list_units[-recent_n:]

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.")