pycityagent 2.0.0a84__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a85__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/economy/econ_client.py

@@ -3,7 +3,7 @@ import logging
 import re
 import time
 from collections.abc import Sequence
-from typing import Any, Literal, Union
+from typing import Any, Literal, Union, cast
 
 import grpc
 import pycityproto.city.economy.v2.economy_pb2 as economyv2
@@ -83,17 +83,18 @@ class EconomyClient:
         - It establishes an asynchronous connection and provides methods to communicate with the service.
     """
 
-    def __init__(self, server_address: str, secure: bool = False):
+    def __init__(
+        self,
+        server_address: str,
+    ):
         """
         Initialize the EconomyClient.
 
         - **Args**:
             - `server_address` (`str`): The address of the Economy server to connect to.
-            - `secure` (`bool`, optional): Whether to use a secure connection. Defaults to `False`.
 
         - **Attributes**:
             - `server_address` (`str`): The address of the Economy server.
-            - `secure` (`bool`): A flag indicating if a secure connection should be used.
             - `_aio_stub` (`OrgServiceStub`): A gRPC stub used to make remote calls to the Economy service.
 
         - **Description**:
@@ -102,6 +103,8 @@ class EconomyClient:
             - Instantiates a gRPC stub (`_aio_stub`) for interacting with the Economy service.
         """
         self.server_address = server_address
+
+        secure = self.server_address.startswith("https")
         self.secure = secure
         aio_channel = _create_aio_channel(server_address, secure)
         self._aio_stub = org_grpc.OrgServiceStub(aio_channel)
@@ -149,7 +152,9 @@ class EconomyClient:
         self._agent_ids = agent_ids
         self._org_ids = org_ids
 
-    async def get_agent(self, id: Union[list[int], int]) -> economyv2.Agent:
+    async def get_agent(
+        self, id: Union[list[int], int]
+    ) -> Union[dict[str, Any], list[dict[str, Any]]]:
         """
         Get agent by id
 
@@ -160,17 +165,10 @@ class EconomyClient:
             - `economyv2.Agent`: The agent object.
         """
         start_time = time.time()
-        log = {
-            "req": "get_agent",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "get_agent", "start_time": start_time, "consumption": 0}
         if isinstance(id, list):
             agents = await self._aio_stub.BatchGet(
-                org_service.BatchGetRequest(
-                    ids=id,
-                    type="agent"
-                )
+                org_service.BatchGetRequest(ids=id, type="agent")
             )
             agents = MessageToDict(agents)["agents"]
             agent_dicts = [camel_to_snake(agent) for agent in agents]
@@ -179,16 +177,16 @@ class EconomyClient:
             return agent_dicts
         else:
             agent = await self._aio_stub.GetAgent(
-                org_service.GetAgentRequest(
-                    agent_id=id
-                )
+                org_service.GetAgentRequest(agent_id=id)
             )
             agent_dict = MessageToDict(agent)["agent"]
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
             return camel_to_snake(agent_dict)
 
-    async def get_org(self, id: Union[list[int], int]) -> economyv2.Org:
+    async def get_org(
+        self, id: Union[list[int], int]
+    ) -> Union[dict[str, Any], list[dict[str, Any]]]:
         """
         Get org by id
 
@@ -199,17 +197,10 @@ class EconomyClient:
             - `economyv2.Org`: The org object.
         """
         start_time = time.time()
-        log = {
-            "req": "get_org",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "get_org", "start_time": start_time, "consumption": 0}
         if isinstance(id, list):
             orgs = await self._aio_stub.BatchGet(
-                org_service.BatchGetRequest(
-                    ids=id,
-                    type="org"
-                )
+                org_service.BatchGetRequest(ids=id, type="org")
             )
             orgs = MessageToDict(orgs)["orgs"]
             org_dicts = [camel_to_snake(org) for org in orgs]
@@ -217,11 +208,7 @@ class EconomyClient:
             self._log_list.append(log)
             return org_dicts
         else:
-            org = await self._aio_stub.GetOrg(
-                org_service.GetOrgRequest(
-                    org_id=id
-                )
-            )
+            org = await self._aio_stub.GetOrg(org_service.GetOrgRequest(org_id=id))
             org_dict = MessageToDict(org)["org"]
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
@@ -243,18 +230,15 @@ class EconomyClient:
             - Any
         """
         start_time = time.time()
