pycityagent 2.0.0a49__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a51__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/{agent.py → agent/agent_base.py}

@@ -1,33 +1,29 @@
 from __future__ import annotations
+
 import asyncio
 import inspect
 import json
 import logging
-import random
 import uuid
 from abc import ABC, abstractmethod
-from copy import deepcopy
 from datetime import datetime, timezone
 from enum import Enum
-from typing import Any,  Optional, Type, get_type_hints
-from uuid import UUID
+from typing import Any, Optional, Union, get_type_hints
 
 import fastavro
-from pyparsing import Dict
 import ray
-from mosstool.util.format_converter import dict2pb
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
+from pyparsing import Dict
 
-from pycityagent.workflow import Block
-
-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 .metrics import MlflowClient
-from .utils import DIALOG_SCHEMA, SURVEY_SCHEMA, process_survey_for_llm
+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 ..metrics import MlflowClient
+from ..utils import DIALOG_SCHEMA, SURVEY_SCHEMA, process_survey_for_llm
+from ..workflow import Block
 
 logger = logging.getLogger("pycityagent")
 
@@ -49,6 +45,7 @@ class Agent(ABC):
     """
     Agent base class
     """
+
     configurable_fields: list[str] = []
     default_values: dict[str, Any] = {}
 
@@ -108,11 +105,7 @@ class Agent(ABC):
 
     @classmethod
     def export_class_config(cls) -> dict[str, Dict]:
-        result = {
-            "agent_name": cls.__name__,
-            "config": {},
-            "blocks": []
-        }
+        result = {"agent_name": cls.__name__, "config": {}, "blocks": []}
         config = {
             field: cls.default_values.get(field, "default_value")
             for field in cls.configurable_fields
@@ -123,25 +116,29 @@ class Agent(ABC):
         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,
+                        "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():
             if inspect.isclass(attr_type) and issubclass(attr_type, Block):
                 block_config = attr_type.export_class_config()
-                children.append({
-                    "name": attr_name,
-                    "config": block_config,
-                    "children": cls._export_subblocks(attr_type)
-                })
+                children.append(
+                    {
+                        "name": attr_name,
+                        "config": block_config,
+                        "children": cls._export_subblocks(attr_type),
+                    }
+                )
         return children
 
     @classmethod
@@ -151,28 +148,29 @@ class Agent(ABC):
             json.dump(config, f, indent=4)
 
     @classmethod
-    def import_block_config(cls, config: dict[str, list[Dict]]) -> "Agent":
-        agent = cls(name=config["agent_name"])
+    def import_block_config(cls, config: dict[str, Union[list[dict], str]]) -> Agent:
+        agent = cls(name=config["agent_name"])  # type:ignore
 
-        def build_block(block_data: Dict) -> Block:
+        def build_block(block_data: dict[str, Any]) -> Block:
             block_cls = globals()[block_data["name"]]
             block_instance = block_cls.import_config(block_data)
             return block_instance
 
         # 创建顶层Block
         for block_data in config["blocks"]:
+            assert isinstance(block_data, dict)
             block = build_block(block_data)
             setattr(agent, block.name.lower(), block)
 
         return agent
 
     @classmethod
-    def import_from_file(cls, filepath: str) -> "Agent":
+    def import_from_file(cls, filepath: str) -> Agent:
         with open(filepath, "r") as f:
             config = json.load(f)
             return cls.import_block_config(config)
-        
-    def load_from_config(self, config: dict[str, list[Dict]]) -> None:
+
+    def load_from_config(self, config: dict[str, list[dict]]) -> None:
         """
         使用配置更新当前Agent实例的Block层次结构。
         """
@@ -185,13 +183,15 @@ class Agent(ABC):
         # 递归更新或创建顶层Block
         for block_data in config.get("blocks", []):
             block_name = block_data["name"]
-            existing_block = getattr(self, block_name, None) # type:ignore
+            existing_block = getattr(self, block_name, None)  # type:ignore
 
             if existing_block:
                 # 如果Block已经存在，则递归更新
                 existing_block.load_from_config(block_data)
             else:
-                raise KeyError(f"Block '{block_name}' not found in agent '{self.__class__.__name__}'")
+                raise KeyError(
+                    f"Block '{block_name}' not found in agent '{self.__class__.__name__}'"
+                )
 
     def load_from_file(self, filepath: str) -> None:
         with open(filepath, "r") as f:
@@ -315,7 +315,7 @@ class Agent(ABC):
                 f"Copy Writer access before assignment, please `set_pgsql_writer` first!"
             )
         return self._pgsql_writer
-    
+
     async def messager_ping(self):
         if self._messager is None:
             raise RuntimeError("Messager is not set")
