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

pycityagent/environment/utils/map_utils.py

@@ -0,0 +1,157 @@
+import math
+from typing import Dict, List, Literal, Optional, Tuple, Union
+
+import numpy as np
+
+
+def get_angle(x, y):
+    return math.atan2(y, x) * 180 / math.pi
+
+
+def point_on_line_given_distance(start_node, end_node, distance):
+    """
+    Given two points (start_point and end_point) defining a line, and a distance s to travel along the line,
+    return the coordinates of the point reached after traveling s units along the line, starting from start_point.
+
+    Args:
+        start_point (tuple): Tuple of (x, y) representing the starting point on the line.
+        end_point (tuple): Tuple of (x, y) representing the ending point on the line.
+        distance (float): Distance to travel along the line, starting from start_point.
+
+    Returns:
+        tuple: Tuple of (x, y) representing the new point reached after traveling s units along the line.
+    """
+
+    x1, y1 = start_node["x"], start_node["y"]
+    x2, y2 = end_node["x"], end_node["y"]
+
+    # Calculate the slope m and the y-intercept b of the line
+    if x1 == x2:
+        # Vertical line, distance is only along the y-axis
+        return (x1, y1 + distance if distance >= 0 else y1 - abs(distance))
+    else:
+        m = (y2 - y1) / (x2 - x1)
+        b = y1 - m * x1
+
+        # Calculate the direction vector (dx, dy) along the line
+        dx = (x2 - x1) / math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+        dy = (y2 - y1) / math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
+
+        # Scale the direction vector by the given distance
+        scaled_dx = dx * distance
+        scaled_dy = dy * distance
+
+        # Calculate the new point's coordinates
+        x = x1 + scaled_dx
+        y = y1 + scaled_dy
+
+        return [x, y]
+
+
+def get_key_index_in_lane(
+    nodes: List[Dict[str, float]],
+    distance: float,
+    direction: Union[Literal["front"], Literal["back"]],
+) -> int:
+    if direction == "front":
+        _nodes = [n for n in nodes]
+        _index_offset, _index_factor = 0, 1
+    elif direction == "back":
+        _nodes = [n for n in nodes[::-1]]
+        _index_offset, _index_factor = len(_nodes) - 1, -1
+    else:
+        raise ValueError(f"Invalid direction type {direction}!")
+    _lane_points: List[Tuple[float, float, float]] = [
+        (n["x"], n["y"], n.get("z", 0)) for n in _nodes
+    ]
+    _line_lengths: List[float] = [0.0 for _ in range(len(_nodes))]
+    _s = 0.0
+    for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
+        _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])
+        _line_lengths[i + 1] = _s
+    s = np.clip(distance, _line_lengths[0], _line_lengths[-1])
+    _key_index = 0
+    for (
+        prev_s,
+        cur_s,
+    ) in zip(_line_lengths[:-1], _line_lengths[1:]):
+        if prev_s <= s < cur_s:
+            break
+        _key_index += 1
+    return _index_offset + _index_factor * _key_index
+
+
+def get_xy_in_lane(
+    nodes: List[Dict[str, float]],
+    distance: float,
+    direction: Union[Literal["front"], Literal["back"]],
+) -> Tuple[float, float]:
+    if direction == "front":
+        _nodes = [n for n in nodes]
+    elif direction == "back":
+        _nodes = [n for n in nodes[::-1]]
+    else:
+        raise ValueError(f"Invalid direction type {direction}!")
+    _lane_points: List[Tuple[float, float, float]] = [
+        (n["x"], n["y"], n.get("z", 0)) for n in _nodes
+    ]
+    _line_lengths: List[float] = [0.0 for _ in range(len(_nodes))]
+    _s = 0.0
+    for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
+        _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])
+        _line_lengths[i + 1] = _s
+    s = np.clip(distance, _line_lengths[0], _line_lengths[-1])
+    for prev_s, prev_p, cur_s, cur_p in zip(
+        _line_lengths[:-1],
+        _lane_points[:-1],
+        _line_lengths[1:],
+        _lane_points[1:],
+    ):
+        if prev_s <= s < cur_s:
+            _delta_x, _delta_y, _delta_z = [
+                cur_p[_idx] - prev_p[_idx] for _idx in [0, 1, 2]
+            ]
+            _blend_x, _blend_y, _blend_z = [prev_p[_idx] for _idx in [0, 1, 2]]
+            _ratio = (s - prev_s) / (cur_s - prev_s)
+            return (
+                _blend_x + _ratio * _delta_x,
+                _blend_y + _ratio * _delta_y,
+                _blend_z + _ratio * _delta_z,
+            )[:2]
+    return _lane_points[-1][:2]
+
+
+def get_direction_by_s(
+    nodes: List[Dict[str, float]],
+    distance: float,
+    direction: Union[Literal["front"], Literal["back"]],
+) -> float:
+    if direction == "front":
+        _nodes = [n for n in nodes]
+    elif direction == "back":
+        _nodes = [n for n in nodes[::-1]]
+    else:
+        raise ValueError(f"Invalid direction type {direction}!")
+    _lane_points: List[Tuple[float, float, float]] = [
+        (n["x"], n["y"], n.get("z", 0)) for n in _nodes
+    ]
+    _line_lengths: List[float] = [0.0 for _ in range(len(_nodes))]
+    _line_directions: List[Tuple[float, float]] = []
+    _s = 0.0
+    for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
+        _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])
+        _line_lengths[i + 1] = _s
+    for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
+        _direction = math.atan2(next_p[1] - cur_p[1], next_p[0] - cur_p[0])
+        _pitch = math.atan2(
+            next_p[2] - cur_p[2],
+            math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1]),
+        )
+        _line_directions.append((_direction / math.pi * 180, _pitch / math.pi * 180))
+    s = np.clip(distance, _line_lengths[0], _line_lengths[-1])
+    for prev_s, cur_s, direcs in zip(
+        _line_lengths[:-1], _line_lengths[1:], _line_directions
+    ):
+        if prev_s <= s < cur_s:
+            return direcs[0]
+    return _line_directions[-1][0]

