pycityagent 1.0.0__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

pycityagent/economy/econ_client.py

@@ -0,0 +1,307 @@
+import asyncio
+import logging
+from typing import Any, Literal, Union
+
+import grpc
+import pycityproto.city.economy.v2.economy_pb2 as economyv2
+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
+
+__all__ = [
+    "EconomyClient",
+]
+
+
+def _snake_to_pascal(snake_str):
+    _res = "".join(word.capitalize() or "_" for word in snake_str.split("_"))
+    for _word in {
+        "Gdp",
+    }:
+        if _word in _res:
+            _res = _res.replace(_word, _word.upper())
+    return _res
+
+
+def _get_field_type_and_repeated(message, field_name: str) -> tuple[Any, bool]:
+    try:
+        field_descriptor = message.DESCRIPTOR.fields_by_name[field_name]
+        field_type = field_descriptor.type
+        _type_mapping = {
+            descriptor.FieldDescriptor.TYPE_FLOAT: float,
+            descriptor.FieldDescriptor.TYPE_INT32: int,
+        }
+        is_repeated = (
+            field_descriptor.label == descriptor.FieldDescriptor.LABEL_REPEATED
+        )
+        return (_type_mapping.get(field_type), is_repeated)
+    except KeyError:
+        raise KeyError(f"Invalid message {message} and filed name {field_name}!")
+
+
+def _create_aio_channel(server_address: str, secure: bool = False) -> grpc.aio.Channel:
+    """
+    Create a grpc asynchronous channel
+
+    Args:
+    - server_address (str): server address.
+    - secure (bool, optional): Defaults to False. Whether to use a secure connection. Defaults to False.
+
+    Returns:
+    - grpc.aio.Channel: grpc asynchronous channel.
+    """
+    if server_address.startswith("http://"):
+        server_address = server_address.split("//")[1]
+        if secure:
+            raise ValueError("secure channel must use `https` or not use `http`")
+    elif server_address.startswith("https://"):
+        server_address = server_address.split("//")[1]
+        if not secure:
+            secure = True
+
+    if secure:
+        return grpc.aio.secure_channel(server_address, grpc.ssl_channel_credentials())
+    else:
+        return grpc.aio.insecure_channel(server_address)
+
+
+class EconomyClient:
+    """
+    Client side of Economy service
+    """
+
+    def __init__(self, server_address: str, secure: bool = False):
+        """
+        Constructor of EconomyClient
+
+        Args:
+        - server_address (str): Economy server address
+        - secure (bool, optional): Defaults to False. Whether to use a secure connection. Defaults to False.
+        """
+        self.server_address = server_address
+        self.secure = secure
+        aio_channel = _create_aio_channel(server_address, secure)
+        self._aio_stub = org_grpc.OrgServiceStub(aio_channel)
+
+    def __getstate__(self):
+        """
+        Copy the object's state from self.__dict__ which contains
+        all our instance attributes. Always use the dict.copy()
+        method to avoid modifying the original state.
+        """
+        state = self.__dict__.copy()
+        # Remove the non-picklable entries.
+        del state["_aio_stub"]
+        return state
+
+    def __setstate__(self, state):
+        """ "
+        Restore instance attributes (i.e., filename and mode) from the
+        unpickled state dictionary.
+        """
+        self.__dict__.update(state)
+        # Re-initialize the channel after unpickling
+        aio_channel = _create_aio_channel(self.server_address, self.secure)
+        self._aio_stub = org_grpc.OrgServiceStub(aio_channel)
+
+    async def get(
+        self,
+        id: int,
+        key: str,
+    ) -> Any:
+        """
+        Get specific value
+
+        Args:
+        - id (int): the id of `Org` or `Agent`.
+        - key (str): the attribute to fetch.
+
+        Returns:
+        - Any
+        """
+        pascal_key = _snake_to_pascal(key)
+        _request_type = getattr(org_service, f"Get{pascal_key}Request")
+        _request_func = getattr(self._aio_stub, f"Get{pascal_key}")
+        response = await _request_func(_request_type(org_id=id))
+        value_type, is_repeated = _get_field_type_and_repeated(response, field_name=key)
+        if is_repeated:
+            return list(getattr(response, key))
+        else:
+            return value_type(getattr(response, key))
+
+    async def update(
+        self,
+        id: int,
+        key: str,
+        value: Any,
+        mode: Union[Literal["replace"], Literal["merge"]] = "replace",
+    ) -> Any:
+        """
+        Update key-value pair
+
+        Args:
+        - id (int): the id of `Org` or `Agent`.
+        - key (str): the attribute to update.
+        - mode (Union[Literal["replace"], Literal["merge"]], optional): Update mode. Defaults to "replace".
+
+
+        Returns:
+        - Any
+        """
+        pascal_key = _snake_to_pascal(key)
+        _request_type = getattr(org_service, f"Set{pascal_key}Request")
+        _request_func = getattr(self._aio_stub, f"Set{pascal_key}")
+        if mode == "merge":
+            orig_value = await self.get(id, key)
+            _orig_type = type(orig_value)
+            _new_type = type(value)
+            if _orig_type != _new_type:
+                logging.debug(
+                    f"Inconsistent type of original value {_orig_type.__name__} and to-update value {_new_type.__name__}"
+                )
+            else:
+                if isinstance(orig_value, set):
+                    orig_value.update(set(value))
+                    value = orig_value
+                elif isinstance(orig_value, dict):
+                    orig_value.update(dict(value))
+                    value = orig_value
+                elif isinstance(orig_value, list):
+                    orig_value.extend(list(value))
+                    value = orig_value
+                else:
+                    logging.warning(
+                        f"Type of {type(orig_value)} does not support mode `merge`, using `replace` instead!"
+                    )
+        return await _request_func(
+            _request_type(
+                **{
+                    "org_id": id,
+                    key: value,
+                }
+            )
+        )
+
+    async def add_agents(self, configs: Union[list[dict], dict]):
+        if isinstance(configs, dict):
+            configs = [configs]
+        tasks = [
+            self._aio_stub.AddAgent(
+                org_service.AddAgentRequest(
+                    agent=economyv2.Agent(
+                        id=config["id"],
+                        currency=config.get("currency", 0.0),
+                    )
+                )
+            )
+            for config in configs
+        ]
+        responses = await asyncio.gather(*tasks)
+
+    async def add_orgs(self, configs: Union[list[dict], dict]):
+        if isinstance(configs, dict):
+            configs = [configs]
+        tasks = [
+            self._aio_stub.AddOrg(
+                org_service.AddOrgRequest(
+                    org=economyv2.Org(
+                        id=config["id"],
+                        type=config["type"],
+                        nominal_gdp=config.get("nominal_gdp", []),
+                        real_gdp=config.get("real_gdp", []),
+                        unemployment=config.get("unemployment", []),
+                        wages=config.get("wages", []),
+                        prices=config.get("prices", []),
+                        inventory=config.get("inventory", 0),
+                        price=config.get("price", 0),
+                        currency=config.get("currency", 0.0),
+                        interest_rate=config.get("interest_rate", 0.0),
+                        bracket_cutoffs=config.get("bracket_cutoffs", []),
+                        bracket_rates=config.get("bracket_rates", []),
+                    )
+                )
+            )
+            for config in configs
+        ]
+        responses = await asyncio.gather(*tasks)
+
+    async def calculate_taxes_due(
+        self,
+        org_id: int,
+        agent_ids: list[int],
+        incomes: list[float],
+        enable_redistribution: bool,
+    ):
+        request = org_service.CalculateTaxesDueRequest(
+            government_id=org_id,
+            agent_ids=agent_ids,
+            incomes=incomes,
+            enable_redistribution=enable_redistribution,
+        )
+        response: org_service.CalculateTaxesDueResponse = (
+            await self._aio_stub.CalculateTaxesDue(request)
+        )
+        return (float(response.taxes_due), list(response.updated_incomes))
+
+    async def calculate_consumption(
+        self, org_id: int, agent_ids: list[int], demands: list[int]
+    ):
+        request = org_service.CalculateConsumptionRequest(
+            firm_id=org_id,
+            agent_ids=agent_ids,
+            demands=demands,
+        )
+        response: org_service.CalculateConsumptionResponse = (
+            await self._aio_stub.CalculateConsumption(request)
+        )
+        return (int(response.remain_inventory), list(response.updated_currencies))
+
+    async def calculate_interest(self, org_id: int, agent_ids: list[int]):
+        request = org_service.CalculateInterestRequest(
+            bank_id=org_id,
+            agent_ids=agent_ids,
+        )
+        response: org_service.CalculateInterestResponse = (
+            await self._aio_stub.CalculateInterest(request)
+        )
+        return (float(response.total_interest), list(response.updated_currencies))
+
+    async def remove_agents(self, agent_ids: Union[int, list[int]]):
+        if isinstance(agent_ids, int):
+            agent_ids = [agent_ids]
+        tasks = [
+            self._aio_stub.RemoveAgent(
+                org_service.RemoveAgentRequest(agent_id=agent_id)
+            )
+            for agent_id in agent_ids
+        ]
+        responses = await asyncio.gather(*tasks)
+
+    async def remove_orgs(self, org_ids: Union[int, list[int]]):
+        if isinstance(org_ids, int):
+            org_ids = [org_ids]
+        tasks = [
+            self._aio_stub.RemoveOrg(org_service.RemoveOrgRequest(org_id=org_id))
+            for org_id in org_ids
+        ]
+        responses = await asyncio.gather(*tasks)
+
+    async def save(self, file_path: str) -> tuple[list[int], list[int]]:
+        request = org_service.SaveEconomyEntitiesRequest(
+            file_path=file_path,
+        )
+        response: org_service.SaveEconomyEntitiesResponse = (
+            await self._aio_stub.SaveEconomyEntities(request)
+        )
+        # current agent ids and org ids
+        return (list(response.agent_ids), list(response.org_ids))
+
+    async def load(self, file_path: str):
+        request = org_service.LoadEconomyEntitiesRequest(
+            file_path=file_path,
+        )
+        response: org_service.LoadEconomyEntitiesResponse = (
+            await self._aio_stub.LoadEconomyEntities(request)
+        )
+        # current agent ids and org ids
+        return (list(response.agent_ids), list(response.org_ids))