@@ -633,307 +633,3 @@ class Agent(ABC):
             await self._messager.ping.remote()
         if not self._blocked:
             await self.forward()
-
-
-class CitizenAgent(Agent):
-    """
-    CitizenAgent: 城市居民智能体类及其定义
-    """
-
-    def __init__(
-        self,
-        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
-        avro_file: Optional[dict] = None,
-    ) -> None:
-        super().__init__(
-            name=name,
-            type=AgentType.Citizen,
-            llm_client=llm_client,
-            economy_client=economy_client,
-            messager=messager,
-            simulator=simulator,
-            mlflow_client=mlflow_client,
-            memory=memory,
-            avro_file=avro_file,
-        )
-
-    async def bind_to_simulator(self):
-        await self._bind_to_simulator()
-        await self._bind_to_economy()
-
-    async def _bind_to_simulator(self):
-        """
-        Bind Agent to Simulator
-
-        Args:
-            person_template (dict, optional): The person template in dict format. Defaults to PersonService.default_dict_person().
-        """
-        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)
-
-    async def _bind_to_economy(self):
-        if self._economy_client is None:
-            logger.warning("Economy client is not set")
-            return
-        if not self._has_bound_to_economy:
-            if self._has_bound_to_simulator:
-                try:
-                    await self._economy_client.remove_agents([self._agent_id])
-                except:
-                    pass
-                person_id = await self.memory.get("id")
-                currency = await self.memory.get("currency")
-                await self._economy_client.add_agents(
-                    {
-                        "id": person_id,
-                        "currency": currency,
-                    }
-                )
-                self._has_bound_to_economy = True
-            else:
-                logger.debug(
-                    f"Binding to Economy before binding to Simulator, skip binding to Economy Simulator"
-                )
-
-    async def handle_gather_message(self, payload: dict):
-        """处理收到的消息，识别发送者"""
-        # 从消息中解析发送者 ID 和消息内容
-        target = payload["target"]
-        sender_id = payload["from"]
-        content = await self.memory.get(f"{target}")
-        payload = {
-            "from": self._uuid,
-            "content": content,
-        }
-        await self._send_message(sender_id, payload, "gather")
-
-
-class InstitutionAgent(Agent):
-    """
-    InstitutionAgent: 机构智能体类及其定义
-    """
-
-    def __init__(
-        self,
-        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
-        avro_file: Optional[dict] = None,
-    ) -> None:
-        super().__init__(
-            name=name,
-            type=AgentType.Institution,
-            llm_client=llm_client,
-            economy_client=economy_client,
-            mlflow_client=mlflow_client,
-            messager=messager,
-            simulator=simulator,
-            memory=memory,
-            avro_file=avro_file,
-        )
-        # 添加响应收集器
-        self._gather_responses: dict[str, asyncio.Future] = {}
-
-    async def bind_to_simulator(self):
-        await self._bind_to_economy()
-
-    async def _bind_to_economy(self):
-        print("Debug:", self._economy_client, self._has_bound_to_economy)
-        if self._economy_client is None:
-            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
-            self.memory.set_agent_id(_id)
-            map_header = self.simulator.map.header
-            # TODO: remove random position assignment
-            await self.memory.update(
-                "position",
-                {
-                    "xy_position": {
-                        "x": float(
-                            random.randrange(
-                                start=int(map_header["west"]),
-                                stop=int(map_header["east"]),
-                            )
-                        ),
-                        "y": float(
-                            random.randrange(
-                                start=int(map_header["south"]),
-                                stop=int(map_header["north"]),
-                            )
-                        ),
-                    }
-                },
-                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:
-                _memory = self.memory
-                _id = await _memory.get("id")
-                _type = await _memory.get("type")
-                try:
-                    nominal_gdp = await _memory.get("nominal_gdp")
-                except:
-                    nominal_gdp = []
-                try:
-                    real_gdp = await _memory.get("real_gdp")
-                except:
-                    real_gdp = []
-                try:
-                    unemployment = await _memory.get("unemployment")
-                except:
-                    unemployment = []
-                try:
-                    wages = await _memory.get("wages")
-                except:
-                    wages = []
-                try:
-                    prices = await _memory.get("prices")
-                except:
-                    prices = []
-                try:
-                    inventory = await _memory.get("inventory")
-                except:
-                    inventory = 0
-                try:
-                    price = await _memory.get("price")
-                except:
-                    price = 0
-                try:
-                    currency = await _memory.get("currency")
-                except:
-                    currency = 0.0
-                try:
-                    interest_rate = await _memory.get("interest_rate")
-                except:
-                    interest_rate = 0.0
-                try:
-                    bracket_cutoffs = await _memory.get("bracket_cutoffs")
-                except:
-                    bracket_cutoffs = []
-                try:
-                    bracket_rates = await _memory.get("bracket_rates")
-                except:
-                    bracket_rates = []
-                await self._economy_client.add_orgs(
-                    {
-                        "id": _id,
-                        "type": _type,
-                        "nominal_gdp": nominal_gdp,
-                        "real_gdp": real_gdp,
-                        "unemployment": unemployment,
-                        "wages": wages,
-                        "prices": prices,
-                        "inventory": inventory,
-                        "price": price,
-                        "currency": currency,
-                        "interest_rate": interest_rate,
-                        "bracket_cutoffs": bracket_cutoffs,
-                        "bracket_rates": bracket_rates,
-                    }
-                )
-            except Exception as e:
-                logger.error(f"Failed to bind to Economy: {e}")
-            self._has_bound_to_economy = True
-
-    async def handle_gather_message(self, payload: dict):
-        """处理收到的消息，识别发送者"""
-        content = payload["content"]
-        sender_id = payload["from"]
-
-        # 将响应存储到对应的Future中
-        response_key = str(sender_id)
-        if response_key in self._gather_responses:
-            self._gather_responses[response_key].set_result(
-                {
-                    "from": sender_id,
-                    "content": content,
-                }
-            )
-
-    async def gather_messages(self, agent_uuids: list[str], target: str) -> list[dict]:
-        """从多个智能体收集消息
-
-        Args:
-            agent_uuids: 目标智能体UUID列表
-            target: 要收集的信息类型
-
-        Returns:
-            list[dict]: 收集到的所有响应
-        """
-        # 为每个agent创建Future
-        futures = {}
-        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_uuid in agent_uuids:
-            await self._send_message(agent_uuid, payload, "gather")
-
-        try:
-            # 等待所有响应
-            responses = await asyncio.gather(*futures.values())
-            return responses
-        finally:
-            # 清理Future
-            for key in futures:
-                self._gather_responses.pop(key, None)

pycityagent/cityagent/bankagent.py

@@ -11,22 +11,32 @@ import logging
 
 logger = logging.getLogger("pycityagent")
 
-class BankAgent(InstitutionAgent):    
-    def __init__(self, 
-                 name: str,
-                 llm_client: Optional[LLM] = None,
-                 simulator: Optional[Simulator] = None,
-                 memory: Optional[Memory] = None,
-                 economy_client: Optional[EconomyClient] = None,
-                 messager: Optional[Messager] = None,
-                 avro_file: Optional[dict] = None,
-                 ) -> None:
-        super().__init__(name=name, llm_client=llm_client, simulator=simulator, memory=memory, economy_client=economy_client, messager=messager, avro_file=avro_file)
+
+class BankAgent(InstitutionAgent):
+    def __init__(
+        self,
+        name: str,
+        llm_client: Optional[LLM] = None,
+        simulator: Optional[Simulator] = None,
+        memory: Optional[Memory] = None,
+        economy_client: Optional[EconomyClient] = None,
+        messager: Optional[Messager] = None,
+        avro_file: Optional[dict] = None,
+    ) -> None:
+        super().__init__(
+            name=name,
+            llm_client=llm_client,
+            simulator=simulator,
+            memory=memory,
+            economy_client=economy_client,
+            messager=messager,
+            avro_file=avro_file,
+        )
         self.initailzed = False
         self.last_time_trigger = None
         self.time_diff = 30 * 24 * 60 * 60
         self.forward_times = 0
-        
+
     async def month_trigger(self):
         now_time = await self.simulator.get_time()
         if self.last_time_trigger is None:
@@ -36,19 +46,21 @@ class BankAgent(InstitutionAgent):
             self.last_time_trigger = now_time
             return True
         return False
-    
+
     async def gather_messages(self, agent_ids, content):
         infos = await super().gather_messages(agent_ids, content)
-        return [info['content'] for info in infos]
+        return [info["content"] for info in infos]
 
     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')
+                agents_forward = await self.gather_messages(citizens, "forward")
                 if np.all(np.array(agents_forward) > self.forward_times):
                     break
                 await asyncio.sleep(1)
             self.forward_times += 1
             for uuid in citizens:
-                await self.send_message_to_agent(uuid, f"bank_forward@{self.forward_times}")
+                await self.send_message_to_agent(
+                    uuid, f"bank_forward@{self.forward_times}"
+                )

pycityagent/cityagent/firmagent.py

@@ -11,7 +11,8 @@ import logging
 
 logger = logging.getLogger("pycityagent")
 
-class FirmAgent(InstitutionAgent):    
+
+class FirmAgent(InstitutionAgent):
     configurable_fields = ["time_diff", "max_price_inflation", "max_wage_inflation"]
     default_values = {
         "time_diff": 30 * 24 * 60 * 60,
@@ -19,16 +20,25 @@ class FirmAgent(InstitutionAgent):
         "max_wage_inflation": 0.05,
     }
 
-    def __init__(self, 
-                 name: str,
-                 llm_client: Optional[LLM] = None,
-                 simulator: Optional[Simulator] = None,
-                 memory: Optional[Memory] = None,
-                 economy_client: Optional[EconomyClient] = None,
-                 messager: Optional[Messager] = None,
-                 avro_file: Optional[dict] = None,
-                 ) -> None:
-        super().__init__(name=name, llm_client=llm_client, simulator=simulator, memory=memory, economy_client=economy_client, messager=messager, avro_file=avro_file)
+    def __init__(
+        self,
+        name: str,
+        llm_client: Optional[LLM] = None,
+        simulator: Optional[Simulator] = None,
+        memory: Optional[Memory] = None,
+        economy_client: Optional[EconomyClient] = None,
+        messager: Optional[Messager] = None,
+        avro_file: Optional[dict] = None,
+    ) -> None:
+        super().__init__(
+            name=name,
+            llm_client=llm_client,
+            simulator=simulator,
+            memory=memory,
+            economy_client=economy_client,
+            messager=messager,
+            avro_file=avro_file,
+        )
         self.initailzed = False
         self.last_time_trigger = None
         self.forward_times = 0
@@ -45,31 +55,59 @@ class FirmAgent(InstitutionAgent):
             self.last_time_trigger = now_time
             return True
         return False
-    
+
     async def gather_messages(self, agent_ids, content):
         infos = await super().gather_messages(agent_ids, content)
-        return [info['content'] for info in infos]
+        return [info["content"] for info in infos]
 
     async def forward(self):
         if await self.month_trigger():
             employees = await self.memory.get("employees")
             while True:
-                agents_forward = await self.gather_messages(employees, 'forward')
+                agents_forward = await self.gather_messages(employees, "forward")
                 if np.all(np.array(agents_forward) > self.forward_times):
                     break
                 await asyncio.sleep(1)
-            goods_demand = await self.gather_messages(employees, 'goods_demand')
-            goods_consumption = await self.gather_messages(employees, 'goods_consumption')
-            print(f'goods_demand: {goods_demand}, goods_consumption: {goods_consumption}')
+            goods_demand = await self.gather_messages(employees, "goods_demand")
+            goods_consumption = await self.gather_messages(
+                employees, "goods_consumption"
+            )
+            print(
+                f"goods_demand: {goods_demand}, goods_consumption: {goods_consumption}"
+            )
             total_demand = sum(goods_demand)
-            last_inventory = sum(goods_consumption) + await self.economy_client.get(self._agent_id, "inventory")
-            print(f'total_demand: {total_demand}, last_inventory: {last_inventory}, goods_contumption: {sum(goods_consumption)}')
-            max_change_rate = (total_demand - last_inventory)/(max(total_demand, last_inventory)+1e-8)
-            skills = await self.gather_messages(employees, 'work_skill')
+            last_inventory = sum(goods_consumption) + await self.economy_client.get(
+                self._agent_id, "inventory"
+            )
+            print(
+                f"total_demand: {total_demand}, last_inventory: {last_inventory}, goods_contumption: {sum(goods_consumption)}"
+            )
+            max_change_rate = (total_demand - last_inventory) / (
+                max(total_demand, last_inventory) + 1e-8
+            )
+            skills = await self.gather_messages(employees, "work_skill")
             for skill, uuid in zip(skills, employees):
-                await self.send_message_to_agent(uuid, f"work_skill@{max(skill*(1 + np.random.uniform(0, max_change_rate*self.max_wage_inflation)), 1)}")
-            price = await self.economy_client.get(self._agent_id, 'price')
-            await self.economy_client.update(self._agent_id, 'price', max(price*(1 + np.random.uniform(0, max_change_rate*self.max_price_inflation)), 1))
+                await self.send_message_to_agent(
+                    uuid,
+                    f"work_skill@{max(skill*(1 + np.random.uniform(0, max_change_rate*self.max_wage_inflation)), 1)}",
+                )
+            price = await self.economy_client.get(self._agent_id, "price")
+            await self.economy_client.update(
+                self._agent_id,
+                "price",
+                max(
+                    price
+                    * (
+                        1
+                        + np.random.uniform(
+                            0, max_change_rate * self.max_price_inflation
+                        )
+                    ),
+                    1,
+                ),
+            )
             self.forward_times += 1
             for uuid in employees:
-                await self.send_message_to_agent(uuid, f"firm_forward@{self.forward_times}")
+                await self.send_message_to_agent(
+                    uuid, f"firm_forward@{self.forward_times}"
+                )