-        log = {
-            "req": "get",
-            "start_time": start_time,
-            "consumption": 0
-        }
-        if isinstance(id, list):
+        log = {"req": "get", "start_time": start_time, "consumption": 0}
+        if isinstance(id, Sequence):
             requests = "Org" if id[0] in self._org_ids else "Agent"
             if requests == "Org":
                 response = await self.get_org(id)
             else:
                 response = await self.get_agent(id)
             results = []
+            response = cast(list[dict[str, Any]], response)
             for res in response:
                 results.append(res[key])
             log["consumption"] = time.time() - start_time
@@ -268,10 +252,11 @@ class EconomyClient:
                 response = await self.get_org(id)
             else:
                 response = await self.get_agent(id)
+            response = cast(dict[str, Any], response)
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
             return response[key]
-        
+
     def _merge(self, original_value, key, value):
         try:
             orig_value = original_value[key]
@@ -321,16 +306,14 @@ class EconomyClient:
             - Any
         """
         start_time = time.time()
-        log = {
-            "req": "update",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "update", "start_time": start_time, "consumption": 0}
         if isinstance(id, list):
             if not isinstance(value, list):
                 raise ValueError(f"Invalid value, the value must be a list!")
             if len(id) != len(value):
-                raise ValueError(f"Invalid ids and values, the length of ids and values must be the same!")
+                raise ValueError(
+                    f"Invalid ids and values, the length of ids and values must be the same!"
+                )
             request_type = "Org" if id[0] in self._org_ids else "Agent"
         else:
             if id not in self._agent_ids and id not in self._org_ids:
@@ -363,16 +346,11 @@ class EconomyClient:
                 #     ))
                 # await asyncio.gather(*batch_update_tasks)
                 await self._aio_stub.BatchUpdate(
-                    org_service.BatchUpdateRequest(
-                        orgs=original_value,
-                        agents=None
-                    )
+                    org_service.BatchUpdateRequest(orgs=original_value, agents=None)
                 )
             else:
                 await self._aio_stub.UpdateOrg(
-                    org_service.UpdateOrgRequest(
-                        org=original_value
-                    )
+                    org_service.UpdateOrgRequest(org=original_value)
                 )
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
@@ -387,16 +365,11 @@ class EconomyClient:
                 #     ))
                 # await asyncio.gather(*batch_update_tasks)
                 await self._aio_stub.BatchUpdate(
-                    org_service.BatchUpdateRequest(
-                        orgs=None,
-                        agents=original_value
-                    )
+                    org_service.BatchUpdateRequest(orgs=None, agents=original_value)
                 )
             else:
                 await self._aio_stub.UpdateAgent(
-                    org_service.UpdateAgentRequest(
-                        agent=original_value
-                    )
+                    org_service.UpdateAgentRequest(agent=original_value)
                 )
             log["consumption"] = time.time() - start_time
             self._log_list.append(log)
@@ -572,18 +545,11 @@ class EconomyClient:
             return response.actual_consumption
         else:
             return -1
-        
-    
+
     async def calculate_real_gdp(self, nbs_id: int):
         start_time = time.time()
-        log = {
-            "req": "calculate_real_gdp",
-            "start_time": start_time,
-            "consumption": 0
-        }
-        request = org_service.CalculateRealGDPRequest(
-            nbs_agent_id=nbs_id
-        )
+        log = {"req": "calculate_real_gdp", "start_time": start_time, "consumption": 0}
+        request = org_service.CalculateRealGDPRequest(nbs_agent_id=nbs_id)
         response: org_service.CalculateRealGDPResponse = (
             await self._aio_stub.CalculateRealGDP(request)
         )

pycityagent/environment/simulator.py

@@ -87,49 +87,63 @@ class Simulator:
         - 模拟器配置
         - simulator config
         """
