pycityagent 2.0.0a94__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a96__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/configs/exp_config.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
+
+from pydantic import BaseModel, Field
+
+from ..utils import WorkflowType
+
+if TYPE_CHECKING:
+    from ..simulation import AgentSimulation
+
+
+class WorkflowStep(BaseModel):
+    type: WorkflowType
+    days: int = Field(1, description="Number of simulation days")
+    times: int = Field(1, description="Step Execution Times")
+    description: str = Field("no description")
+    func: Optional[Callable] = None
+
+
+class AgentConfig(BaseModel):
+
+    number_of_citizen: int = Field(1, description="Number of citizens")
+    number_of_firm: int = Field(1, description="Number of firms")
+    number_of_government: int = Field(1, description="Number of governments")
+    number_of_bank: int = Field(1, description="Number of banks")
+    number_of_nbs: int = Field(1, description="Number of neighborhood-based services")
+    group_size: int = Field(100, description="Size of agent groups")
+    embedding_model: Any = Field(None, description="Embedding model")
+    agent_class_configs: Optional[dict[Any, dict[str, Any]]] = None
+    memory_config_func: Optional[dict[type["Any"], Callable]] = None
+    memory_config_init_func: Optional[Callable] = Field(None)
+    init_func: Optional[list[Callable[["AgentSimulation"], None]]] = None
+    enable_institution: bool = Field(
+        True, description="Whether institutions are enabled in the experiment"
+    )
+
+    @classmethod
+    def create(
+        cls,
+        number_of_citizen: int = 1,
+        number_of_firm: int = 1,
+        number_of_government: int = 1,
+        number_of_bank: int = 1,
+        number_of_nbs: int = 1,
+        group_size: int = 100,
+        embedding_model: Any = None,
+        agent_class_configs: Optional[dict[Any, dict[str, Any]]] = None,
+        enable_institution: bool = True,
+        memory_config_func: Optional[dict[type["Any"], Callable]] = None,
+        memory_config_init_func: Optional[Callable] = None,
+        init_func: Optional[list[Callable[["AgentSimulation"], None]]] = None,
+    ) -> "AgentConfig":
+        return cls(
+            number_of_citizen=number_of_citizen,
+            number_of_firm=number_of_firm,
+            number_of_government=number_of_government,
+            number_of_bank=number_of_bank,
+            number_of_nbs=number_of_nbs,
+            group_size=group_size,
+            embedding_model=embedding_model,
+            agent_class_configs=agent_class_configs,
+            enable_institution=enable_institution,
+            memory_config_func=memory_config_func,
+            memory_config_init_func=memory_config_init_func,
+            init_func=init_func,
+        )
+
+
+class EnvironmentConfig(BaseModel):
+    weather: str = Field(default="The weather is normal")
+    crime: str = Field(default="The crime rate is low")
+    pollution: str = Field(default="The pollution level is low")
+    temperature: str = Field(default="The temperature is normal")
+    day: str = Field(default="Workday")
+
+    @classmethod
+    def create(
+        cls,
+        weather: str = "The weather is normal",
+        crime: str = "The crime rate is low",
+        pollution: str = "The pollution level is low",
+        temperature: str = "The temperature is normal",
+        day: str = "Workday",
+    ) -> "EnvironmentConfig":
+        return cls(
+            weather=weather,
+            crime=crime,
+            pollution=pollution,
+            temperature=temperature,
+            day=day,
+        )
+
+
+class MessageInterceptConfig(BaseModel):
+    mode: str = Field(..., pattern="^(point|edge)$")
+    max_violation_time: int = Field(default=3)
+
+    @classmethod
+    def create(cls, mode: str, max_violation_time: int = 3) -> "MessageInterceptConfig":
+        return cls(mode=mode, max_violation_time=max_violation_time)
+
+
+class ExpConfig(BaseModel):
+    agent_config: Optional[AgentConfig] = None
+    workflow: Optional[list[WorkflowStep]] = None
+    environment: Optional[EnvironmentConfig] = EnvironmentConfig()
+    message_intercept: Optional[MessageInterceptConfig] = None
+    metric_extractors: Optional[list[tuple[int, Callable]]] = None
+    logging_level: int = Field(logging.WARNING)
+    exp_name: str = Field("default_experiment")
+    llm_semaphore: int = Field(200)
+
+    @property
+    def prop_agent_config(self) -> AgentConfig:
+        return self.agent_config  # type:ignore
+
+    @property
+    def prop_workflow(self) -> list[WorkflowStep]:
+        return self.workflow  # type:ignore
+
+    @property
+    def prop_environment(self) -> EnvironmentConfig:
+        return self.environment  # type:ignore
+
+    @property
+    def prop_message_intercept(self) -> MessageInterceptConfig:
+        return self.message_intercept  # type:ignore
+
+    @property
+    def prop_metric_extractors(
+        self,
+    ) -> list[tuple[int, Callable]]:
+        return self.metric_extractors  # type:ignore
+
+    def SetAgentConfig(
+        self,
+        number_of_citizen: int = 1,
+        number_of_firm: int = 1,
+        number_of_government: int = 1,
+        number_of_bank: int = 1,
+        number_of_nbs: int = 1,
+        group_size: int = 100,
+        embedding_model: Any = None,
+        agent_class_configs: Optional[dict[Any, dict[str, list[dict]]]] = None,
+        enable_institution: bool = True,
+        memory_config_func: Optional[dict[type["Any"], Callable]] = None,
+        memory_config_init_func: Optional[Callable] = None,
+        init_func: Optional[list[Callable[["AgentSimulation"], None]]] = None,
+    ) -> "ExpConfig":
+        self.agent_config = AgentConfig.create(
+            number_of_citizen=number_of_citizen,
+            number_of_firm=number_of_firm,
+            number_of_government=number_of_government,
+            number_of_bank=number_of_bank,
+            number_of_nbs=number_of_nbs,
+            group_size=group_size,
+            embedding_model=embedding_model,
+            agent_class_configs=agent_class_configs,
+            enable_institution=enable_institution,
+            memory_config_func=memory_config_func,
+            memory_config_init_func=memory_config_init_func,
+            init_func=init_func,
+        )
+        return self
+
+    def SetEnvironment(
+        self,
+        weather: str = "The weather is normal",
+        crime: str = "The crime rate is low",
+        pollution: str = "The pollution level is low",
+        temperature: str = "The temperature is normal",
+        day: str = "Workday",
+    ) -> "ExpConfig":
+        self.environment = EnvironmentConfig.create(
+            weather=weather,
+            crime=crime,
+            pollution=pollution,
+            temperature=temperature,
+            day=day,
+        )
+        return self
+
+    def SetMessageIntercept(
+        self,
+        mode: Union[Literal["point"], Literal["edge"]],
+        max_violation_time: int = 3,
+    ) -> "ExpConfig":
+        self.message_intercept = MessageInterceptConfig.create(
+            mode=mode, max_violation_time=max_violation_time
+        )
+        return self
+
+    def SetMetricExtractors(self, metric_extractors: list[tuple[int, Callable]]):
+        self.metric_extractors = metric_extractors
+        return self
+
+    def SetWorkFlow(self, workflows: list[WorkflowStep]):
+        self.workflow = workflows
+        return self

