pycityagent 2.0.0a52__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a54__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/agent/agent.py

@@ -13,7 +13,7 @@ from ..economy import EconomyClient
 from ..environment import Simulator
 from ..llm import LLM
 from ..memory import Memory
-from ..message.messager import Messager
+from ..message import MessageInterceptor, Messager
 from ..metrics import MlflowClient
 from .agent_base import Agent, AgentType
 
@@ -30,10 +30,10 @@ class CitizenAgent(Agent):
         name: str,
         llm_client: Optional[LLM] = None,
         simulator: Optional[Simulator] = None,
-        mlflow_client: Optional[MlflowClient] = None,
         memory: Optional[Memory] = None,
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,  # type:ignore
+        message_interceptor: Optional[MessageInterceptor] = None,  # type:ignore
         avro_file: Optional[dict] = None,
     ) -> None:
         super().__init__(
@@ -42,11 +42,12 @@ class CitizenAgent(Agent):
             llm_client=llm_client,
             economy_client=economy_client,
             messager=messager,
+            message_interceptor=message_interceptor,
             simulator=simulator,
-            mlflow_client=mlflow_client,
             memory=memory,
             avro_file=avro_file,
         )
+        self._mlflow_client = None
 
     async def bind_to_simulator(self):
         await self._bind_to_simulator()
@@ -62,42 +63,31 @@ class CitizenAgent(Agent):
         if self._simulator is None:
             logger.warning("Simulator is not set")
             return
-        if not self._has_bound_to_simulator:
-            FROM_MEMORY_KEYS = {
-                "attribute",
-                "home",
-                "work",
-                "vehicle_attribute",
-                "bus_attribute",
-                "pedestrian_attribute",
-                "bike_attribute",
-            }
-            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)
-                logger.debug(f"Binding to Person `{person_id}` already in Simulator")
-            else:
-                dict_person = deepcopy(self._person_template)
-                for _key in FROM_MEMORY_KEYS:
-                    try:
-                        _value = await memory.get(_key)
-                        if _value:
-                            dict_person[_key] = _value
-                    except KeyError as e:
-                        continue
-                resp = await simulator.add_person(
-                    dict2pb(dict_person, person_pb2.Person())
-                )
-                person_id = resp["person_id"]
-                await memory.update("id", person_id, protect_llm_read_only_fields=False)
-                logger.debug(f"Binding to Person `{person_id}` just added to Simulator")
-                # 防止模拟器还没有到prepare阶段导致get_person出错
-            self._has_bound_to_simulator = True
-            self._agent_id = person_id
-            self.memory.set_agent_id(person_id)
+        FROM_MEMORY_KEYS = {
+            "attribute",
+            "home",
+            "work",
+            "vehicle_attribute",
+            "bus_attribute",
+            "pedestrian_attribute",
+            "bike_attribute",
+        }
+        simulator = self.simulator
+        status = self.status
+        dict_person = deepcopy(self._person_template)
+        for _key in FROM_MEMORY_KEYS:
+            try:
+                _value = await status.get(_key)
+                if _value:
+                    dict_person[_key] = _value
+            except KeyError as e:
+                continue
+        resp = await simulator.add_person(dict2pb(dict_person, person_pb2.Person()))
+        person_id = resp["person_id"]
+        await status.update("id", person_id, protect_llm_read_only_fields=False)
+        logger.debug(f"Binding to Person `{person_id}` just added to Simulator")
+        self._agent_id = person_id
+        self.status.set_agent_id(person_id)
 
     async def _bind_to_economy(self):
         if self._economy_client is None:
@@ -109,8 +99,8 @@ class CitizenAgent(Agent):
                     await self._economy_client.remove_agents([self._agent_id])
                 except:
                     pass
-                person_id = await self.memory.get("id")
-                currency = await self.memory.get("currency")
+                person_id = await self.status.get("id")
+                currency = await self.status.get("currency")
                 await self._economy_client.add_agents(
                     {
                         "id": person_id,
@@ -128,17 +118,32 @@ class CitizenAgent(Agent):
         # 从消息中解析发送者 ID 和消息内容
         target = payload["target"]
         sender_id = payload["from"]
-        content = await self.memory.get(f"{target}")
+        content = await self.status.get(f"{target}")
         payload = {
             "from": self._uuid,
             "content": content,
         }
         await self._send_message(sender_id, payload, "gather")
 
+    @property
+    def mlflow_client(self) -> MlflowClient:
+        """The Agent's MlflowClient"""
+        if self._mlflow_client is None:
+            raise RuntimeError(
+                f"MlflowClient access before assignment, please `set_mlflow_client` first!"
+            )
+        return self._mlflow_client
+
+    def set_mlflow_client(self, mlflow_client: MlflowClient):
+        """
+        Set the mlflow_client of the agent.
+        """
+        self._mlflow_client = mlflow_client
+
 
 class InstitutionAgent(Agent):
     """
-    InstitutionAgent: 机构智能体类及其定义
+    InstitutionAgent: Institution agent class and definition
     """
 
     def __init__(
@@ -146,10 +151,10 @@ class InstitutionAgent(Agent):
         name: str,
         llm_client: Optional[LLM] = None,
         simulator: Optional[Simulator] = None,
-        mlflow_client: Optional[MlflowClient] = None,
         memory: Optional[Memory] = None,
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,  # type:ignore
+        message_interceptor: Optional[MessageInterceptor] = None,  # type:ignore
         avro_file: Optional[dict] = None,
     ) -> None:
         super().__init__(
@@ -157,12 +162,13 @@ class InstitutionAgent(Agent):
             type=AgentType.Institution,
             llm_client=llm_client,
             economy_client=economy_client,
-            mlflow_client=mlflow_client,
             messager=messager,
+            message_interceptor=message_interceptor,
             simulator=simulator,
             memory=memory,
             avro_file=avro_file,
         )
+        self._mlflow_client = None
         # 添加响应收集器
         self._gather_responses: dict[str, asyncio.Future] = {}
 
@@ -178,10 +184,10 @@ class InstitutionAgent(Agent):
             # TODO: More general id generation
             _id = random.randint(100000, 999999)
             self._agent_id = _id
-            self.memory.set_agent_id(_id)
+            self.status.set_agent_id(_id)
             map_header = self.simulator.map.header
             # TODO: remove random position assignment
-            await self.memory.update(
+            await self.status.update(
                 "position",
                 {
                     "xy_position": {
@@ -201,57 +207,57 @@ class InstitutionAgent(Agent):
                 },
                 protect_llm_read_only_fields=False,
             )
-            await self.memory.update("id", _id, protect_llm_read_only_fields=False)
+            await self.status.update("id", _id, protect_llm_read_only_fields=False)
             try:
                 await self._economy_client.remove_orgs([self._agent_id])
             except:
                 pass
             try:
-                _memory = self.memory
-                _id = await _memory.get("id")
-                _type = await _memory.get("type")
+                _status = self.status
+                _id = await _status.get("id")
+                _type = await _status.get("type")
                 try:
-                    nominal_gdp = await _memory.get("nominal_gdp")
+                    nominal_gdp = await _status.get("nominal_gdp")
                 except:
                     nominal_gdp = []
                 try:
-                    real_gdp = await _memory.get("real_gdp")
+                    real_gdp = await _status.get("real_gdp")
                 except:
                     real_gdp = []
                 try:
-                    unemployment = await _memory.get("unemployment")
+                    unemployment = await _status.get("unemployment")
                 except:
                     unemployment = []
                 try:
-                    wages = await _memory.get("wages")
+                    wages = await _status.get("wages")
                 except:
                     wages = []
                 try:
-                    prices = await _memory.get("prices")
+                    prices = await _status.get("prices")
                 except:
                     prices = []
                 try:
-                    inventory = await _memory.get("inventory")
+                    inventory = await _status.get("inventory")
                 except:
                     inventory = 0
                 try:
-                    price = await _memory.get("price")
+                    price = await _status.get("price")
                 except:
                     price = 0
                 try:
-                    currency = await _memory.get("currency")
+                    currency = await _status.get("currency")
                 except:
                     currency = 0.0
                 try:
-                    interest_rate = await _memory.get("interest_rate")
+                    interest_rate = await _status.get("interest_rate")
                 except:
                     interest_rate = 0.0
                 try:
-                    bracket_cutoffs = await _memory.get("bracket_cutoffs")
+                    bracket_cutoffs = await _status.get("bracket_cutoffs")
                 except:
                     bracket_cutoffs = []
                 try:
-                    bracket_rates = await _memory.get("bracket_rates")
+                    bracket_rates = await _status.get("bracket_rates")
                 except:
                     bracket_rates = []
                 await self._economy_client.add_orgs(
@@ -322,3 +328,18 @@ class InstitutionAgent(Agent):
             # 清理Future
             for key in futures:
                 self._gather_responses.pop(key, None)
+
+    @property
+    def mlflow_client(self) -> MlflowClient:
+        """The Agent's MlflowClient"""
+        if self._mlflow_client is None:
+            raise RuntimeError(
+                f"MlflowClient access before assignment, please `set_mlflow_client` first!"
+            )
+        return self._mlflow_client
+
+    def set_mlflow_client(self, mlflow_client: MlflowClient):
+        """
+        Set the mlflow_client of the agent.
+        """
+        self._mlflow_client = mlflow_client

pycityagent/agent/agent_base.py

@@ -13,14 +13,13 @@ from typing import Any, Optional, Union, get_type_hints
 import fastavro
 import ray
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
-from pyparsing import Dict
 
 from ..economy import EconomyClient
 from ..environment import Simulator
 from ..environment.sim.person_service import PersonService
 from ..llm import LLM
 from ..memory import Memory
-from ..message.messager import Messager
+from ..message import MessageInterceptor, Messager
 from ..metrics import MlflowClient
 from ..utils import DIALOG_SCHEMA, SURVEY_SCHEMA, process_survey_for_llm
 from ..workflow import Block
@@ -48,6 +47,7 @@ class Agent(ABC):
 
     configurable_fields: list[str] = []
     default_values: dict[str, Any] = {}
+    fields_description: dict[str, str] = {}
 
     def __init__(
         self,
@@ -55,9 +55,9 @@ class Agent(ABC):
         type: AgentType = AgentType.Unspecified,
         llm_client: Optional[LLM] = None,
         economy_client: Optional[EconomyClient] = None,
-        messager: Optional[Messager] = None,  # type:ignore
+        messager: Optional[ray.ObjectRef] = None,
+        message_interceptor: Optional[ray.ObjectRef] = None,
         simulator: Optional[Simulator] = None,
-        mlflow_client: Optional[MlflowClient] = None,
         memory: Optional[Memory] = None,
         avro_file: Optional[dict[str, str]] = None,
         copy_writer: Optional[ray.ObjectRef] = None,
@@ -72,7 +72,6 @@ class Agent(ABC):
             economy_client (EconomyClient): The `EconomySim` client. Defaults to None.
             messager (Messager, optional): The messager object. Defaults to None.
             simulator (Simulator, optional): The simulator object. Defaults to None.
-            mlflow_client (MlflowClient, optional): The Mlflow 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.
             copy_writer (ray.ObjectRef): The copy_writer of the agent. Defaults to None.
@@ -83,8 +82,8 @@ class Agent(ABC):
         self._llm_client = llm_client
         self._economy_client = economy_client
         self._messager = messager
+        self._message_interceptor = message_interceptor
         self._simulator = simulator
-        self._mlflow_client = mlflow_client
         self._memory = memory
         self._exp_id = -1
         self._agent_id = -1
@@ -104,13 +103,22 @@ class Agent(ABC):
         return state
 
     @classmethod
-    def export_class_config(cls) -> dict[str, Dict]:
-        result = {"agent_name": cls.__name__, "config": {}, "blocks": []}
+    def export_class_config(cls) -> dict[str, dict]:
+        result = {
+            "agent_name": cls.__name__,
+            "config": {},
+            "description": {},
+            "blocks": [],
+        }
         config = {
             field: cls.default_values.get(field, "default_value")
             for field in cls.configurable_fields
         }
         result["config"] = config
+        result["description"] = {
+            field: cls.fields_description.get(field, "")
+            for field in cls.configurable_fields
+        }
         # 解析类中的注解，找到Block类型的字段
         hints = get_type_hints(cls)
         for attr_name, attr_type in hints.items():
@@ -119,14 +127,15 @@ class Agent(ABC):
                 result["blocks"].append(
                     {
                         "name": attr_name,
-                        "config": block_config,
+                        "config": block_config[0],  # type:ignore
+                        "description": block_config[1],  # type:ignore
                         "children": cls._export_subblocks(attr_type),
                     }
                 )
         return result
 
     @classmethod
-    def _export_subblocks(cls, block_cls: type[Block]) -> list[Dict]:
+    def _export_subblocks(cls, block_cls: type[Block]) -> list[dict]:
         children = []
         hints = get_type_hints(block_cls)  # 获取类的注解
         for attr_name, attr_type in hints.items():
@@ -135,7 +144,8 @@ class Agent(ABC):
                 children.append(
                     {
                         "name": attr_name,
-                        "config": block_config,
+                        "config": block_config[0],  # type:ignore
+                        "description": block_config[1],  # type:ignore
                         "children": cls._export_subblocks(attr_type),
                     }
                 )
@@ -216,12 +226,6 @@ class Agent(ABC):
         """
         self._simulator = simulator
 
-    def set_mlflow_client(self, mlflow_client: MlflowClient):
-        """
-        Set the mlflow_client of the agent.
-        """
-        self._mlflow_client = mlflow_client
-
     def set_economy_client(self, economy_client: EconomyClient):
         """
         Set the economy_client of the agent.
@@ -252,6 +256,12 @@ class Agent(ABC):
         """
         self._pgsql_writer = pgsql_writer
 
+    def set_message_interceptor(self, message_interceptor: ray.ObjectRef):
+        """
+        Set the PostgreSQL copy writer of the agent.
+        """
+        self._message_interceptor = message_interceptor
+
     @property
     def uuid(self):
         """The Agent's UUID"""
@@ -279,16 +289,7 @@ class Agent(ABC):
                 f"EconomyClient access before assignment, please `set_economy_client` first!"
             )
         return self._economy_client
-
-    @property
-    def mlflow_client(self):
-        """The Agent's MlflowClient"""
-        if self._mlflow_client is None:
-            raise RuntimeError(
-                f"MlflowClient access before assignment, please `set_mlflow_client` first!"
-            )
-        return self._mlflow_client
-
+    
     @property
     def memory(self):
         """The Agent's Memory"""
@@ -298,6 +299,24 @@ class Agent(ABC):
             )
         return self._memory
 
+    @property
+    def status(self):
+        """The Agent's Status Memory"""
+        if self.memory.status is None:
+            raise RuntimeError(
+                f"Status access before assignment, please `set_memory` first!"
+            )
+        return self.memory.status
+
+    @property
+    def stream(self):
+        """The Agent's Stream Memory"""
+        if self.memory.stream is None:
+            raise RuntimeError(
+                f"Stream access before assignment, please `set_memory` first!"
+            )
+        return self.memory.stream
+
     @property
     def simulator(self):
         """The Simulator"""
@@ -315,11 +334,17 @@ class Agent(ABC):
                 f"Copy Writer access before assignment, please `set_pgsql_writer` first!"
             )
         return self._pgsql_writer
+    
+    @property
+    def messager(self):
+        if self._messager is None:
+            raise RuntimeError("Messager is not set")
+        return self._messager
 
     async def messager_ping(self):
         if self._messager is None:
             raise RuntimeError("Messager is not set")
-        return await self._messager.ping()
+        return await self._messager.ping.remote()  # type:ignore
 
     async def generate_user_survey_response(self, survey: dict) -> str:
         """生成回答 —— 可重写
@@ -338,22 +363,15 @@ class Agent(ABC):
 
         # 添加记忆上下文
         if self._memory:
-            relevant_memories = await self.memory.search(survey_prompt)
-
-            formatted_results = []
-            # for result in top_results:
-            #     formatted_results.append(
-            #         f"- [{result['type']}] {result['content']} "
-            #         f"(相关度: {result['similarity']:.2f})"
-            #     )
-
-            if relevant_memories:
-                dialog.append(
-                    {
-                        "role": "system",
-                        "content": f"Answer based on these memories:\n{relevant_memories}",
-                    }
-                )
+            profile_and_states = await self.status.search(survey_prompt)
+            relevant_activities = await self.stream.search(survey_prompt)
+
+            dialog.append(
+                {
+                    "role": "system",
+                    "content": f"Answer based on following profile and states:\n{profile_and_states}\n Related activities:\n{relevant_activities}",
+                }
+            )
 
         # 添加问卷问题
         dialog.append({"role": "user", "content": survey_prompt})
@@ -412,6 +430,7 @@ class Agent(ABC):
                     _data_tuples
                 )
             )
+        await self.messager.send_message.remote(f"exps/{self._exp_id}/user_payback", {"count": 1})# type:ignore
 
     async def generate_user_chat_response(self, question: str) -> str:
         """生成回答 —— 可重写
@@ -430,14 +449,15 @@ class Agent(ABC):
 
         # 添加记忆上下文
         if self._memory:
-            relevant_memories = await self._memory.search(question)
-            if relevant_memories:
-                dialog.append(
-                    {
-                        "role": "system",
-                        "content": f"Answer based on these memories:\n{relevant_memories}",
-                    }
-                )
+            profile_and_states = await self.status.search(question, top_k=10)
+            relevant_activities = await self.stream.search(question, top_k=10)
+
+            dialog.append(
+                {
+                    "role": "system",
+                    "content": f"Answer based on following profile and states:\n{profile_and_states}\n Related activities:\n{relevant_activities}",
+                }
+            )
 
         # 添加用户问题
         dialog.append({"role": "user", "content": question})
@@ -497,6 +517,8 @@ class Agent(ABC):
                     _data
                 )
             )
+        await self.messager.send_message.remote(f"exps/{self._exp_id}/user_payback", {"count": 1})# type:ignore
+        print(f"Sent payback message to {self._exp_id}")
 
     async def process_agent_chat_response(self, payload: dict) -> str:
         resp = f"Agent {self._uuid} received agent chat response: {payload}"
@@ -566,7 +588,12 @@ class Agent(ABC):
         if self._messager is None:
             raise RuntimeError("Messager is not set")
         topic = f"exps/{self._exp_id}/agents/{to_agent_uuid}/{sub_topic}"
-        await self._messager.send_message.remote(topic, payload)
+        await self._messager.send_message.remote(  # type:ignore
+            topic,
+            payload,
+            self._uuid,
+            to_agent_uuid,
+        )
 
     async def send_message_to_agent(
         self, to_agent_uuid: str, content: str, type: str = "social"
@@ -630,6 +657,6 @@ class Agent(ABC):
         当_blocked为True时，不执行forward方法
         """
         if self._messager is not None:
-            await self._messager.ping.remote()
+            await self._messager.ping.remote()  # type:ignore
         if not self._blocked:
             await self.forward()

pycityagent/cityagent/bankagent.py

@@ -53,14 +53,12 @@ class BankAgent(InstitutionAgent):
 
     async def forward(self):
         if await self.month_trigger():
-            citizens = await self.memory.get("citizens")
-            while True:
-                agents_forward = await self.gather_messages(citizens, "forward")
-                if np.all(np.array(agents_forward) > self.forward_times):
-                    break
-                await asyncio.sleep(1)
+            citizens = await self.memory.status.get("citizens")
+            agents_forward = []
+            if not np.all(np.array(agents_forward) > self.forward_times):
+                return
             self.forward_times += 1
             for uuid in citizens:
                 await self.send_message_to_agent(
-                    uuid, f"bank_forward@{self.forward_times}"
+                    uuid, f"bank_forward@{self.forward_times}", "economy"
                 )

pycityagent/cityagent/blocks/__init__.py

@@ -5,7 +5,6 @@ from .needs_block import NeedsBlock
 from .social_block import SocialBlock
 from .economy_block import EconomyBlock
 from .other_block import OtherBlock
-from .time_block import TimeBlock
 
 __all__ = [
     "MobilityBlock",
@@ -16,5 +15,4 @@ __all__ = [
     "EconomyBlock",
     "OtherBlock",
     "LongTermDecisionBlock",    
-    "TimeBlock",
 ]