-        _mongo_uri, _mongo_db, _mongo_coll, _map_cache_dir = (
-            config["map_request"]["mongo_uri"],
-            config["map_request"]["mongo_db"],
-            config["map_request"]["mongo_coll"],
-            config["map_request"]["cache_dir"],
-        )
-        _mongo_client = MongoClient(_mongo_uri)
-        os.makedirs(_map_cache_dir, exist_ok=True)
-        _map_pb_path = os.path.join(_map_cache_dir, f"{_mongo_db}.{_mongo_coll}.pb")  # type: ignore
-        _map_pb = map_pb2.Map()
-        if os.path.exists(_map_pb_path):
-            with open(_map_pb_path, "rb") as f:
-                _map_pb.ParseFromString(f.read())
+        _map_request = config["map_request"]
+        if "file_path" not in _map_request:
+            # from mongo db
+            _mongo_uri, _mongo_db, _mongo_coll, _map_cache_dir = (
+                _map_request["mongo_uri"],
+                _map_request["mongo_db"],
+                _map_request["mongo_coll"],
+                _map_request["cache_dir"],
+            )
+            _mongo_client = MongoClient(_mongo_uri)
+            os.makedirs(_map_cache_dir, exist_ok=True)
+            _map_pb_path = os.path.join(_map_cache_dir, f"{_mongo_db}.{_mongo_coll}.pb")  # type: ignore
+            _map_pb = map_pb2.Map()
+            if os.path.exists(_map_pb_path):
+                with open(_map_pb_path, "rb") as f:
+                    _map_pb.ParseFromString(f.read())
+            else:
+                _map_pb = coll2pb(_mongo_client[_mongo_db][_mongo_coll], _map_pb)
+                with open(_map_pb_path, "wb") as f:
+                    f.write(_map_pb.SerializeToString())
         else:
-            _map_pb = coll2pb(_mongo_client[_mongo_db][_mongo_coll], _map_pb)
-            with open(_map_pb_path, "wb") as f:
-                f.write(_map_pb.SerializeToString())
+            # from local file
+            _mongo_uri, _mongo_db, _mongo_coll, _map_cache_dir = "", "", "", ""
+            _map_pb_path = _map_request["file_path"]
 
         if "simulator" in config:
             if config["simulator"] is None:
                 config["simulator"] = {}
-            if "server" not in config["simulator"]:
+            if not config["simulator"].get("_server_activated", False):
                 self._sim_env = sim_env = ControlSimEnv(
                     task_name=config["simulator"].get("task", "citysim"),
                     map_file=_map_pb_path,
                     max_day=config["simulator"].get("max_day", 1000),
                     start_step=config["simulator"].get("start_step", 28800),
-                    total_step=config["simulator"].get("total_step", 24*60*60*365),
+                    total_step=config["simulator"].get(
+                        "total_step", 24 * 60 * 60 * 365
+                    ),
                     log_dir=config["simulator"].get("log_dir", "./log"),
                     min_step_time=config["simulator"].get("min_step_time", 1000),
                     sim_addr=config["simulator"].get("server", None),
                 )
                 self.server_addr = sim_env.sim_addr
                 config["simulator"]["server"] = self.server_addr
+                config["simulator"]["_server_activated"] = True
                 # using local client
-                self._client = CityClient(sim_env.sim_addr, secure=False)
+                self._client = CityClient(
+                    sim_env.sim_addr, secure=self.server_addr.startswith("https")
+                )
                 """
                 - 模拟器grpc客户端
                 - grpc client of simulator
                 """
             else:
-                self._client = CityClient(config["simulator"]["server"], secure=False)
                 self.server_addr = config["simulator"]["server"]