pycityagent/environment/__init__.py

@@ -0,0 +1,7 @@
+"""Environment"""
+
+from .sence.static import LEVEL_ONE_PRE, POI_TYPE_DICT
+from .sim import AoiService, PersonService
+from .simulator import Simulator
+
+__all__ = ["Simulator", "POI_TYPE_DICT", "LEVEL_ONE_PRE", "PersonService", "AoiService"]

pycityagent/environment/interact/interact.py

@@ -0,0 +1,141 @@
+"""环境相关的Interaction定义"""
+from enum import Enum
+from typing import Callable, Optional, Any
+from abc import ABC, abstractmethod
+from typing import Callable, Any
+
+class ActionType(Enum):
+    """
+    行动类型枚举 所有行动本质上为数据推送
+    Action Type enumeration, all actions are essentially data push
+
+    Types:
+    - Sim = 1, 用于表示与模拟器对接的行动
+    - Hub = 2, 用于表示与AppHub(前端)对接的行动
+    - Comp = 3, 表示综合类型 (可能同时包含与Sim以及Hub的交互)
+    """
+    Sim = 1
+    Hub = 2
+    Comp = 3
+    
+class Action:
+    """
+    - Action
+    """
+    def __init__(self, agent, type:ActionType, source: Optional[str] = None, before:Optional[Callable[[list], Any]] = None) -> None:
+        '''
+        默认初始化
+        
+        Args:
+        - agent (Agent): the related agent
+        - type (ActionType)
+        - source (str): 数据来源, 默认为None, 如果为None则会从接收用户传入的数据作为Forward函数参数, 否则从WM.Reason数据缓存中取对应数据作为参数
+        - before (function): 数据处理方法, 用于当Reason缓存中的参数与标准格式不符时使用
+        '''
+        self._agent = agent
+        self._type = type
+        self._source = source
+        self._before = before
+
+    def get_source(self):
+        """
+        获取source数据
+        """
+        if self._source != None:
+            source = self._agent.Brain.Memory.Working.Reason[self._source]
+            if self._before != None:
+                source = self._before(source)
+            return source
+        else:
+            return None
+
+    @abstractmethod
+    async def Forward(self):
+        '''接口函数'''
+
+class SimAction(Action):
+    """SimAction: 模拟器关联Action"""
+    def __init__(self, agent, source: Optional[str] = None, before:Optional[Callable[[list], Any]] = None) -> None:
+        super().__init__(agent, ActionType.Sim, source, before)
+
+class HubAction(Action):
+    """HubAction: Apphub关联Action"""
+    def __init__(self, agent, source: Optional[str] = None, before:Optional[Callable[[list], Any]] = None) -> None:
+        super().__init__(agent, ActionType.Hub, source, before)
+
+
+class SetSchedule(SimAction):
+    """
+    用于将agent的行程信息同步至模拟器 —— 仅对citizen类型agent适用
+    Synchronize agent's schedule to simulator —— only avalable for citizen type of agent
+    """
+    def __init__(self, agent, source: Optional[str] = None, before:Optional[Callable[[list], Any]] = None) -> None:
+        super().__init__(agent, source, before)
+
+    async def Forward(self, schedule = None):
+        """
+        如果当前行程已经同步至模拟器: 跳过同步, 否则同步至模拟器
+        If current schedule has been synchronized to simulator: skip, else sync
+        """
+        if not schedule == None:
+            if not schedule.is_set:
+                '''同步schedule至模拟器'''
+                self._agent.Scheduler.now.is_set = True
+                departure_time = schedule.time
+                mode = schedule.mode
+                aoi_id = schedule.target_id_aoi
+                poi_id = schedule.target_id_poi
+                end = {'aoi_position': {'aoi_id': aoi_id, 'poi_id': poi_id}}
+                activity = schedule.description
+                trips = [{'mode': mode, 'end': end, 'departure_time': departure_time, 'activity': activity}]
+                set_schedule = [{'trips': trips, 'loop_count': 1, 'departure_time': departure_time}]
+
+                # * 与模拟器对接
+                req = {'person_id': self._agent._id, 'schedules': set_schedule}
+                await self._agent._client.person_service.SetSchedule(req)
+        elif self._source != None:
+            schedule = self.get_source()
+            if schedule != None and not schedule.is_set:
+                '''同步schedule至模拟器'''
+                self._agent.Scheduler.now.is_set = True
+                departure_time = schedule.time
+                mode = schedule.mode
+                aoi_id = schedule.target_id_aoi
+                poi_id = schedule.target_id_poi
+                end = {'aoi_position': {'aoi_id': aoi_id, 'poi_id': poi_id}}
+                activity = schedule.description
+                trips = [{'mode': mode, 'end': end, 'departure_time': departure_time, 'activity': activity}]
+                set_schedule = [{'trips': trips, 'loop_count': 1, 'departure_time': departure_time}]
+
+                # * 与模拟器对接
+                req = {'person_id': self._agent._id, 'schedules': set_schedule}
+                await self._agent._client.person_service.SetSchedule(req)
+
+
+class SendAgentMessage(SimAction):
+    """
+    发送信息给其他agent
+    Send messages to other agents
+    """
+    def __init__(self, agent, source: Optional[str] = None, before:Optional[Callable[[list], Any]] = None) -> None:
+        super().__init__(agent, source, before)
+
+    async def Forward(self, messages: Optional[dict] = None):
+        if not messages == None and len(messages) > 0:
+            req = {'messages': []}
+            for message in messages:
+                from_id = self._agent._id
+                to_id = message['id']
+                mes = message['message']
+                req['messages'].append({'from': from_id, 'to': to_id, 'message': mes})
+            await self._agent._client.social_service.Send(req)
+        elif self._source != None:
+            messages = self.get_source()
+            if not messages == None and len(messages) > 0:
+                req = {'messages': []}
+                for message in messages:
+                    from_id = self._agent._id
+                    to_id = message['id']
+                    mes = message['message']
+                    req['messages'].append({'from': from_id, 'to': to_id, 'message': mes})
+                await self._agent._client.social_service.Send(req)