pycityagent/environment/utils/port.py

@@ -0,0 +1,11 @@
+import socket
+from contextlib import closing
+
+__all__ = ["find_free_port"]
+
+
+def find_free_port():
+    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
+        s.bind(("", 0))
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        return s.getsockname()[1]

pycityagent/environment/utils/protobuf.py

@@ -0,0 +1,39 @@
+from typing import Any, Awaitable, TypeVar, Union, Dict
+from google.protobuf.message import Message
+from google.protobuf.json_format import MessageToDict
+
+__all__ = ["parse", "async_parse"]
+
+T = TypeVar("T", bound=Message)
+
+
+def parse(res: T, dict_return: bool) -> Union[Dict[str, Any], T]:
+    """
+    将Protobuf返回值转换为dict或者原始值
+    Convert Protobuf return value to dict or original value
+    """
+    if dict_return:
+        return MessageToDict(
+            res,
+            including_default_value_fields=True,
+            preserving_proto_field_name=True,
+            use_integers_for_enums=True,
+        )
+    else:
+        return res
+
+
+async def async_parse(res: Awaitable[T], dict_return: bool) -> Union[Dict[str, Any], T]:
+    """
+    将Protobuf await返回值转换为dict或者原始值
+    Convert Protobuf await return value to dict or original value
+    """
+    if dict_return:
+        return MessageToDict(
+            await res,
+            including_default_value_fields=True,
+            preserving_proto_field_name=True,
+            use_integers_for_enums=True,
+        )
+    else:
+        return await res

pycityagent/llm/__init__.py

@@ -0,0 +1,6 @@
+"""LLM相关模块"""
+
+from .llm import LLM, LLMConfig
+from .embedding import SimpleEmbedding
+
+__all__ = ["LLM", "LLMConfig", "SimpleEmbedding"]

pycityagent/llm/embedding.py

