pycityagent 2.0.0a73__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a75__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/environment/simulator.py

@@ -3,12 +3,13 @@
 import asyncio
 import logging
 import os
-from datetime import datetime, timedelta
 import time
-from typing import Any, Optional, Union, cast
+from datetime import datetime, timedelta
+from typing import Optional, Union, cast
 
+import ray
 from mosstool.type import TripMode
-from mosstool.util.format_converter import coll2pb
+from mosstool.util.format_converter import coll2pb, dict2pb
 from pycitydata.map import Map as SimMap
 from pycityproto.city.map.v2 import map_pb2 as map_pb2
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
@@ -26,6 +27,51 @@ __all__ = [
 ]
 
 
+@ray.remote
+class CityMap:
+    def __init__(self, mongo_input: tuple[str, str, str, str], map_cache_path: str):
+        if map_cache_path:
+            self.map = SimMap(
+                pb_path=map_cache_path,
+            )
+        else:
+            mongo_uri, mongo_db, mongo_coll, cache_dir = mongo_input
+            self.map = SimMap(
+                mongo_uri=mongo_uri,
+                mongo_db=mongo_db,
+                mongo_coll=mongo_coll,
+                cache_dir=cache_dir,
+            )
+        self.poi_cate = POI_CATG_DICT
+
+    def get_aoi(self, aoi_id: Optional[int] = None):
+        if aoi_id is None:
+            return list(self.map.aois.values())
+        else:
+            return self.map.aois[aoi_id]
+
+    def get_poi(self, poi_id: Optional[int] = None):
+        if poi_id is None:
+            return list(self.map.pois.values())
+        else:
+            return self.map.pois[poi_id]
+
+    def query_pois(self, **kwargs):
+        return self.map.query_pois(**kwargs)
+
+    def get_poi_cate(self):
+        return self.poi_cate
+
+    def get_map(self):
+        return self.map
+
+    def get_map_header(self):
+        return self.map.header
+
+    def get_projector(self):
+        return self.map.header["projection"]
+
+
 class Simulator:
     """
     Main class of the simulator.
@@ -35,7 +81,7 @@ class Simulator:
         - It reads parameters from a configuration dictionary, initializes map data, and starts or connects to a simulation server as needed.
     """
 
-    def __init__(self, config: dict, secure: bool = False) -> None:
+    def __init__(self, config: dict, create_map: bool = False) -> None:
         self.config = config
         """
         - 模拟器配置
@@ -66,7 +112,8 @@ class Simulator:
                 self._sim_env = sim_env = ControlSimEnv(
                     task_name=config["simulator"].get("task", "citysim"),
                     map_file=_map_pb_path,
-                    start_step=config["simulator"].get("start_step", 0),
+                    max_day=config["simulator"].get("max_day", 1000),
+                    start_step=config["simulator"].get("start_step", 28800),
                     total_step=2147000000,
                     log_dir=config["simulator"].get("log_dir", "./log"),
                     min_step_time=config["simulator"].get("min_step_time", 1000),
@@ -88,16 +135,18 @@ class Simulator:
             logger.warning(
                 "No simulator config found, no simulator client will be used"
             )
-        self.map = SimMap(
-            mongo_uri=_mongo_uri,
-            mongo_db=_mongo_db,
-            mongo_coll=_mongo_coll,
-            cache_dir=_map_cache_dir,
-        )
+        self._map = None
         """
         - 模拟器地图对象
         - 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,
+            )
+            self._create_poi_id_2_aoi_id()
 
         self.time: int = 0
         """
@@ -109,19 +158,32 @@ class Simulator:
         self.map_y_gap = None
         self._bbox: tuple[float, float, float, float] = (-1, -1, -1, -1)
         self._lock = asyncio.Lock()
-        # poi id dict
+        self._environment_prompt: dict[str, str] = {}
+        self._log_list = []
+
+    def set_map(self, map: ray.ObjectRef):
+        self._map = map
+        self._create_poi_id_2_aoi_id()
+
+    def _create_poi_id_2_aoi_id(self):
+        pois = ray.get(self._map.get_poi.remote())  # type:ignore
         self.poi_id_2_aoi_id: dict[int, int] = {
-            poi["id"]: poi["aoi_id"] for _, poi in self.map.pois.items()
+            poi["id"]: poi["aoi_id"] for poi in pois
         }
-        self._environment_prompt:dict[str, str] = {}
-        self._log_list = []
+
+    @property
+    def map(self):
+        return self._map
 
     def get_log_list(self):
         return self._log_list
-    
+
     def clear_log_list(self):
         self._log_list = []
 
+    def get_poi_cate(self):
+        return self.poi_cate
+
     @property
     def environment(self) -> dict[str, str]:
         """
@@ -129,6 +191,9 @@ class Simulator:
         """
         return self._environment_prompt
 
+    def get_server_addr(self):
+        return self.server_addr
+
     def set_environment(self, environment: dict[str, str]):
         """
         Set the entire dictionary of environment variables.
@@ -177,11 +242,7 @@ class Simulator:
               Refer to https://cityproto.sim.fiblab.net/#city.person.1.GetPersonByLongLatBBoxResponse.
         """
         start_time = time.time()
-        log = {
-            "req": "find_agents_by_area",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "find_agents_by_area", "start_time": start_time, "consumption": 0}
         loop = asyncio.get_event_loop()
         resp = loop.run_until_complete(
             self._client.person_service.GetPersonByLongLatBBox(req=req)
@@ -216,18 +277,16 @@ class Simulator:
             - `List[str]`: A list of unique POI category names.
         """
         start_time = time.time()
-        log = {
-            "req": "get_poi_categories",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "get_poi_categories", "start_time": start_time, "consumption": 0}
         categories: list[str] = []
         if center is None:
             center = (0, 0)
-        _pois: list[dict] = self.map.query_pois(  # type:ignore
-            center=center,
-            radius=radius,
-            return_distance=False,
+        _pois: list[dict] = ray.get(
+            self.map.query_pois.remote(  # type:ignore
+                center=center,
+                radius=radius,
+                return_distance=False,
+            )
         )
         for poi in _pois:
             catg = poi["category"]
@@ -252,11 +311,7 @@ class Simulator:
             - `Union[int, str]`: The current simulation time either as an integer representing seconds since midnight or as a formatted string.
         """
         start_time = time.time()
-        log = {
-            "req": "get_time",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "get_time", "start_time": start_time, "consumption": 0}
         now = await self._client.clock_service.Now({})
         now = cast(dict[str, int], now)
         self.time = now["t"]
@@ -280,11 +335,7 @@ class Simulator:
         This method sends a request to the simulator's pause service to pause the simulation.
         """
         start_time = time.time()
-        log = {
-            "req": "pause",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "pause", "start_time": start_time, "consumption": 0}
         await self._client.pause_service.pause()
         log["consumption"] = time.time() - start_time
         self._log_list.append(log)
@@ -296,11 +347,7 @@ class Simulator:
         This method sends a request to the simulator's pause service to resume the simulation.
         """
         start_time = time.time()
-        log = {
-            "req": "resume",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "resume", "start_time": start_time, "consumption": 0}
         await self._client.pause_service.resume()
         log["consumption"] = time.time() - start_time
         self._log_list.append(log)
@@ -313,11 +360,7 @@ class Simulator:
             - `int`: The day number since the start of the simulation.
         """
         start_time = time.time()
-        log = {
-            "req": "get_simulator_day",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        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 = now["day"]
@@ -336,7 +379,7 @@ class Simulator:
         log = {
             "req": "get_simulator_second_from_start_of_day",
             "start_time": start_time,
-            "consumption": 0
+            "consumption": 0,
         }
         now = await self._client.clock_service.Now({})
         now = cast(dict[str, int], now)
@@ -355,43 +398,35 @@ class Simulator:
             - `Dict`: Information about the specified person.
         """
         start_time = time.time()
-        log = {
-            "req": "get_person",
-            "start_time": start_time,
-            "consumption": 0
-        }
-        person = await self._client.person_service.GetPerson(
+        log = {"req": "get_person", "start_time": start_time, "consumption": 0}
+        person: dict = await self._client.person_service.GetPerson(
             req={"person_id": person_id}
         )  # type:ignore
         log["consumption"] = time.time() - start_time
         self._log_list.append(log)
         return person
 
-    async def add_person(self, person: Any) -> dict:
+    async def add_person(self, dict_person: dict) -> dict:
         """
         Add a new person to the simulation.
 
         - **Args**:
-            - `person` (`Any`): The person object to add. If it's an instance of `person_pb2.Person`,
-              it will be wrapped in an `AddPersonRequest`. Otherwise, `person` is expected to already be a valid request object.
+            - `dict_person` (`dict`): The person object to add.
 
         - **Returns**:
             - `Dict`: Response from adding the person.
         """
         start_time = time.time()
-        log = {
-            "req": "add_person",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "add_person", "start_time": start_time, "consumption": 0}
+        person = dict2pb(dict_person, person_pb2.Person())
         if isinstance(person, person_pb2.Person):
             req = person_service.AddPersonRequest(person=person)
         else:
             req = person
-        person_id = await self._client.person_service.AddPerson(req)  # type:ignore
+        resp: dict = await self._client.person_service.AddPerson(req)  # type:ignore
         log["consumption"] = time.time() - start_time
         self._log_list.append(log)
-        return person_id
+        return resp
 
     async def set_aoi_schedules(
         self,
@@ -411,15 +446,11 @@ class Simulator:
               A list of AOI or POI IDs or tuples of (AOI ID, POI ID) that the person will visit.
             - `departure_times` (`Optional[List[float]]`): Departure times for each trip in the schedule.
               If not provided, current time will be used for all trips.
-            - `modes` (`Optional[List[TripMode]]`): Travel modes for each trip.
+            - `modes` (`Optional[List[int]]`): Travel modes for each trip.
               Defaults to `TRIP_MODE_DRIVE_ONLY` if not specified.
         """
         start_time = time.time()
-        log = {
-            "req": "set_aoi_schedules",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "set_aoi_schedules", "start_time": start_time, "consumption": 0}
         cur_time = float(await self.get_time())
         if not isinstance(target_positions, list):
             target_positions = [target_positions]
@@ -494,7 +525,7 @@ class Simulator:
         log = {
             "req": "reset_person_position",
             "start_time": start_time,
-            "consumption": 0
+            "consumption": 0,
         }
         reset_position = {}
         if aoi_id is not None:
@@ -545,11 +576,7 @@ class Simulator:
             - `List[Dict]`: A list of dictionaries containing information about the POIs found.
         """
         start_time = time.time()
-        log = {
-            "req": "get_around_poi",
-            "start_time": start_time,
-            "consumption": 0
-        }
+        log = {"req": "get_around_poi", "start_time": start_time, "consumption": 0}
         if isinstance(poi_type, str):
             poi_type = [poi_type]
         transformed_poi_type: list[str] = []
@@ -560,10 +587,12 @@ class Simulator:
                 transformed_poi_type += self.poi_cate[t]
         poi_type_set = set(transformed_poi_type)
         # 获取半径内的poi
-        _pois: list[dict] = self.map.query_pois(  # type:ignore
-            center=center,
-            radius=radius,
-            return_distance=False,
+        _pois: list[dict] = ray.get(
+            self.map.query_pois.remote(  # type:ignore
+                center=center,
+                radius=radius,
+                return_distance=False,
+            )
         )
         # 过滤掉不满足类别前缀的poi
         pois = []
@@ -574,4 +603,4 @@ class Simulator:
             pois.append(poi)
         log["consumption"] = time.time() - start_time
         self._log_list.append(log)
-        return pois
+        return pois

pycityagent/llm/llm.py

@@ -232,111 +232,112 @@ class LLM:
         """
         start_time = time.time()
         log = {"request_time": start_time}
-        if (
-            self.config.text["request_type"] == "openai"
-            or self.config.text["request_type"] == "deepseek"
-            or self.config.text["request_type"] == "qwen"
-        ):
-            for attempt in range(retries):
-                try:
-                    client = self._get_next_client()
-                    response = await client.chat.completions.create(
-                        model=self.config.text["model"],
-                        messages=dialog,
-                        temperature=temperature,
-                        max_tokens=max_tokens,
-                        top_p=top_p,
-                        frequency_penalty=frequency_penalty,  # type: ignore
-                        presence_penalty=presence_penalty,  # type: ignore
-                        stream=False,
-                        timeout=timeout,
-                        tools=tools,
-                        tool_choice=tool_choice,
-                    )  # type: ignore
-                    self._client_usage[self._current_client_index]["prompt_tokens"] += response.usage.prompt_tokens  # type: ignore
-                    self._client_usage[self._current_client_index]["completion_tokens"] += response.usage.completion_tokens  # type: ignore
-                    self._client_usage[self._current_client_index]["request_number"] += 1
-                    end_time = time.time()
-                    log["consumption"] = end_time - start_time
-                    log["input_tokens"] = response.usage.prompt_tokens
-                    log["output_tokens"] = response.usage.completion_tokens
-                    self._log_list.append(log)
-                    if tools and response.choices[0].message.tool_calls:
-                        return json.loads(
-                            response.choices[0]
-                            .message.tool_calls[0]
-                            .function.arguments
-                        )
-                    else:
-                        return response.choices[0].message.content
-                except APIConnectionError as e:
-                    print("API connection error:", e)
-                    if attempt < retries - 1:
-                        await asyncio.sleep(2**attempt)
-                    else:
-                        raise e
-                except OpenAIError as e:
-                    if hasattr(e, "http_status"):
-                        print(f"HTTP status code: {e.http_status}")  # type: ignore
-                    else:
-                        print("An error occurred:", e)
-                    if attempt < retries - 1:
-                        await asyncio.sleep(2**attempt)
-                    else:
-                        raise e
-        elif self.config.text["request_type"] == "zhipuai":
-            for attempt in range(retries):
-                try:
-                    client = self._get_next_client()
-                    response = client.chat.asyncCompletions.create(  # type: ignore
-                        model=self.config.text["model"],
-                        messages=dialog,
-                        temperature=temperature,
-                        top_p=top_p,
-                        timeout=timeout,
-                        tools=tools,
-                        tool_choice=tool_choice,
-                    )
-                    task_id = response.id
-                    task_status = ""
-                    get_cnt = 0
-                    cnt_threshold = int(timeout / 0.5)
-                    while (
-                        task_status != "SUCCESS"
-                        and task_status != "FAILED"
-                        and get_cnt <= cnt_threshold
-                    ):
-                        result_response = client.chat.asyncCompletions.retrieve_completion_result(id=task_id)  # type: ignore
-                        task_status = result_response.task_status
-                        await asyncio.sleep(0.5)
-                        get_cnt += 1
-                    if task_status != "SUCCESS":
-                        raise Exception(f"Task failed with status: {task_status}")
-
-                    self._client_usage[self._current_client_index]["prompt_tokens"] += result_response.usage.prompt_tokens  # type: ignore
-                    self._client_usage[self._current_client_index]["completion_tokens"] += result_response.usage.completion_tokens  # type: ignore
-                    self._client_usage[self._current_client_index]["request_number"] += 1
-                    end_time = time.time()
-                    log["used_time"] = end_time - start_time
-                    log["token_consumption"] = result_response.usage.prompt_tokens + result_response.usage.completion_tokens
-                    self._log_list.append(log)
-                    if tools and result_response.choices[0].message.tool_calls:  # type: ignore
-                        return json.loads(
-                            result_response.choices[0]  # type: ignore
-                            .message.tool_calls[0]
-                            .function.arguments
+        async with self.semaphore:
+            if (
+                self.config.text["request_type"] == "openai"
+                or self.config.text["request_type"] == "deepseek"
+                or self.config.text["request_type"] == "qwen"
+            ):
+                for attempt in range(retries):
+                    try:
+                        client = self._get_next_client()
+                        response = await client.chat.completions.create(
+                            model=self.config.text["model"],
+                            messages=dialog,
+                            temperature=temperature,
+                            max_tokens=max_tokens,
+                            top_p=top_p,
+                            frequency_penalty=frequency_penalty,  # type: ignore
+                            presence_penalty=presence_penalty,  # type: ignore
+                            stream=False,
+                            timeout=timeout,
+                            tools=tools,
+                            tool_choice=tool_choice,
+                        )  # type: ignore
+                        self._client_usage[self._current_client_index]["prompt_tokens"] += response.usage.prompt_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["completion_tokens"] += response.usage.completion_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["request_number"] += 1
+                        end_time = time.time()
+                        log["consumption"] = end_time - start_time
+                        log["input_tokens"] = response.usage.prompt_tokens
+                        log["output_tokens"] = response.usage.completion_tokens
+                        self._log_list.append(log)
+                        if tools and response.choices[0].message.tool_calls:
+                            return json.loads(
+                                response.choices[0]
+                                .message.tool_calls[0]
+                                .function.arguments
+                            )
+                        else:
+                            return response.choices[0].message.content
+                    except APIConnectionError as e:
+                        print("API connection error:", e)
+                        if attempt < retries - 1:
+                            await asyncio.sleep(2**attempt)
+                        else:
+                            raise e
+                    except OpenAIError as e:
+                        if hasattr(e, "http_status"):
+                            print(f"HTTP status code: {e.http_status}")  # type: ignore
+                        else:
+                            print("An error occurred:", e)
+                        if attempt < retries - 1:
+                            await asyncio.sleep(2**attempt)
+                        else:
+                            raise e
+            elif self.config.text["request_type"] == "zhipuai":
+                for attempt in range(retries):
+                    try:
+                        client = self._get_next_client()
+                        response = client.chat.asyncCompletions.create(  # type: ignore
+                            model=self.config.text["model"],
+                            messages=dialog,
+                            temperature=temperature,
+                            top_p=top_p,
+                            timeout=timeout,
+                            tools=tools,
+                            tool_choice=tool_choice,
                         )
-                    else:
-                        return result_response.choices[0].message.content  # type: ignore
-                except APIConnectionError as e:
-                    print("API connection error:", e)
-                    if attempt < retries - 1:
-                        await asyncio.sleep(2**attempt)
-                    else:
-                        raise e
-        else:
-            print("ERROR: Wrong Config")
-            return "wrong config"
+                        task_id = response.id
+                        task_status = ""
+                        get_cnt = 0
+                        cnt_threshold = int(timeout / 0.5)
+                        while (
+                            task_status != "SUCCESS"
+                            and task_status != "FAILED"
+                            and get_cnt <= cnt_threshold
+                        ):
+                            result_response = client.chat.asyncCompletions.retrieve_completion_result(id=task_id)  # type: ignore
+                            task_status = result_response.task_status
+                            await asyncio.sleep(0.5)
+                            get_cnt += 1
+                        if task_status != "SUCCESS":
+                            raise Exception(f"Task failed with status: {task_status}")
+
+                        self._client_usage[self._current_client_index]["prompt_tokens"] += result_response.usage.prompt_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["completion_tokens"] += result_response.usage.completion_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["request_number"] += 1
+                        end_time = time.time()
+                        log["used_time"] = end_time - start_time
+                        log["token_consumption"] = result_response.usage.prompt_tokens + result_response.usage.completion_tokens
+                        self._log_list.append(log)
+                        if tools and result_response.choices[0].message.tool_calls:  # type: ignore
+                            return json.loads(
+                                result_response.choices[0]  # type: ignore
+                                .message.tool_calls[0]
+                                .function.arguments
+                            )
+                        else:
+                            return result_response.choices[0].message.content  # type: ignore
+                    except APIConnectionError as e:
+                        print("API connection error:", e)
+                        if attempt < retries - 1:
+                            await asyncio.sleep(2**attempt)
+                        else:
+                            raise e
+            else:
+                print("ERROR: Wrong Config")
+                return "wrong config"
 
     async def img_understand(
         self, img_path: Union[str, list[str]], prompt: Optional[str] = None

pycityagent/memory/memory.py

@@ -341,7 +341,7 @@ class StreamMemory:
         """获取指定ID的记忆"""
         memories = [memory for memory in self._memories if memory.id in memory_ids]
         sorted_results = sorted(memories, key=lambda x: (x.day, x.t), reverse=True)
-        return self.format_memory(sorted_results)
+        return await self.format_memory(sorted_results)
 
     async def search(
         self,