pycityagent/{brain → environment/sence}/static.py

@@ -1,5 +1,5 @@
 """
-Static Resources: Poi Type association; Type prefix.
+静态数据支持; Static Resources: Poi Type association; Type prefix.
 """
 POI_TYPE_DICT = {
     "100000": "\u7f8e\u98df",

pycityagent/environment/sidecar/__init__.py

@@ -0,0 +1,8 @@
+"""
+用于接入模拟器同步框架的Sidecar框架服务（仅客户端）
+Sidecar framework service for accessing the simulator synchronization framework (client only)
+"""
+
+from .sidecarv2 import *
+
+__all__ = ["OnlyClientSidecar"]

pycityagent/environment/sidecar/sidecarv2.py

@@ -0,0 +1,109 @@
+import logging
+from time import sleep
+from typing import cast
+
+import grpc
+from pycityproto.city.sync.v2 import sync_service_pb2 as sync_service
+from pycityproto.city.sync.v2 import sync_service_pb2_grpc as sync_grpc
+
+from ..utils.grpc import create_channel
+
+__all__ = ["OnlyClientSidecar"]
+
+
+class OnlyClientSidecar:
+    """
+    Sidecar框架服务（仅支持作为客户端，不支持对外提供gRPC服务）
+    Sidecar framework service (only supported as a client, does not support external gRPC services)
+    """
+
+    def __init__(self, name: str, syncer_address: str, secure: bool = False):
+        """
+        Args:
+        - name (str): 本服务在etcd上的注册名。The registered name of this service on etcd.
+        - server_address (str): syncer地址。syncer address.
+        - listen_address (str): sidecar监听地址。sidecar listening address.
+        - secure (bool, optional): 是否使用安全连接. Defaults to False. Whether to use a secure connection. Defaults to False.
+        """
+        self._name = name
+        channel = create_channel(syncer_address, secure)
+        self._sync_stub = sync_grpc.SyncServiceStub(channel)
+
+    def wait_url(self, name: str) -> str:
+        """
+        获取服务的uri
+        Get the uri of the service
+
+        Args:
+        - name (str): 服务的注册名。Service registration name.
+
+        Returns:
+        - str: 服务的url。service url.
+        """
+        while True:
+            try:
+                resp = cast(
+                    sync_service.GetURLResponse,
+                    self._sync_stub.GetURL(sync_service.GetURLRequest(name=name)),
+                )
+                url = resp.url
+                break
+            except grpc.RpcError as e:
+                logging.warning("get uri failed, retrying..., %s", e)
+                sleep(1)
+
+        logging.debug("get uri: %s for name=%s", url, name)
+        return url
+
+    def step(self, close: bool = False) -> bool:
+        """
+        同步器步进
+        synchronizer step up
+
+        Args:
+        - close (bool): 是否退出模拟。Whether the simulation exited.
+
+        Returns:
+        - close (bool): 是否退出模拟。Whether the simulation exited.
+        """
+        self._sync_stub.EnterStepSync(
+            sync_service.EnterStepSyncRequest(name=self._name)
+        )
+        response = self._sync_stub.ExitStepSync(
+            sync_service.ExitStepSyncRequest(name=self._name, close=close)
+        )
+        return response.close
+
+    def init(self) -> bool:
+        """
+        同步器初始化
+        Synchronizer initialization
+
+        Returns:
+        - close (bool): 是否退出模拟。Whether the simulation exited.
+
+        Examples:
+        ```python
+        close = client.init()
+        print(close)
+        # > False
+        ```
+        """
+        return self.step()
+
+    def close(self) -> bool:
+        """
+        同步器关闭
+        Synchronizer close
+
+        Returns:
+        - close (bool): 是否退出模拟。Whether the simulation exited.
+        """
+        return self.step(True)
+
+    def notify_step_ready(self):
+        """
+        通知prepare阶段已完成
+        Notify that the prepare phase is completed
+        """
+        ...

pycityagent/environment/sim/__init__.py

@@ -0,0 +1,27 @@
+"""
+模拟器gRPC接入客户端
+Simulator gRPC access client
+"""
+
+from .person_service import PersonService
+from .aoi_service import AoiService
+from .client import CityClient
+from .clock_service import ClockService
+from .economy_services import EconomyOrgService, EconomyPersonService
+from .lane_service import LaneService
+from .road_service import RoadService
+from .social_service import SocialService
+from .light_service import LightService
+
+__all__ = [
+    "CityClient",
+    "ClockService",
+    "PersonService",
+    "AoiService",
+    "LaneService",
+    "RoadService",
+    "SocialService",
+    "EconomyPersonService",
+    "EconomyOrgService",
+    "LightService",
+]

pycityagent/environment/sim/aoi_service.py

@@ -0,0 +1,38 @@
+from typing import Any, Awaitable, Coroutine, cast, Union, Dict
+
+import grpc
+from google.protobuf.json_format import ParseDict
+from pycityproto.city.map.v2 import aoi_service_pb2 as aoi_service
+from pycityproto.city.map.v2 import aoi_service_pb2_grpc as aoi_grpc
+
+from ..utils.protobuf import async_parse
+
+__all__ = ["AoiService"]
+
+
+class AoiService:
+    """
+    aoi服务
+    AOI service
+    """
+
+    def __init__(self, aio_channel: grpc.aio.Channel):
+        self._aio_stub = aoi_grpc.AoiServiceStub(aio_channel)
+
+    def GetAoi(
+        self, req: Union[aoi_service.GetAoiRequest, dict], dict_return: bool = True
+    ) -> Coroutine[Any, Any, Union[Dict[str, Any], aoi_service.GetAoiResponse]]:
+        """
+        获取AOI信息
+        get AOI information
+        
+        Args:
+        - req (dict): https://cityproto.sim.fiblab.net/#city.map.v2.GetAoiRequest
+
+        Returns:
+        - https://cityproto.sim.fiblab.net/#city.map.v2.GetAoiResponse
+        """
+        if type(req) != aoi_service.GetAoiRequest:
+            req = ParseDict(req, aoi_service.GetAoiRequest())
+        res = cast(Awaitable[aoi_service.GetAoiResponse], self._aio_stub.GetAoi(req))
+        return async_parse(res, dict_return)