pycityagent 2.0.0a1__tar.gz → 2.0.0a3__tar.gz

{pycityagent-2.0.0a1 → pycityagent-2.0.0a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a1
+Version: 2.0.0a3
 Summary: LLM-based城市环境agent构建库
 License: MIT
 Author: Yuwei Yan

{pycityagent-2.0.0a1 → pycityagent-2.0.0a3}/pycityagent/agent.py

@@ -126,7 +126,7 @@ class Agent(ABC):
         dialog = []
 
         # 添加系统提示
-        system_prompt = f"你是一个名叫{self._name}的{self._type.value}。请以第一人称的方式回答问题,保持回答简洁明了。"
+        system_prompt = f"请以第一人称的方式回答问题,保持回答简洁明了。"
         dialog.append({"role": "system", "content": system_prompt})
 
         # 添加记忆上下文

{pycityagent-2.0.0a1 → pycityagent-2.0.0a3}/pycityagent/environment/sim/__init__.py

@@ -12,6 +12,7 @@ from .lane_service import LaneService
 from .road_service import RoadService
 from .social_service import SocialService
 from .light_service import LightService
+from .sim_env import ControlSimEnv
 
 __all__ = [
     "CityClient",
@@ -24,4 +25,5 @@ __all__ = [
     "EconomyPersonService",
     "EconomyOrgService",
     "LightService",
+    "ControlSimEnv",
 ]

pycityagent-2.0.0a3/pycityagent/environment/sim/sim_env.py

@@ -0,0 +1,145 @@
+import logging
+import os
+import time
+import atexit
+from subprocess import DEVNULL, Popen
+from typing import Optional
+
+from pycitydata.map import Map
+
+from ..utils import encode_to_base64, find_free_port
+
+__all__ = ["ControlSimEnv"]
+
+
+def _generate_yaml_config(map_file: str, start_step: int, total_step: int) -> str:
+    map_file = os.path.abspath(map_file)
+    return f"""
+input:
+  # 地图
+  map:
+    file: "{map_file}"
+
+control:
+  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:
+"""
+
+
+class ControlSimEnv:
+    def __init__(
+        self,
+        task_name: str,
+        map_file: str,
+        # person_file,
+        start_step: int,
+        total_step: int,
+        log_dir: str,
+        min_step_time: int = 1000,
+        timeout: int = 5,
+        simuletgo_addr: Optional[str] = None,
+    ):
+        self._task_name = task_name
+        self._map_file = map_file
+        # self._person_file = person_file
+        self._start_step = start_step
+        self._total_step = total_step
+        self._log_dir = log_dir
+        self._min_step_time = min_step_time
+        self._timeout = timeout
+        self.m = Map(pb_path=map_file)
+        """
+        地图数据
+        """
+
+        # 检查二进制文件是否存在
+        # simulet-go: ~/.local/bin/simulet-go
+        self._simuletgo_path = os.path.expanduser("~/.local/bin/simulet-go")
+        if not os.path.exists(os.path.expanduser(self._simuletgo_path)):
+            raise FileNotFoundError("simulet-go not found, please install it first")
+
+        self._simuletgo_config = _generate_yaml_config(map_file, start_step, total_step)
+        self.simuletgo_port = None
+        self._simuletgo_proc = None
+        self._traffic_client = None
+
+        os.makedirs(log_dir, exist_ok=True)
+
+        self.simuletgo_addr = self.reset(simuletgo_addr)
+
+    def reset(
+        self,
+        simuletgo_addr: Optional[str] = None,
+    ):
+        """
+        Args:
+        - simuletgo_addr: str, simulet-go的地址。如果为None，则启动一个新的simulet-go
+        """
+
+        # 三个地址必须同时为None或者同时不为None
+        if simuletgo_addr is None:
+            # 1. 启动simulet-go
+            # ./simulet-go -config-data configbase64 -job test -syncer http://localhost:53001 -listen :51102
+            assert self.simuletgo_port is None
+            assert self._simuletgo_proc is None
+            self.simuletgo_port = find_free_port()
+            config_base64 = encode_to_base64(self._simuletgo_config)
+            self._simuletgo_proc = Popen(
+                [
+                    self._simuletgo_path,
+                    "-config-data",
+                    config_base64,
+                    "-job",
+                    self._task_name,
+                    "-listen",
+                    f":{self.simuletgo_port}",
+                    "-run.min_step_time",
+                    f"{self._min_step_time}",
+                    "-output",
+                    self._log_dir,
+                    "-cache",
+                    "",
+                    "-log.level",
+                    "error",
+                ],
+                # 忽略输出
+                # stdout=DEVNULL,
+            )
+            logging.info(
+                f"start simulet-go at localhost:{self.simuletgo_port}, PID={self._simuletgo_proc.pid}"
+            )
+            simuletgo_addr = f"http://localhost:{self.simuletgo_port}"
+            atexit.register(self.close)
+            time.sleep(1)
+        elif simuletgo_addr is not None:
+            pass
+        else:
+            # raise ValueError(
+            #     "simuletgo_addr, syncer_addr, routing_addr must be all None or all not None"
+            # )
+            pass
+        return simuletgo_addr
+
+    def close(self):
+        if self._simuletgo_proc is not None:
+            self._simuletgo_proc.terminate()
+            simuletgo_code = self._simuletgo_proc.wait()
+            logging.info(f"simulet-go exit with code {simuletgo_code}")
+
+        self.simuletgo_port = None
+        self._simuletgo_proc = None
+        self._traffic_client = None

{pycityagent-2.0.0a1 → pycityagent-2.0.0a3}/pycityagent/environment/simulator.py

@@ -2,15 +2,24 @@
 
 import asyncio
 import logging
+import os
+from collections import defaultdict
+from collections.abc import Sequence
 from datetime import datetime, timedelta
 from typing import Any, Optional, Tuple, Union, cast
 
 from mosstool.type import TripMode
+from mosstool.util.format_converter import coll2pb
 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
 from pycityproto.city.person.v2 import person_service_pb2 as person_service
+from pymongo import MongoClient
+from shapely.geometry import Point
+from shapely.strtree import STRtree
 
-from .sim import CityClient
+from .sim import CityClient, ControlSimEnv
+from .utils.const import *
 
 
 class Simulator:
@@ -25,18 +34,52 @@ 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())
+        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())
 