pycityagent/configs/sim_config.py

@@ -0,0 +1,251 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field
+
+from ..utils import LLMRequestType
+
+__all__ = [
+    "SimConfig",
+]
+
+
+class LLMRequestConfig(BaseModel):
+    request_type: LLMRequestType = Field(
+        ..., description="The type of the request or provider"
+    )
+    api_key: list[str] = Field(..., description="API key for accessing the service")
+    model: str = Field(..., description="The model to use")
+
+    @classmethod
+    def create(
+        cls, request_type: LLMRequestType, api_key: list[str], model: str
+    ) -> "LLMRequestConfig":
+        return cls(request_type=request_type, api_key=api_key, model=model)
+
+
+class MQTTConfig(BaseModel):
+    server: str = Field(..., description="MQTT server address")
+    port: int = Field(..., description="Port number for MQTT connection")
+    password: Optional[str] = Field(None, description="Password for MQTT connection")
+    username: Optional[str] = Field(None, description="Username for MQTT connection")
+
+    @classmethod
+    def create(
+        cls,
+        server: str,
+        port: int,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+    ) -> "MQTTConfig":
+        return cls(server=server, username=username, port=port, password=password)
+
+
+class SimulatorRequestConfig(BaseModel):
+    task_name: str = Field("citysim", description="Name of the simulation task")
+    max_day: int = Field(1000, description="Maximum number of days to simulate")
+    start_step: int = Field(28800, description="Starting step of the simulation")
+    total_step: int = Field(
+        24 * 60 * 60 * 365, description="Total number of steps in the simulation"
+    )
+    log_dir: str = Field("./log", description="Directory path for saving logs")
+    min_step_time: int = Field(
+        1000, description="Minimum time (in seconds) between simulation steps"
+    )
+    primary_node_ip: str = Field(
+        "localhost", description="Primary node IP address for distributed simulation"
+    )
+
+    @classmethod
+    def create(
+        cls,
+        task_name: str = "citysim",
+        max_day: int = 1000,
+        start_step: int = 28800,
+        total_step: int = 24 * 60 * 60 * 365,
+        log_dir: str = "./log",
+        min_step_time: int = 1000,
+        primary_node_ip: str = "localhost",
+    ) -> "SimulatorRequestConfig":
+        return cls(
+            task_name=task_name,
+            max_day=max_day,
+            start_step=start_step,
+            total_step=total_step,
+            log_dir=log_dir,
+            min_step_time=min_step_time,
+            primary_node_ip=primary_node_ip,
+        )
+
+
+class MapRequestConfig(BaseModel):
+    file_path: str = Field(..., description="Path to the map file")
+
+    @classmethod
+    def create(cls, file_path: str) -> "MapRequestConfig":
+        return cls(file_path=file_path)
+
+
+class MlflowConfig(BaseModel):
+    username: Optional[str] = Field(None, description="Username for MLflow")
+    password: Optional[str] = Field(None, description="Password for MLflow")
+    mlflow_uri: str = Field(..., description="URI for MLflow server")
+
+    @classmethod
+    def create(cls, username: str, password: str, mlflow_uri: str) -> "MlflowConfig":
+        return cls(username=username, password=password, mlflow_uri=mlflow_uri)
+
+
+class PostgreSQLConfig(BaseModel):
+    enabled: Optional[bool] = Field(
+        True, description="Whether PostgreSQL storage is enabled"
+    )
+    dsn: str = Field(..., description="Data source name for PostgreSQL")
+
+    @classmethod
+    def create(cls, dsn: str, enabled: bool = False) -> "PostgreSQLConfig":
+        return cls(dsn=dsn, enabled=enabled)
+
+
+class AvroConfig(BaseModel):
+    enabled: Optional[bool] = Field(False, description="Whether Avro storage is enabled")
+    path: str = Field(..., description="Avro file storage path")
+
+    @classmethod
+    def create(cls, path: Optional[str] = None, enabled: bool = False) -> "AvroConfig":
+        return cls(enabled=enabled, path=path)
+
+
+class MetricRequest(BaseModel):
+    mlflow: Optional[MlflowConfig] = Field(None)
+
+
+class SimStatus(BaseModel):
+    simulator_activated: bool = False
+
+
+class SimConfig(BaseModel):
+    llm_request: Optional["LLMRequestConfig"] = None
+    simulator_request: Optional["SimulatorRequestConfig"] = None
+    mqtt: Optional["MQTTConfig"] = None
+    map_request: Optional["MapRequestConfig"] = None
+    metric_request: Optional[MetricRequest] = None
+    pgsql: Optional["PostgreSQLConfig"] = None
+    avro: Optional["AvroConfig"] = None
+    simulator_server_address: Optional[str] = None
+    status: Optional["SimStatus"] = SimStatus()
+
+    @property
+    def prop_llm_request(self) -> "LLMRequestConfig":
+        return self.llm_request  # type:ignore
+
+    @property
+    def prop_status(self) -> "SimStatus":
+        return self.status  # type:ignore
+
+    @property
+    def prop_simulator_request(self) -> "SimulatorRequestConfig":
+        return self.simulator_request  # type:ignore
+
+    @property
+    def prop_mqtt(self) -> "MQTTConfig":
+        return self.mqtt  # type:ignore
+
+    @property
+    def prop_map_request(self) -> "MapRequestConfig":
+        return self.map_request  # type:ignore
+
+    @property
+    def prop_avro_config(self) -> "AvroConfig":
+        return self.avro  # type:ignore
+
+    @property
+    def prop_postgre_sql_config(self) -> "PostgreSQLConfig":
+        return self.pgsql  # type:ignore
+
+    @property
+    def prop_simulator_server_address(self) -> str:
+        return self.simulator_server_address  # type:ignore
+
+    @property
+    def prop_metric_request(self) -> MetricRequest:
+        return self.metric_request  # type:ignore
+
+    def SetLLMRequest(
+        self, request_type: LLMRequestType, api_key: list[str], model: str
+    ) -> "SimConfig":
+        self.llm_request = LLMRequestConfig.create(request_type, api_key, model)
+        return self
+
+    def SetSimulatorRequest(
+        self,
+        task_name: str = "citysim",
+        max_day: int = 1000,
+        start_step: int = 28800,
+        total_step: int = 24 * 60 * 60 * 365,
+        log_dir: str = "./log",
+        min_step_time: int = 1000,
+        primary_node_ip: str = "localhost",
+    ) -> "SimConfig":
+        self.simulator_request = SimulatorRequestConfig.create(
+            task_name=task_name,
+            max_day=max_day,
+            start_step=start_step,
+            total_step=total_step,
+            log_dir=log_dir,
+            min_step_time=min_step_time,
+            primary_node_ip=primary_node_ip,
+        )
+        return self
+
+    def SetMQTT(
+        self,
+        server: str,
+        port: int,
+        username: Optional[str] = None,
+        password: Optional[str] = None,
+    ) -> "SimConfig":
+        self.mqtt = MQTTConfig.create(server, port, username, password)
+        return self
+
+    def SetMapRequest(self, file_path: str) -> "SimConfig":
+        self.map_request = MapRequestConfig.create(file_path)
+        return self
+
+    def SetMetricRequest(
+        self, username: str, password: str, mlflow_uri: str
+    ) -> "SimConfig":
+        self.metric_request = MetricRequest(
+            mlflow=MlflowConfig.create(username, password, mlflow_uri)
+        )
+        return self
+
+    def SetAvro(self, path: str, enabled: bool = False) -> "SimConfig":
+        self.avro = AvroConfig.create(path, enabled)
+        return self
+
+    def SetPostgreSql(self, path: str, enabled: bool = False) -> "SimConfig":
+        self.pgsql = PostgreSQLConfig.create(path, enabled)
+        return self
+
+    def SetServerAddress(self, simulator_server_address: str) -> "SimConfig":
+        self.simulator_server_address = simulator_server_address
+        return self
+
+    def model_dump(self, *args, **kwargs):
+        exclude_fields = {
+            "status",
+        }
+        data = super().model_dump(*args, **kwargs)
+        return {k: v for k, v in data.items() if k not in exclude_fields}
+
+
+if __name__ == "__main__":
+    config = (
+        SimConfig()
+        .SetLLMRequest("openai", "key", "model")  # type:ignore
+        .SetMQTT("server", 1883, "username", "password")
+        .SetMapRequest("./path/to/map")
+        .SetMetricRequest("username", "password", "uri")
+        .SetPostgreSql("dsn", True)
+    )
+    print(config.llm_request)