@@ -0,0 +1,136 @@
+"""简单的基于内存的embedding实现"""
+
+import numpy as np
+from typing import List, Dict, Optional
+import hashlib
+import json
+
+class SimpleEmbedding:
+    """简单的基于内存的embedding实现
+    
+    使用简单的词袋模型(Bag of Words)和TF-IDF来生成文本的向量表示。
+    所有向量都保存在内存中，适用于小规模应用。
+    """
+    
+    def __init__(self, vector_dim: int = 128, cache_size: int = 1000):
+        """初始化
+        
+        Args:
+            vector_dim: 向量维度
+            cache_size: 缓存大小，超过此大小将清除最早的缓存
+        """
+        self.vector_dim = vector_dim
+        self.cache_size = cache_size
+        self._cache: Dict[str, np.ndarray] = {}
+        self._vocab: Dict[str, int] = {}  # 词汇表
+        self._idf: Dict[str, float] = {}  # 逆文档频率
+        self._doc_count = 0  # 文档总数
+        
+    def _text_to_hash(self, text: str) -> str:
+        """将文本转换为hash值"""
+        return hashlib.md5(text.encode()).hexdigest()
+        
+    def _tokenize(self, text: str) -> List[str]:
+        """简单的分词"""
+        # 这里使用简单的空格分词，实际应用中可以使用更复杂的分词方法
+        return text.lower().split()
+        
+    def _update_vocab(self, tokens: List[str]):
+        """更新词汇表"""
+        for token in set(tokens):  # 使用set去重
+            if token not in self._vocab:
+                self._vocab[token] = len(self._vocab)
+                
+    def _update_idf(self, tokens: List[str]):
+        """更新IDF值"""
+        self._doc_count += 1
+        unique_tokens = set(tokens)
+        for token in unique_tokens:
+            self._idf[token] = self._idf.get(token, 0) + 1
+            
+    def _calculate_tf(self, tokens: List[str]) -> Dict[str, float]:
+        """计算词频(TF)"""
+        tf = {}
+        total_tokens = len(tokens)
+        for token in tokens:
+            tf[token] = tf.get(token, 0) + 1
+        # 归一化
+        for token in tf:
+            tf[token] /= total_tokens
+        return tf
+        
+    def _calculate_tfidf(self, tokens: List[str]) -> np.ndarray:
+        """计算TF-IDF向量"""
+        vector = np.zeros(self.vector_dim)
+        tf = self._calculate_tf(tokens)
+        
+        for token, tf_value in tf.items():
+            if token in self._idf:
+                idf = np.log(self._doc_count / self._idf[token])
+                idx = self._vocab[token] % self.vector_dim  # 使用取模运算来控制向量维度
+                vector[idx] += tf_value * idf
+                
+        # L2归一化
+        norm = np.linalg.norm(vector)
+        if norm > 0:
+            vector /= norm
+            
+        return vector
+        
+    async def embed(self, text: str) -> np.ndarray:
+        """生成文本的向量表示
+        
+        Args:
+            text: 输入文本
+            
+        Returns:
+            np.ndarray: 文本的向量表示
+        """
+        # 检查缓存
+        text_hash = self._text_to_hash(text)
+        if text_hash in self._cache:
+            return self._cache[text_hash]
+            
+        # 分词
+        tokens = self._tokenize(text)
+        if not tokens:
+            return np.zeros(self.vector_dim)
+            
+        # 更新词汇表和IDF
+        self._update_vocab(tokens)
+        self._update_idf(tokens)
+        
+        # 计算向量
+        vector = self._calculate_tfidf(tokens)
+        
+        # 更新缓存
+        if len(self._cache) >= self.cache_size:
+            # 删除最早的缓存
+            oldest_key = next(iter(self._cache))
+            del self._cache[oldest_key]
+        self._cache[text_hash] = vector
+        
+        return vector
+        
+    def save(self, file_path: str):
+        """保存模型"""
+        state = {
+            'vector_dim': self.vector_dim,
+            'cache_size': self.cache_size,
+            'vocab': self._vocab,
+            'idf': self._idf,
+            'doc_count': self._doc_count
+        }
+        with open(file_path, 'w') as f:
+            json.dump(state, f)
+            
+    def load(self, file_path: str):
+        """加载模型"""
+        with open(file_path, 'r') as f:
+            state = json.load(f)
+        self.vector_dim = state['vector_dim']
+        self.cache_size = state['cache_size']
+        self._vocab = state['vocab']
+        self._idf = state['idf']
+        self._doc_count = state['doc_count']
+        self._cache = {}  # 清空缓存