pycityagent 2.0.0a52__cp39-cp39-macosx_11_0_arm64.whl → 2.0.0a53__cp39-cp39-macosx_11_0_arm64.whl

pycityagent/agent/agent.py

@@ -14,7 +14,6 @@ from ..environment import Simulator
 from ..llm import LLM
 from ..memory import Memory
 from ..message.messager import Messager
-from ..metrics import MlflowClient
 from .agent_base import Agent, AgentType
 
 logger = logging.getLogger("pycityagent")
@@ -30,7 +29,6 @@ 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
@@ -43,7 +41,6 @@ class CitizenAgent(Agent):
             economy_client=economy_client,
             messager=messager,
             simulator=simulator,
-            mlflow_client=mlflow_client,
             memory=memory,
             avro_file=avro_file,
         )
@@ -62,42 +59,33 @@ 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 +97,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,7 +116,7 @@ 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,
@@ -138,7 +126,7 @@ class CitizenAgent(Agent):
 
 class InstitutionAgent(Agent):
     """
-    InstitutionAgent: 机构智能体类及其定义
+    InstitutionAgent: Institution agent class and definition
     """
 
     def __init__(
@@ -146,7 +134,6 @@ 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
@@ -157,7 +144,6 @@ class InstitutionAgent(Agent):
             type=AgentType.Institution,
             llm_client=llm_client,
             economy_client=economy_client,
-            mlflow_client=mlflow_client,
             messager=messager,
             simulator=simulator,
             memory=memory,
@@ -178,10 +164,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 +187,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(

pycityagent/agent/agent_base.py

@@ -48,6 +48,7 @@ class Agent(ABC):
 
     configurable_fields: list[str] = []
     default_values: dict[str, Any] = {}
+    fields_description: dict[str, str] = {}
 
     def __init__(
         self,
@@ -57,7 +58,6 @@ class Agent(ABC):
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,  # type:ignore
         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.
@@ -84,7 +83,6 @@ class Agent(ABC):
         self._economy_client = economy_client
         self._messager = messager
         self._simulator = simulator
-        self._mlflow_client = mlflow_client
         self._memory = memory
         self._exp_id = -1
         self._agent_id = -1
@@ -104,25 +102,33 @@ 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():
             if inspect.isclass(attr_type) and issubclass(attr_type, Block):
                 block_config = attr_type.export_class_config()
-                result["blocks"].append(
-                    {
-                        "name": attr_name,
-                        "config": block_config,
-                        "children": cls._export_subblocks(attr_type),
-                    }
-                )
+                result["blocks"].append({
+                    "name": attr_name,
+                    "config": block_config[0],
+                    "description": block_config[1],
+                    "children": cls._export_subblocks(attr_type)
+                })
         return result
 
     @classmethod
@@ -135,7 +141,8 @@ class Agent(ABC):
                 children.append(
                     {
                         "name": attr_name,
-                        "config": block_config,
+                        "config": block_config[0],
+                    "description": block_config[1],
                         "children": cls._export_subblocks(attr_type),
                     }
                 )
@@ -216,12 +223,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.
@@ -279,16 +280,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"""
@@ -297,6 +289,24 @@ class Agent(ABC):
                 f"Memory access before assignment, please `set_memory` first!"
             )
         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):
@@ -315,6 +325,12 @@ 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:
@@ -338,22 +354,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 +421,7 @@ class Agent(ABC):
                     _data_tuples
                 )
             )
+        await self.messager.send_message.remote(f"exps/{self._exp_id}/user_payback", {"count": 1})
 
     async def generate_user_chat_response(self, question: str) -> str:
         """生成回答 —— 可重写
@@ -430,14 +440,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 +508,8 @@ class Agent(ABC):
                     _data
                 )
             )
+        await self.messager.send_message.remote(f"exps/{self._exp_id}/user_payback", {"count": 1})
+        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}"

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",
 ]