pycityagent/configs/utils.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Union
+
+import yaml
+
+if TYPE_CHECKING:
+    from .exp_config import ExpConfig
+    from .sim_config import SimConfig
+
+
+def load_config_from_file(
+    filepath: str, config_type: Union[type[SimConfig], type[ExpConfig]]
+) -> Union[SimConfig, ExpConfig]:
+    with open(filepath, "r") as file:
+        data = yaml.safe_load(file)
+    return config_type(**data)

pycityagent/environment/__init__.py

@@ -2,9 +2,11 @@
 
 from .sim import AoiService, PersonService
 from .simulator import Simulator
+from .economy import EconomyClient
 
 __all__ = [
     "Simulator",
     "PersonService",
     "AoiService",
+    "EconomyClient",
 ]

pycityagent/{economy → environment/economy}/econ_client.py

@@ -9,7 +9,6 @@ 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
 from google.protobuf.json_format import MessageToDict
 
 logger = logging.getLogger("pycityagent")
@@ -19,27 +18,6 @@ __all__ = [
 ]
 
 
-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 camel_to_snake(d):
-    if not isinstance(d, dict):
-        return d
-    return {
-        re.sub("([a-z0-9])([A-Z])", r"\1_\2", k).lower(): (
-            camel_to_snake(v) if isinstance(v, dict) else v
-        )
-        for k, v in d.items()
-    }
-
-
 def _create_aio_channel(server_address: str, secure: bool = False) -> grpc.aio.Channel:
     """
     Create a gRPC asynchronous channel.
@@ -170,8 +148,9 @@ class EconomyClient:
             agents = await self._aio_stub.BatchGet(
                 org_service.BatchGetRequest(ids=id, type="agent")
             )
-            agents = MessageToDict(agents)["agents"]
-            agent_dicts = [camel_to_snake(agent) for agent in agents]
+            agent_dicts = MessageToDict(agents, preserving_proto_field_name=True)[
+                "agents"
+            ]
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
             return agent_dicts
@@ -179,10 +158,12 @@ class EconomyClient:
             agent = await self._aio_stub.GetAgent(
                 org_service.GetAgentRequest(agent_id=id)
             )
-            agent_dict = MessageToDict(agent)["agent"]
+            agent_dict: dict = MessageToDict(agent, preserving_proto_field_name=True)[
+                "agent"
+            ]
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
-            return camel_to_snake(agent_dict)
+            return agent_dict
 
     async def get_org(
         self, id: Union[list[int], int]
@@ -202,17 +183,16 @@ class EconomyClient:
             orgs = await self._aio_stub.BatchGet(
                 org_service.BatchGetRequest(ids=id, type="org")
             )
-            orgs = MessageToDict(orgs)["orgs"]
-            org_dicts = [camel_to_snake(org) for org in orgs]
+            org_dicts = MessageToDict(orgs, preserving_proto_field_name=True)["orgs"]
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
             return org_dicts
         else:
             org = await self._aio_stub.GetOrg(org_service.GetOrgRequest(org_id=id))
-            org_dict = MessageToDict(org)["org"]
+            org_dict = MessageToDict(org, preserving_proto_field_name=True)["org"]
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
-            return camel_to_snake(org_dict)
+            return org_dict
 
     async def get(
         self,
@@ -510,7 +490,7 @@ class EconomyClient:
         return (float(response.taxes_due), list(response.updated_incomes))
 
     async def calculate_consumption(
-        self, org_ids: Union[int, list[int]], agent_id: int, demands: list[int]
+        self, org_ids: Union[int, list[int]], agent_id: int, demands: list[int], consumption_accumulation: bool = True
     ):
         """
         Calculate consumption for agents based on their demands.