-        self._client = CityClient(self.config["simulator"]["server"], secure=secure)
-        """
-        - 模拟器grpc客户端
-        - grpc client of simulator
-        """
+        if 'simulator' in config:
+            if 'server' not in config["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),
+                    total_step=2147000000,
+                    log_dir=config["simulator"].get("log_dir", "./log"),
+                    min_step_time=config["simulator"].get("min_step_time", 1000),
+                    simuletgo_addr=config["simulator"].get("server", None),
+                )
+
+                # using local client
+                self._client = CityClient(sim_env.simuletgo_addr, secure=False)
+                """
+                - 模拟器grpc客户端
+                - grpc client of simulator
+                """
+            else:
+                self._client = CityClient(config['simulator']['server'], secure=False)
+        else:
+            logging.warning("No simulator config found, no simulator client will be used")
 
         self.map = SimMap(
-            mongo_uri=config["map_request"]["mongo_uri"],
-            mongo_db=config["map_request"]["mongo_db"],
-            mongo_coll=config["map_request"]["mongo_coll"],
-            cache_dir=config["map_request"]["cache_dir"],
+            mongo_uri=_mongo_uri,
+            mongo_db=_mongo_db,
+            mongo_coll=_mongo_coll,
+            cache_dir=_map_cache_dir,
         )
         """
         - 模拟器地图对象
@@ -54,24 +97,14 @@ class Simulator:
         - 模拟城市当前时间
         - The current time of simulator
         """
-        self.poi_cate = {
-            "10": "eat",
-            "13": "shopping",
-            "18": "sports",
-            "22": "excursion",
-            "16": "entertainment",
-            "20": "medical tratment",
-            "14": "trivialities",
-            "25": "financial",
-            "12": "government and political services",
-            "23": "cultural institutions",
-            "28": "residence",
-        }
+        self.poi_cate = POI_CATG_DICT
         self.map_x_gap = None
         self.map_y_gap = None
-        self._bbox: Tuple[float, float, float, float] = (-1, -1, -1, -1)
+        self._bbox: tuple[float, float, float, float] = (-1, -1, -1, -1)
         self.poi_matrix_centers = []
         self._lock = asyncio.Lock()
+        # poi STRtree
+        self.set_poi_tree()
 
     # * Agent相关
     def FindAgentsByArea(self, req: dict, status=None):
@@ -101,146 +134,36 @@ class Simulator:
             resp.motions = motions  # type: ignore
             return resp
 
-    def set_poi_matrix(
-        self, row_number: int = 12, col_number: int = 10, radius: int = 10000
-    ):
-        """
-        初始化pois_matrix
-
-        Args:
-        - map (dict): 地图参数
-            east, west, north, south
-        - row_number (int): 行数
-        - col_number (int): 列数
-        - radius (int): 搜索半径, 单位m
-        """
-        self.matrix_map = self.map
-        map_header = self.matrix_map.header
-        min_x, min_y, max_x, max_y = (
-            map_header[k] for k in ("west", "south", "east", "north")
-        )
-        self._bbox = (min_x, min_y, max_x, max_y)
-        logging.info(
-            f"Building Poi searching matrix, Row_number: {row_number}, Col_number: {col_number}, Radius: {radius}m"
-        )
-        self.map_x_gap = map_x_gap = (max_x - min_x) / col_number
-        self.map_y_gap = map_y_gap = (max_y - min_y) / row_number
-        self.poi_matrix_centers = poi_matrix_centers = [[] for _ in range(row_number)]
-        for i, i_centers in enumerate(poi_matrix_centers):
-            for j in range(col_number):
-                center_x = map_header["west"] + map_x_gap * j + map_x_gap / 2
-                center_y = map_header["south"] + map_y_gap * i + map_y_gap / 2
-                i_centers.append((center_x, center_y))
-
-        for pre, _ in self.poi_cate.items():
-            logging.info(f"Building matrix for Poi category: {pre}")
-            self.pois_matrix[pre] = []
-            for row_centers in self.poi_matrix_centers:
-                row_pois = []
-                for center in row_centers:
-                    pois = self.map.query_pois(
-                        center=center, radius=radius, category_prefix=pre
-                    )
-                    row_pois.append(pois)
-                self.pois_matrix[pre].append(row_pois)
-        logging.info("Finished")
-
-    def get_pois_from_matrix(self, center: Tuple[float, float], prefix: str):
-        """
-        从poi搜索矩阵中快速获取poi
-
-        Args:
-        - center (Tuple[float, float]): 位置信息
-        - prefix (str): 类型前缀
-        """
-        (min_x, min_y, max_x, max_y) = self._bbox
-        _x, _y = center
-        if self.map_x_gap is None or self.map_y_gap is None:
-            logging.warning("Set Poi Matrix first")
-            return
-        elif prefix not in self.poi_cate:
-            logging.warning(f"Wrong prefix, only {self.poi_cate.keys()} is usable")
-            return
-        elif not (min_x < _x < max_x) or not (min_y < _y < max_y):
-            logging.warning("Wrong center")
-            return
-
-        # 矩阵匹配
-        rows = int((_y - min_y) / self.map_y_gap)
-        cols = int((_x - min_x) / self.map_x_gap)
-        pois = self.pois_matrix[prefix][rows][cols]
-        return pois
-
-    def get_cat_from_pois(self, pois: list):
-        cat_2_num = {}
-        for poi in pois:
-            cate = poi["category"][:2]
-            if cate not in self.poi_cate:
-                continue
-            if cate in cat_2_num:
-                cat_2_num[cate] += 1
-            else:
-                cat_2_num[cate] = 1
-        max_cat = ""
-        max_num = 0
-        for key in cat_2_num.keys():
-            if cat_2_num[key] > max_num:
-                max_num = cat_2_num[key]
-                max_cat = self.poi_cate[key]
-        return max_cat
-
-    def get_poi_matrix_in_rec(
+    def set_poi_tree(
         self,
-        center: Tuple[float, float],
-        radius: int = 2500,
-        rows: int = 5,
-        cols: int = 5,
     ):
         """
-        获取以center为中心的正方形区域内的poi集合
-
-        Args:
-        - center (Tuple[float, float]): 中心位置点
-        - radius (int): 半径
+        初始化pois_tree
         """
-        north = center[1] + radius
-        south = center[1] - radius
-        west = center[0] - radius
-        east = center[0] + radius
-        x_gap = (east - west) / cols
-        y_gap = (north - south) / rows
-        matrix = []
-        for i in range(rows):
-            matrix.append([])
-            for j in range(cols):
-                matrix[i].append([])
-        for poi in self.map.pois.values():
-            x = poi["position"]["x"]
-            y = poi["position"]["y"]
-            if west < x < east and south < y < north:
-                row_index = int((y - south) / x_gap)
-                col_index = int((x - west) / y_gap)
-                matrix[row_index][col_index].append(poi)
-        matrix_type = []
-        for i in range(rows):
-            for j in range(cols):
-                matrix_type.append(self.get_cat_from_pois(matrix[i][j]))
-        poi_total_number = []
-        poi_type_number = []
-        for i in range(rows):
-            for j in range(cols):
-                poi_total_number.append(len(matrix[i][j]))
-                number = 0
-                for poi in matrix[i][j]:
-                    if (
-                        poi["category"][:2] in self.poi_cate.keys()
-                        and self.poi_cate[poi["category"][:2]]
-                        == matrix_type[i * cols + j]
-                    ):
-                        number += 1
-                poi_type_number.append(number)
+        poi_geos = []
+        tree_id_2_poi_and_catg: dict[int, tuple[dict, str]] = {}
+        for tree_id, poi in enumerate(self.map.pois.values()):
+            tree_id_2_poi_and_catg[tree_id] = (poi, poi["category"])
+            poi_geos.append(Point([poi["position"][k] for k in ["x", "y"]]))
+        self.tree_id_2_poi_and_catg = tree_id_2_poi_and_catg
+        self.pois_tree = STRtree(poi_geos)
 
-        return matrix, matrix_type, poi_total_number, poi_type_number
+    def GetPoiCategories(
+        self,
+        center: Optional[Union[tuple[float, float], Point]] = None,
+        radius: Optional[float] = None,
+    ) -> list[str]:
+        if center is not None and radius is not None:
+            if not isinstance(center, Point):
+                center = Point(center)
+            indices = self.pois_tree.query(center.buffer(radius))
+        else:
+            indices = list(self.tree_id_2_poi_and_catg.keys())
+        categories = []
+        for index in indices:
+            _, catg = self.tree_id_2_poi_and_catg[index]
+            categories.append(catg.split("|")[-1])
+        return list(set(categories))
 
     async def GetTime(
         self, format_time: bool = False, format: str = "%H:%M:%S"
@@ -367,3 +290,31 @@ class Simulator:
             logging.debug(
                 f"Neither aoi or lane pos provided for person {person_id} position reset!!"
             )
+
+    def GetAroundPoi(
+        self,
+        center: Union[tuple[float, float], Point],
+        radius: float,
+        poi_type: Union[str, list[str]],
+    ):
+        if not isinstance(poi_type, Sequence):
+            poi_type = [poi_type]
+        transformed_poi_type = []
+        for t in poi_type:
+            if t not in self.poi_cate:
+                transformed_poi_type.append(t)
+            else:
+                transformed_poi_type += self.poi_cate[t]
+        poi_type_set = set(transformed_poi_type)
+        if not isinstance(center, Point):
+            center = Point(center)
+        # 获取半径内的poi
+        indices = self.pois_tree.query(center.buffer(radius))
+        # 过滤掉不满足类别前缀的poi
+        pois = []
+        for index in indices:
+            poi, catg = self.tree_id_2_poi_and_catg[index]
+            if catg.split("|")[-1] not in poi_type_set:
+                continue
+            pois.append(poi)
+        return pois

pycityagent-2.0.0a3/pycityagent/environment/utils/__init__.py

@@ -0,0 +1,10 @@
+"""
+实用工具
+utilities
+"""
+
+from .geojson import wrap_feature_collection
+from .base64 import encode_to_base64
+from .port import find_free_port
+
+__all__ = ["wrap_feature_collection","find_free_port","find_free_port",]

pycityagent-2.0.0a3/pycityagent/environment/utils/base64.py

@@ -0,0 +1,16 @@
+import base64
+
+__all__ = ["encode_to_base64"]
+
+
+def encode_to_base64(input_string: str) -> str:
+    # 将字符串转换为字节
+    input_bytes = input_string.encode("utf-8")
+
+    # 将字节编码为Base64
+    base64_bytes = base64.b64encode(input_bytes)
+
+    # 将Base64字节转换为字符串
+    base64_string = base64_bytes.decode("utf-8")
+
+    return base64_string