+                self._client = CityClient(
+                    self.server_addr, secure=self.server_addr.startswith("https")
+                )
         else:
             self.server_addr = None
             logger.warning(
@@ -141,10 +155,9 @@ class Simulator:
         - Simulator map object
         """
         if create_map:
-            _map_cache_path = ""  # 地图pb文件路径
             self._map = CityMap.remote(
                 (_mongo_uri, _mongo_db, _mongo_coll, _map_cache_dir),
-                _map_cache_path,
+                _map_pb_path,
             )
             self._create_poi_id_2_aoi_id()
 
@@ -363,7 +376,7 @@ class Simulator:
         log = {"req": "get_simulator_day", "start_time": start_time, "consumption": 0}
         now = await self._client.clock_service.Now({})
         now = cast(dict[str, int], now)
-        day = int(now["t"]//(24*60*60))
+        day = int(now["t"] // (24 * 60 * 60))
         log["consumption"] = time.time() - start_time
         self._log_list.append(log)
         return day
@@ -385,7 +398,7 @@ class Simulator:
         now = cast(dict[str, int], now)
         log["consumption"] = time.time() - start_time
         self._log_list.append(log)
-        return now["t"]%(24*60*60)
+        return now["t"] % (24 * 60 * 60)
 
     async def get_person(self, person_id: int) -> dict:
         """

pycityagent/simulation/agentgroup.py

@@ -22,13 +22,8 @@ from ..llm.llmconfig import LLMConfig
 from ..memory import FaissQuery, Memory
 from ..message import Messager
 from ..metrics import MlflowClient
-from ..utils import (
-    DIALOG_SCHEMA,
-    INSTITUTION_STATUS_SCHEMA,
-    PROFILE_SCHEMA,
-    STATUS_SCHEMA,
-    SURVEY_SCHEMA,
-)
+from ..utils import (DIALOG_SCHEMA, INSTITUTION_STATUS_SCHEMA, PROFILE_SCHEMA,
+                     STATUS_SCHEMA, SURVEY_SCHEMA)
 
 logger = logging.getLogger("pycityagent")
 __all__ = ["AgentGroup"]

pycityagent/simulation/simulation.py

@@ -1021,7 +1021,9 @@ class AgentSimulation:
         try:
             # step
             simulator_day = await self._simulator.get_simulator_day()
-            simulator_time = int(await self._simulator.get_simulator_second_from_start_of_day())
+            simulator_time = int(
+                await self._simulator.get_simulator_second_from_start_of_day()
+            )
             logger.info(
                 f"Start simulation day {simulator_day} at {simulator_time}, step {self._total_steps}"
             )

{pycityagent-2.0.0a84.dist-info → pycityagent-2.0.0a85.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pycityagent
-Version: 2.0.0a84
+Version: 2.0.0a85
 Summary: LLM-based city environment agent building library
 Author-email: Yuwei Yan <pinkgranite86@gmail.com>, Junbo Yan <yanjb20thu@gmali.com>, Jun Zhang <zhangjun990222@gmali.com>
 License: MIT License

{pycityagent-2.0.0a84.dist-info → pycityagent-2.0.0a85.dist-info}/RECORD

@@ -1,10 +1,16 @@
+pycityagent-2.0.0a85.dist-info/RECORD,,
+pycityagent-2.0.0a85.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
+pycityagent-2.0.0a85.dist-info/WHEEL,sha256=NW1RskY9zow1Y68W-gXg0oZyBRAugI1JHywIzAIai5o,109
+pycityagent-2.0.0a85.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
+pycityagent-2.0.0a85.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
+pycityagent-2.0.0a85.dist-info/METADATA,sha256=KKmALYa8DummI6EtxRIqeGaH5Fg_lS4yElcWAOdV2wU,9110
 pycityagent/pycityagent-sim,sha256=Ax6cjHjT8VcT5t07VboL_Ruor2eXSdSEml4Jvv-L5dQ,36972594
 pycityagent/__init__.py,sha256=PUKWTXc-xdMG7px8oTNclodsILUgypANj2Z647sY63k,808
 pycityagent/pycityagent-ui,sha256=Ur95yZygIaZ5l_CDqP9394M5GQ66iV5PkcNPYFWqzvk,41225346
 pycityagent/metrics/mlflow_client.py,sha256=-iyh4BPVnBkluhmfucUzibCUnBX2iftkz4tVJJVxwHw,6958
 pycityagent/metrics/__init__.py,sha256=X08PaBbGVAd7_PRGLREXWxaqm7nS82WBQpD1zvQzcqc,128
 pycityagent/economy/__init__.py,sha256=aonY4WHnx-6EGJ4WKrx4S-2jAkYNLtqUA04jp6q8B7w,75
-pycityagent/economy/econ_client.py,sha256=NUzrFzl0EL4UkMscGRRSu1JQmeCPwIkPvTbJDRefUcQ,29231
+pycityagent/economy/econ_client.py,sha256=KXoRKnOrqYlFqgLHpNim1C5QRM7QPO3GtXG1Aq-yf6U,28621
 pycityagent/tools/__init__.py,sha256=y7sMVMHf0AbivlczM2h-kr7mkgXK-WAx3S9BXLXkWvw,235
 pycityagent/tools/tool.py,sha256=4ZJSHbNM8dfAVwZEw8T0a2_9OuPsPQpKSVL4WxZVBUc,9022
 pycityagent/llm/llmconfig.py,sha256=6AqCMV4B_JqBD2mb98bLGzpUdlOCnziQKae-Hhxxp-E,469
@@ -21,9 +27,9 @@ pycityagent/memory/utils.py,sha256=oJWLdPeJy_jcdKcDTo9JAH9kDZhqjoQhhv_zT9qWC0w,8
 pycityagent/memory/const.py,sha256=nFmjvt-8FEB0hc0glOH3lliqJhkhf3D_NKxWI0pf6TY,936
 pycityagent/memory/faiss_query.py,sha256=KPeyzIjD0dzkxr-TlOeCiIlkdh1WAyyipCAXUEt97Lk,17350
 pycityagent/memory/state.py,sha256=JFCBuK1AS-HqscvdGS9ATF9AUU8t29_2leDY_6iO2_4,5158
-pycityagent/simulation/simulation.py,sha256=08xtIVwmO35r5XbUW9nMDo9rL_6mtDtjWbeXnUyNR6c,49340
+pycityagent/simulation/simulation.py,sha256=bHpR3FtUM6ewurwByRr57OgdTWXTPq9-knIjGt2GHlo,49370
 pycityagent/simulation/__init__.py,sha256=u1WpgwVxPboXWMxrUQKMXJNmTKQYAeCaqZv9HmSxESY,118
-pycityagent/simulation/agentgroup.py,sha256=oXo2iRUThh_96JhXHMufT8_BPx0PrlwqbMpiVqK_viA,36957
+pycityagent/simulation/agentgroup.py,sha256=rPloocCwB6rfjWbIz3T052mYq2lEcBp97mblkMD2y0A,36955
 pycityagent/simulation/storage/pg.py,sha256=xRshSOGttW-p0re0fNBOjOpb-nQ5msIE2LsdT79_E_Y,8425
 pycityagent/message/message_interceptor.py,sha256=QWuTUqi1Cu214fhFs0f78tD2zflMnb6zEAGB4RutXxs,17736
 pycityagent/message/__init__.py,sha256=f5QH7DKPqEAMyfSlBMnl3uouOKlsoel909STlIe7nUk,276
@@ -46,7 +52,7 @@ pycityagent/workflow/prompt.py,sha256=rzenP4EFGxbWE1aq-x2036b6umKvi5cQx2xtWULwgI
 pycityagent/workflow/block.py,sha256=WJfCeL8e117GzkVPJCRNsQZZinccMnVyEubkwrf-17U,12295
 pycityagent/workflow/trigger.py,sha256=4nzAwywGQsFabgo6gzp-vD2EV4wII7Z0LHGEAsflSEY,7608
 pycityagent/environment/__init__.py,sha256=fFIth2jxxgZ92cXm-aoM2igHgaSqsYGwtBhyb7opjzk,166
-pycityagent/environment/simulator.py,sha256=qiqeRQDm7cprFJhAkxtdBNAvohiGGUaNIPGyFXmPdZU,22667
+pycityagent/environment/simulator.py,sha256=sbM0FMjR_q2mssbpMI7Lx9QI2iFj1Hbzg48fbRMd4tY,23206
 pycityagent/environment/utils/port.py,sha256=3OM6kSUt3PxvDUOlgyiendBtETaWU8Mzk_8H0TzTmYg,295
 pycityagent/environment/utils/grpc.py,sha256=_lB4-k4dTKuNvApaDiYgFxiLTPtYG42DVQtG9yOj9pQ,2022
 pycityagent/environment/utils/base64.py,sha256=hoREzQo3FXMN79pqQLO2jgsDEvudciomyKii7MWljAM,374
@@ -89,9 +95,3 @@ pycityagent/cityagent/blocks/mobility_block.py,sha256=qkRiV0nkGOUUoo9lGIjAFE_Hl0
 pycityagent/survey/models.py,sha256=g3xni4GcA1Py3vlGt6z4ltutjgQ4G0uINYAM8vKRJAw,5225
 pycityagent/survey/__init__.py,sha256=rxwou8U9KeFSP7rMzXtmtp2fVFZxK4Trzi-psx9LPIs,153
 pycityagent/survey/manager.py,sha256=tHkdeq4lTfAHwvgf4-udsXri0z2l6E00rEbvwl7SqRs,3439
-pycityagent-2.0.0a84.dist-info/RECORD,,
-pycityagent-2.0.0a84.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
-pycityagent-2.0.0a84.dist-info/WHEEL,sha256=NW1RskY9zow1Y68W-gXg0oZyBRAugI1JHywIzAIai5o,109
-pycityagent-2.0.0a84.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
-pycityagent-2.0.0a84.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
-pycityagent-2.0.0a84.dist-info/METADATA,sha256=wK7JDsxgO5pTlH_YLtftS9-nmD9T97_6-YZp8zru9dA,9110