@@ -519,6 +499,7 @@ class EconomyClient:
             - `org_ids` (`Union[int, list[int]]`): The ID of the firm providing goods or services.
             - `agent_id` (`int`): The ID of the agent whose consumption is being calculated.
             - `demands` (`List[int]`): A list of demand quantities corresponding to each agent.
+            - `consumption_accumulation` (`bool`): Weather accumulation.
 
         - **Returns**:
             - `Tuple[int, List[float]]`: A tuple containing the remaining inventory and updated currencies for each agent.
@@ -535,6 +516,7 @@ class EconomyClient:
             firm_ids=org_ids,
             agent_id=agent_id,
             demands=demands,
+            consumption_accumulation = consumption_accumulation,
         )
         response: org_service.CalculateConsumptionResponse = (
             await self._aio_stub.CalculateConsumption(request)
@@ -707,7 +689,7 @@ class EconomyClient:
         """
         start_time = time.time()
         log = {"req": "add_delta_value", "start_time": start_time, "consumption": 0}
-        pascal_key = _snake_to_pascal(key)
+        pascal_key = ''.join(x.title() for x in key.split('_'))
         _request_type = getattr(org_service, f"Add{pascal_key}Request")
         _request_func = getattr(self._aio_stub, f"Add{pascal_key}")
 

pycityagent/environment/sim/sim_env.py

@@ -6,7 +6,7 @@ import warnings
 from subprocess import Popen
 from typing import Optional
 
-from pycitydata.map import Map
+import yaml
 
 from ..utils import encode_to_base64, find_free_port
 
@@ -16,29 +16,22 @@ __all__ = ["ControlSimEnv"]
 def _generate_yaml_config(
     map_file: str, max_day: int, start_step: int, total_step: int
 ) -> str:
-    map_file = os.path.abspath(map_file)
-    return f"""
-input:
-  # 地图
-  map:
-    file: "{map_file}"
-
-control:
-  day: {max_day}
-  step:
-    start: {start_step}
-    total: {total_step}
-    interval: 1
-  skip_overtime_trip_when_init: true
-  enable_platoon: false
-  enable_indoor: false
-  prefer_fixed_light: true
-  enable_collision_avoidance: false # 计算性能下降10倍，需要保证subloop>=5
-  enable_go_astray: true # 引入串行的路径规划调用，计算性能下降（幅度不确定）
-  lane_change_model: earliest # mobil （主动变道+强制变道，默认值） earliest （总是尽可能早地变道）
-
-output:
-"""
+    config_dict = {
+        "input": {"map": {"file": os.path.abspath(map_file)}},
+        "control": {
+            "day": max_day,
+            "step": {"start": start_step, "total": total_step, "interval": 1},
+            "skip_overtime_trip_when_init": True,
+            "enable_platoon": False,
+            "enable_indoor": False,
+            "prefer_fixed_light": True,
+            "enable_collision_avoidance": False,
+            "enable_go_astray": True,
+            "lane_change_model": "earliest",
+        },
+        "output": None,
+    }
+    return yaml.dump(config_dict, allow_unicode=True)
 
 
 class ControlSimEnv: