pycityagent 2.0.0a13__py3-none-any.whl → 2.0.0a15__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,28 +2,32 @@
 
 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
 
+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
 from .llm import LLM
 from .memory import Memory
 
+logger = logging.getLogger("pycityagent")
+
 
 class AgentType(Enum):
     """
@@ -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,10 +69,11 @@ 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
-        self._uuid = uuid.uuid4()
+        self._uuid = str(uuid.uuid4())
         self._llm_client = llm_client
         self._economy_client = economy_client
         self._messager = messager
@@ -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()
@@ -116,12 +123,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"""
@@ -168,11 +181,61 @@ 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)
+        if self._avro_file is None:
+            return
+        response_to_avro = [{
+            "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),
+        }]
+        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 +245,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 +255,7 @@ class Agent(ABC):
                 dialog.append(
                     {
                         "role": "system",
-                        "content": f"基于以下记忆回答问题:\n{relevant_memories}",
+                        "content": f"Answer based on these memories:\n{relevant_memories}",
                     }
                 )
 
@@ -201,48 +264,82 @@ 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
+    
+    async def _process_interview(self, payload: dict):
+        auros = [{
+            "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"],
+            "created_at": int(datetime.now().timestamp() * 1000),
+        }]
+        question = payload["content"]
+        response = await self.generate_user_chat_response(question)
+        auros.append({
+            "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:
+        logger.info(f"Agent {self._uuid} received agent chat response: {payload}")
+
+    async def _process_agent_chat(self, payload: dict):
+        auros = [{
+            "id": 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))
+        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")
 
-    def get_interview_history(self) -> List[Dict]:
-        """获取采访历史记录"""
-        return self._interview_history
-
+    # 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']}"
-        )
+        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 和消息内容
-        print(
-            f"Agent {self._uuid} received user chat message: '{payload['content']}' from User"
-        )
+        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 和消息内容
-        print(
-            f"Agent {self._uuid} received user survey message: '{payload['content']}' from User"
-        )
+        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):
+        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:
@@ -251,7 +348,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: str, content: str
     ):
         """通过 Messager 发送消息"""
         if self._messager is None:
@@ -259,29 +356,30 @@ 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": self._uuid,
+            "day": await self._simulator.get_simulator_day(),
+            "t": await self._simulator.get_simulator_second_from_start_of_day(),
+            "type": 1,
+            "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")
 
-    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 +403,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 +413,7 @@ class CitizenAgent(Agent):
             messager,
             simulator,
             memory,
+            avro_file,
         )
 
     async def bind_to_simulator(self):
@@ -328,7 +428,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 = {
@@ -346,7 +446,7 @@ class CitizenAgent(Agent):
             # 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:
@@ -361,7 +461,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出错
@@ -370,7 +470,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:
@@ -387,7 +487,7 @@ class CitizenAgent(Agent):
                 )
                 self._has_bound_to_economy = True
             else:
-                logging.debug(
+                logger.debug(
                     f"Binding to Economy before binding to Simulator, skip binding to Economy Simulator"
                 )
 
@@ -417,6 +517,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 +527,7 @@ class InstitutionAgent(Agent):
             messager,
             simulator,
             memory,
+            avro_file,
         )
         # 添加响应收集器
         self._gather_responses: Dict[str, asyncio.Future] = {}
@@ -435,7 +537,7 @@ 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
@@ -511,7 +613,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):
@@ -527,11 +629,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:
@@ -539,18 +641,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/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/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/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/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:]