pycityagent 2.0.0a19__py3-none-any.whl → 2.0.0a20__py3-none-any.whl

pycityagent/agent.py

@@ -8,7 +8,7 @@ from abc import ABC, abstractmethod
 from copy import deepcopy
 from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 from uuid import UUID
 
 import fastavro
@@ -55,7 +55,7 @@ class Agent(ABC):
         simulator: Optional[Simulator] = None,
         mlflow_client: Optional[MlflowClient] = None,
         memory: Optional[Memory] = None,
-        avro_file: Optional[Dict[str, str]] = None,
+        avro_file: Optional[dict[str, str]] = None,
     ) -> None:
         """
         Initialize the Agent.
@@ -69,7 +69,7 @@ class Agent(ABC):
             simulator (Simulator, optional): The simulator object. Defaults to None.
             mlflow_client (MlflowClient, optional): The Mlflow object. Defaults to None.
             memory (Memory, optional): The memory of the agent. Defaults to None.
-            avro_file (Dict[str, str], optional): The avro file of the agent. Defaults to None.
+            avro_file (dict[str, str], optional): The avro file of the agent. Defaults to None.
         """
         self._name = name
         self._type = type
@@ -85,7 +85,7 @@ class Agent(ABC):
         self._has_bound_to_simulator = False
         self._has_bound_to_economy = False
         self._blocked = False
-        self._interview_history: List[Dict] = []  # 存储采访历史
+        self._interview_history: list[dict] = []  # 存储采访历史
         self._person_template = PersonService.default_dict_person()
         self._avro_file = avro_file
 
@@ -137,7 +137,7 @@ class Agent(ABC):
         """
         self._exp_id = exp_id
 
-    def set_avro_file(self, avro_file: Dict[str, str]):
+    def set_avro_file(self, avro_file: dict[str, str]):
         """
         Set the avro file of the agent.
         """
@@ -557,7 +557,7 @@ class InstitutionAgent(Agent):
             avro_file=avro_file,
         )
         # 添加响应收集器
-        self._gather_responses: Dict[str, asyncio.Future] = {}
+        self._gather_responses: dict[str, asyncio.Future] = {}
 
     async def bind_to_simulator(self):
         await self._bind_to_economy()
@@ -659,7 +659,7 @@ class InstitutionAgent(Agent):
                 }
             )
 
-    async def gather_messages(self, agent_uuids: list[str], target: str) -> List[dict]:
+    async def gather_messages(self, agent_uuids: list[str], target: str) -> list[dict]:
         """从多个智能体收集消息
 
         Args:
@@ -667,7 +667,7 @@ class InstitutionAgent(Agent):
             target: 要收集的信息类型
 
         Returns:
-            List[dict]: 收集到的所有响应
+            list[dict]: 收集到的所有响应
         """
         # 为每个agent创建Future
         futures = {}

pycityagent/economy/econ_client.py

@@ -316,3 +316,40 @@ class EconomyClient:
             await self._aio_stub.GetOrgEntityIds(request)
         )
         return list(response.org_ids)
+
+    async def add_delta_value(
+        self,
+        id: int,
+        key: str,
+        value: Any,
+    ) -> Any:
+        """
+        Add key-value pair
+
+        Args:
+        - id (int): the id of `Org` or `Agent`.
+        - key (str): the attribute to update. Can only be `inventory`, `price`, `interest_rate` and `currency`
+
+
+        Returns:
+        - Any
+        """
+        pascal_key = _snake_to_pascal(key)
+        _request_type = getattr(org_service, f"Add{pascal_key}Request")
+        _request_func = getattr(self._aio_stub, f"Add{pascal_key}")
+        _available_keys = {
+            "inventory",
+            "price",
+            "interest_rate",
+            "currency",
+        }
+        if key not in _available_keys:
+            raise ValueError(f"Invalid key `{key}`, can only be {_available_keys}!")
+        return await _request_func(
+            _request_type(
+                **{
+                    "org_id": id,
+                    f"delta_{key}": value,
+                }
+            )
+        )

pycityagent/environment/utils/geojson.py

@@ -1,9 +1,7 @@
-from typing import List
-
 __all__ = ["wrap_feature_collection"]
 
 
-def wrap_feature_collection(features: List[dict], name: str):
+def wrap_feature_collection(features: list[dict], name: str):
     """
     将 GeoJSON Feature 集合包装为 FeatureCollection
     Wrap GeoJSON Feature collection as FeatureCollection

pycityagent/environment/utils/map_utils.py

@@ -1,5 +1,5 @@
 import math
-from typing import Dict, List, Literal, Optional, Tuple, Union
+from typing import Literal,  Union
 
 import numpy as np
 
@@ -14,12 +14,12 @@ def point_on_line_given_distance(start_node, end_node, distance):
     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.
+        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.
+        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"]
@@ -49,7 +49,7 @@ def point_on_line_given_distance(start_node, end_node, distance):
 
 
 def get_key_index_in_lane(
-    nodes: List[Dict[str, float]],
+    nodes: list[dict[str, float]],
     distance: float,
     direction: Union[Literal["front"], Literal["back"]],
 ) -> int:
@@ -61,10 +61,10 @@ def get_key_index_in_lane(
         _index_offset, _index_factor = len(_nodes) - 1, -1
     else:
         raise ValueError(f"Invalid direction type {direction}!")
-    _lane_points: List[Tuple[float, float, float]] = [
+    _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_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])
@@ -82,20 +82,20 @@ def get_key_index_in_lane(
 
 
 def get_xy_in_lane(
-    nodes: List[Dict[str, float]],
+    nodes: list[dict[str, float]],
     distance: float,
     direction: Union[Literal["front"], Literal["back"]],
-) -> Tuple[float, float]:
+) -> 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]] = [
+    _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_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])
@@ -122,7 +122,7 @@ def get_xy_in_lane(
 
 
 def get_direction_by_s(
-    nodes: List[Dict[str, float]],
+    nodes: list[dict[str, float]],
     distance: float,
     direction: Union[Literal["front"], Literal["back"]],
 ) -> float:
@@ -132,11 +132,11 @@ def get_direction_by_s(
         _nodes = [n for n in nodes[::-1]]
     else:
         raise ValueError(f"Invalid direction type {direction}!")
-    _lane_points: List[Tuple[float, float, float]] = [
+    _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]] = []
+    _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])

pycityagent/llm/embedding.py

@@ -1,7 +1,6 @@
 """简单的基于内存的embedding实现"""
 
 import numpy as np
-from typing import List, Dict, Optional
 import hashlib
 import json
 
@@ -22,34 +21,34 @@ class SimpleEmbedding:
         """
         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._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]:
+    def _tokenize(self, text: str) -> list[str]:
         """简单的分词"""
         # 这里使用简单的空格分词，实际应用中可以使用更复杂的分词方法
         return text.lower().split()
 
-    def _update_vocab(self, tokens: List[str]):
+    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]):
+    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]:
+    def _calculate_tf(self, tokens: list[str]) -> dict[str, float]:
         """计算词频(TF)"""
         tf = {}
         total_tokens = len(tokens)
@@ -60,7 +59,7 @@ class SimpleEmbedding:
             tf[token] /= total_tokens
         return tf
 
-    def _calculate_tfidf(self, tokens: List[str]) -> np.ndarray:
+    def _calculate_tfidf(self, tokens: list[str]) -> np.ndarray:
         """计算TF-IDF向量"""
         vector = np.zeros(self.vector_dim)
         tf = self._calculate_tf(tokens)

pycityagent/llm/llm.py

@@ -14,7 +14,7 @@ import requests
 from dashscope import ImageSynthesis
 from PIL import Image
 from io import BytesIO
-from typing import Any, Optional, Union, List, Dict
+from typing import Any, Optional, Union
 from .llmconfig import *
 from .utils import *
 
@@ -117,8 +117,8 @@ Token Usage:
         presence_penalty: Optional[float] = None,
         timeout: int = 300,
         retries=3,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        tool_choice: Optional[Dict[str, Any]] = None,
+        tools: Optional[list[dict[str, Any]]] = None,
+        tool_choice: Optional[dict[str, Any]] = None,
     ):
         """
         异步版文本请求
@@ -227,9 +227,9 @@ Token Usage:
                     self.prompt_tokens_used += result_response.usage.prompt_tokens  # type: ignore
                     self.completion_tokens_used += result_response.usage.completion_tokens  # type: ignore
                     self.request_number += 1
-                    if tools and result_response.choices[0].message.tool_calls:
+                    if tools and result_response.choices[0].message.tool_calls:  # type: ignore
                         return json.loads(
-                            result_response.choices[0]
+                            result_response.choices[0]  # type: ignore
                             .message.tool_calls[0]
                             .function.arguments
                         )

pycityagent/memory/memory.py

@@ -2,7 +2,8 @@ import asyncio
 import logging
 from copy import deepcopy
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
+from typing import Any, Literal, Optional,  Union
+from collections.abc import Sequence,Callable
 
 import numpy as np
 from pyparsing import deque
@@ -27,10 +28,10 @@ class Memory:
 
     def __init__(
         self,
-        config: Optional[Dict[Any, Any]] = None,
-        profile: Optional[Dict[Any, Any]] = None,
-        base: Optional[Dict[Any, Any]] = None,
-        motion: Optional[Dict[Any, Any]] = None,
+        config: Optional[dict[Any, Any]] = None,
+        profile: Optional[dict[Any, Any]] = None,
+        base: Optional[dict[Any, Any]] = None,
+        motion: Optional[dict[Any, Any]] = None,
         activate_timestamp: bool = False,
         embedding_model: Any = None,
     ) -> None:
@@ -38,7 +39,7 @@ class Memory:
         Initializes the Memory with optional configuration.
 
         Args:
-            config (Optional[Dict[Any, Any]], optional):
+            config (Optional[dict[Any, Any]], optional):
                 A configuration dictionary for dynamic memory. The dictionary format is:
                 - Key: The name of the dynamic memory field.
                 - Value: Can be one of two formats:
@@ -46,24 +47,24 @@ class Memory:
                     2. A callable that returns the default value when invoked (useful for complex default values).
                 Note: If a key in `config` overlaps with predefined attributes in `PROFILE_ATTRIBUTES` or `STATE_ATTRIBUTES`, a warning will be logged, and the key will be ignored.
                 Defaults to None.
-            profile (Optional[Dict[Any, Any]], optional): profile attribute dict.
-            base (Optional[Dict[Any, Any]], optional): base attribute dict from City Simulator.
-            motion (Optional[Dict[Any, Any]], optional): motion attribute dict from City Simulator.
+            profile (Optional[dict[Any, Any]], optional): profile attribute dict.
+            base (Optional[dict[Any, Any]], optional): base attribute dict from City Simulator.
+            motion (Optional[dict[Any, Any]], optional): motion attribute dict from City Simulator.
             activate_timestamp (bool): Whether activate timestamp storage in MemoryUnit
             embedding_model (Any): The embedding model for memory search.
         """
-        self.watchers: Dict[str, List[Callable]] = {}
+        self.watchers: dict[str, list[Callable]] = {}
         self._lock = asyncio.Lock()
         self.embedding_model = embedding_model
 
         # 初始化embedding存储
         self._embeddings = {"state": {}, "profile": {}, "dynamic": {}}
 
-        _dynamic_config: Dict[Any, Any] = {}
-        _state_config: Dict[Any, Any] = {}
-        _profile_config: Dict[Any, Any] = {}
+        _dynamic_config: dict[Any, Any] = {}
+        _state_config: dict[Any, Any] = {}
+        _profile_config: dict[Any, Any] = {}
         # 记录哪些字段需要embedding
-        self._embedding_fields: Dict[str, bool] = {}
+        self._embedding_fields: dict[str, bool] = {}
 
         if config is not None:
             for k, v in config.items():
@@ -303,7 +304,7 @@ class Memory:
 
     async def update_batch(
         self,
-        content: Union[Dict, Sequence[Tuple[Any, Any]]],
+        content: Union[dict, Sequence[tuple[Any, Any]]],
         mode: Union[Literal["replace"], Literal["merge"]] = "replace",
         store_snapshot: bool = False,
         protect_llm_read_only_fields: bool = True,
@@ -312,7 +313,7 @@ class Memory:
         Updates multiple values in the memory at once.
 
         Args:
-            content (Union[Dict, Sequence[Tuple[Any, Any]]]): A dictionary or sequence of tuples containing the keys and values to update.
+            content (Union[dict, Sequence[tuple[Any, Any]]]): A dictionary or sequence of tuples containing the keys and values to update.
             mode (Union[Literal["replace"], Literal["merge"]], optional): Update mode. Defaults to "replace".
             store_snapshot (bool): Whether to store a snapshot of the memory after the update.
             protect_llm_read_only_fields (bool): Whether to protect non-self define fields from being updated.
@@ -321,9 +322,9 @@ class Memory:
             TypeError: If the content type is neither a dictionary nor a sequence of tuples.
         """
         if isinstance(content, dict):
-            _list_content: List[Tuple[Any, Any]] = [(k, v) for k, v in content.items()]
+            _list_content: list[tuple[Any, Any]] = [(k, v) for k, v in content.items()]
         elif isinstance(content, Sequence):
-            _list_content: List[Tuple[Any, Any]] = [(k, v) for k, v in content]
+            _list_content: list[tuple[Any, Any]] = [(k, v) for k, v in content]
         else:
             raise TypeError(f"Invalid content type `{type(content)}`!")
         for k, v in _list_content[:1]:
@@ -353,12 +354,12 @@ class Memory:
     @lock_decorator
     async def export(
         self,
-    ) -> Tuple[Sequence[Dict], Sequence[Dict], Sequence[Dict]]:
+    ) -> tuple[Sequence[dict], Sequence[dict], Sequence[dict]]:
         """
         Exports the current state of all memory sections.
 
         Returns:
-            Tuple[Sequence[Dict], Sequence[Dict], Sequence[Dict]]: A tuple containing the exported data of profile, state, and dynamic memory sections.
+            tuple[Sequence[dict], Sequence[dict], Sequence[dict]]: A tuple containing the exported data of profile, state, and dynamic memory sections.
         """
         return (
             await self._profile.export(),
@@ -369,14 +370,14 @@ class Memory:
     @lock_decorator
     async def load(
         self,
-        snapshots: Tuple[Sequence[Dict], Sequence[Dict], Sequence[Dict]],
+        snapshots: tuple[Sequence[dict], Sequence[dict], Sequence[dict]],
         reset_memory: bool = True,
     ) -> None:
         """
         Import the snapshot memories of all sections.
 
         Args:
-            snapshots (Tuple[Sequence[Dict], Sequence[Dict], Sequence[Dict]]): The exported snapshots.
+            snapshots (tuple[Sequence[dict], Sequence[dict], Sequence[dict]]): The exported snapshots.
             reset_memory (bool): Whether to reset previous memory.
         """
         _profile_snapshot, _state_snapshot, _dynamic_snapshot = snapshots

pycityagent/metrics/__init__.py

@@ -1,5 +1,6 @@
-from .mlflow_client import MlflowClient
+from .mlflow_client import MlflowClient,init_mlflow_connection
 
 __all__ = [
     "MlflowClient",
+    "init_mlflow_connection",
 ]

pycityagent/metrics/mlflow_client.py

@@ -18,6 +18,55 @@ from ..utils.decorators import lock_decorator
 logger = logging.getLogger("mlflow")
 
 
+def init_mlflow_connection(
+    config: dict,
+    mlflow_run_name: Optional[str] = None,
+    experiment_name: Optional[str] = None,
+    experiment_description: Optional[str] = None,
+    experiment_tags: Optional[dict[str, Any]] = None,
+) -> tuple[str, tuple[str, mlflow.MlflowClient, Run, str]]:
+
+    os.environ["MLFLOW_TRACKING_USERNAME"] = config.get("username", None)
+    os.environ["MLFLOW_TRACKING_PASSWORD"] = config.get("password", None)
+
+    run_uuid = str(uuid.uuid4())
+    # run name
+    if mlflow_run_name is None:
+        mlflow_run_name = f"exp_{run_uuid}"
+
+    # exp name
+    if experiment_name is None:
+        experiment_name = f"run_{run_uuid}"
+
+    # tags
+    if experiment_tags is None:
+        experiment_tags = {}
+    if experiment_description is not None:
+        experiment_tags["mlflow.note.content"] = experiment_description
+
+    uri = config["mlflow_uri"]
+    client = mlflow.MlflowClient(tracking_uri=uri)
+
+    # experiment
+    try:
+        experiment_id = client.create_experiment(
+            name=experiment_name,
+            tags=experiment_tags,
+        )
+    except Exception as e:
+        experiment = client.get_experiment_by_name(experiment_name)
+        if experiment is None:
+            raise e
+        experiment_id = experiment.experiment_id
+
+    # run
+    run = client.create_run(experiment_id=experiment_id, run_name=mlflow_run_name)
+
+    run_id = run.info.run_id
+
+    return run_id, (uri, client, run, run_uuid)
+
+
 class MlflowClient:
     """
     - Mlflow client
@@ -30,42 +79,30 @@ class MlflowClient:
         experiment_name: Optional[str] = None,
         experiment_description: Optional[str] = None,
         experiment_tags: Optional[dict[str, Any]] = None,
+        run_id: Optional[str] = None,
     ) -> None:
-        os.environ["MLFLOW_TRACKING_USERNAME"] = config.get("username", None)
-        os.environ["MLFLOW_TRACKING_PASSWORD"] = config.get("password", None)
-        self._mlflow_uri = uri = config["mlflow_uri"]
-        self._client = client = mlflow.MlflowClient(tracking_uri=uri)
-        self._run_uuid = run_uuid = str(uuid.uuid4())
-        self._lock = asyncio.Lock()
-        # run name
-        if mlflow_run_name is None:
-            mlflow_run_name = f"exp_{run_uuid}"
-
-        # exp name
-        if experiment_name is None:
-            experiment_name = f"run_{run_uuid}"
-
-        # tags
-        if experiment_tags is None:
-            experiment_tags = {}
-        if experiment_description is not None:
-            experiment_tags["mlflow.note.content"] = experiment_description
-
-        try:
-            self._experiment_id = experiment_id = client.create_experiment(
-                name=experiment_name,
-                tags=experiment_tags,
+        if run_id is None:
+            self._run_id, (
+                self._mlflow_uri,
+                self._client,
+                self._run,
+                self._run_uuid,
+            ) = init_mlflow_connection(
+                config=config,
+                mlflow_run_name=mlflow_run_name,
+                experiment_name=experiment_name,
+                experiment_description=experiment_description,
+                experiment_tags=experiment_tags,
             )
-        except Exception as e:
-            experiment = client.get_experiment_by_name(experiment_name)
-            if experiment is None:
-                raise e
-            self._experiment_id = experiment_id = experiment.experiment_id
-
-        self._run = run = client.create_run(
-            experiment_id=experiment_id, run_name=mlflow_run_name
-        )
-        self._run_id = run.info.run_id
+        else:
+            self._mlflow_uri = uri = config["mlflow_uri"]
+            os.environ["MLFLOW_TRACKING_USERNAME"] = config.get("username", None)
+            os.environ["MLFLOW_TRACKING_PASSWORD"] = config.get("password", None)
+            self._client = client = mlflow.MlflowClient(tracking_uri=uri)
+            self._run = client.get_run(run_id=run_id)
+            self._run_id = run_id
+            self._run_uuid = run_uuid = str(uuid.uuid4())
+        self._lock = asyncio.Lock()
 
     @property
     def client(
@@ -77,6 +114,7 @@ class MlflowClient:
     def run_id(
         self,
     ) -> str:
+        assert self._run_id is not None
         return self._run_id
 
     @lock_decorator

pycityagent/simulation/agentgroup.py

@@ -35,7 +35,8 @@ class AgentGroup:
         enable_avro: bool,
         avro_path: Path,
         enable_pgsql: bool,
-        pgsql_args: tuple[str, str, str, str, str],
+        pgsql_copy_writer: ray.ObjectRef,
+        mlflow_run_id: str,
         logging_level: int,
     ):
         logger.setLevel(logging_level)
@@ -88,6 +89,7 @@ class AgentGroup:
                 config=_mlflow_config,
                 mlflow_run_name=f"EXP_{exp_name}_{1000*int(time.time())}",
                 experiment_name=exp_name,
+                run_id=mlflow_run_id,
             )
         else:
             self.mlflow_client = None

pycityagent/simulation/simulation.py

@@ -3,6 +3,7 @@ import json
 import logging
 import os
 import random
+import time
 import uuid
 from collections.abc import Callable, Sequence
 from concurrent.futures import ThreadPoolExecutor
@@ -11,6 +12,7 @@ from pathlib import Path
 from typing import Any, Optional, Union
 
 import pycityproto.city.economy.v2.economy_pb2 as economyv2
+import ray
 import yaml
 from mosstool.map._map_util.const import AOI_START_ID
 
@@ -18,6 +20,7 @@ from ..agent import Agent, InstitutionAgent
 from ..environment.simulator import Simulator
 from ..memory.memory import Memory
 from ..message.messager import Messager
+from ..metrics import init_mlflow_connection
 from ..survey import Survey
 from .agentgroup import AgentGroup
 
@@ -165,7 +168,8 @@ class AgentSimulation:
         enable_avro: bool,
         avro_path: Path,
         enable_pgsql: bool,
-        pgsql_args: tuple[str, str, str, str, str],
+        pgsql_copy_writer: ray.ObjectRef,
+        mlflow_run_id: str = None,  # type: ignore
         logging_level: int = logging.WARNING,
     ):
         """创建远程组"""
@@ -177,7 +181,8 @@ class AgentSimulation:
             enable_avro,
             avro_path,
             enable_pgsql,
-            pgsql_args,
+            pgsql_copy_writer,
+            mlflow_run_id,
             logging_level,
         )
         return group_name, group, agents
@@ -267,6 +272,16 @@ class AgentSimulation:
 
             class_init_index += agent_count_i
 
+        # 初始化mlflow连接
+        _mlflow_config = self.config.get("metric_request", {}).get("mlflow")
+        if _mlflow_config:
+            mlflow_run_id, _ = init_mlflow_connection(
+                config=_mlflow_config,
+                mlflow_run_name=f"EXP_{self.exp_name}_{1000*int(time.time())}",
+                experiment_name=self.exp_name,
+            )
+        else:
+            mlflow_run_id = None
         # 收集所有创建组的参数
         creation_tasks = []
         for group_name, agents in group_creation_params:
@@ -279,7 +294,10 @@ class AgentSimulation:
                 self.enable_avro,
                 self.avro_path,
                 self.enable_pgsql,
-                self._pgsql_args,
+                # TODO:
+                # self._pgsql_copy_writer, # type:ignore
+                None,
+                mlflow_run_id,
                 self.logging_level,
             )
             creation_tasks.append((group_name, group, agents))

pycityagent/utils/parsers/parser_base.py

@@ -4,7 +4,7 @@ Base class of parser
 
 import re
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Union
 
 
 class ParserBase(ABC):

pycityagent/workflow/prompt.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Optional, Union
 import re
 
 
@@ -10,7 +10,7 @@ class FormatPrompt:
     Attributes:
         template (str): The template string containing placeholders.
         system_prompt (Optional[str]): An optional system prompt to add to the dialog.
-        variables (List[str]): A list of variable names extracted from the template.
+        variables (list[str]): A list of variable names extracted from the template.
         formatted_string (str): The formatted string derived from the template and provided variables.
     """
 
@@ -27,12 +27,12 @@ class FormatPrompt:
         self.variables = self._extract_variables()
         self.formatted_string = ""  # To store the formatted string
 
-    def _extract_variables(self) -> List[str]:
+    def _extract_variables(self) -> list[str]:
         """
         Extracts variable names from the template string.
 
         Returns:
-            List[str]: A list of variable names found within the template.
+            list[str]: A list of variable names found within the template.
         """
         return re.findall(r"\{(\w+)\}", self.template)
 
@@ -51,12 +51,12 @@ class FormatPrompt:
         )  # Store the formatted string
         return self.formatted_string
 
-    def to_dialog(self) -> List[Dict[str, str]]:
+    def to_dialog(self) -> list[dict[str, str]]:
         """
         Converts the formatted prompt and optional system prompt into a dialog format.
 
         Returns:
-            List[Dict[str, str]]: A list representing the dialog with roles and content.
+            list[dict[str, str]]: A list representing the dialog with roles and content.
         """
         dialog = []
         if self.system_prompt:

pycityagent/workflow/tool.py

@@ -1,6 +1,6 @@
 import time
-from typing import Any, Callable, Dict, List, Optional, Union
-
+from typing import Any, Optional, Union
+from collections.abc import Callable
 from mlflow.entities import Metric
 
 from ..agent import Agent
@@ -76,7 +76,7 @@ class SencePOI(Tool):
     Attributes:
         radius (int): The radius within which to search for POIs.
         category_prefix (str): The prefix for the categories of POIs to consider.
-        variables (List[str]): A list of variables relevant to the tool's operation.
+        variables (list[str]): A list of variables relevant to the tool's operation.
 
     Args:
         radius (int, optional): The circular search radius. Defaults to 100.

pycityagent/workflow/trigger.py

@@ -1,5 +1,5 @@
 import asyncio
-from typing import Any, Callable, Dict, List, Optional, Union, Type
+from typing import Optional
 import socket
 from ..memory import Memory
 from ..environment import Simulator
@@ -11,7 +11,7 @@ class EventTrigger:
     """Base class for event triggers that wait for specific conditions to be met."""
 
     # 定义该trigger需要的组件类型
-    required_components: List[Type] = []
+    required_components: list[type] = []
 
     def __init__(self, block=None):
         self.block = block

{pycityagent-2.0.0a19.dist-info → pycityagent-2.0.0a20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a19
+Version: 2.0.0a20
 Summary: LLM-based城市环境agent构建库
 License: MIT
 Author: Yuwei Yan
@@ -34,8 +34,9 @@ Requires-Dist: openai (>=1.58.1,<2.0.0)
 Requires-Dist: pandavro (>=1.8.0,<2.0.0)
 Requires-Dist: poetry (>=1.2.2)
 Requires-Dist: protobuf (<=4.24.0)
+Requires-Dist: psycopg[binary] (>=3.2.3,<4.0.0)
 Requires-Dist: pycitydata (==1.0.0)
-Requires-Dist: pycityproto (>=2.1.4,<3.0.0)
+Requires-Dist: pycityproto (>=2.1.5,<3.0.0)
 Requires-Dist: pyyaml (>=6.0.2,<7.0.0)
 Requires-Dist: ray (>=2.40.0,<3.0.0)
 Requires-Dist: sidecar (==0.7.0)

{pycityagent-2.0.0a19.dist-info → pycityagent-2.0.0a20.dist-info}/RECORD

@@ -1,7 +1,7 @@
 pycityagent/__init__.py,sha256=EDxt3Su3lH1IMh9suNw7GeGL7UrXeWiZTw5KWNznDzc,637
-pycityagent/agent.py,sha256=gy9OlFpnEla7wK45k4m7_ySLp67mBKoG50vxUEmxbJE,24280
+pycityagent/agent.py,sha256=HRFAG_iM1q3nvXtV0T-Dz01foOtty4IWka7h4WD97CU,24268
 pycityagent/economy/__init__.py,sha256=aonY4WHnx-6EGJ4WKrx4S-2jAkYNLtqUA04jp6q8B7w,75
-pycityagent/economy/econ_client.py,sha256=EZDGxM7K83ucYZQ5qdv6HA-jhRCWbR1u5q-kLMqelKc,11192
+pycityagent/economy/econ_client.py,sha256=GuHK9ZBnhqW3Z7F8ViDJn_iN73yOBbbwFyJv1wLEBDk,12211
 pycityagent/environment/__init__.py,sha256=awHxlOud-btWbk0FCS4RmGJ13W84oVCkbGfcrhKqihA,240
 pycityagent/environment/interact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycityagent/environment/interact/interact.py,sha256=ifxPPzuHeqLHIZ_6zvfXMoBOnBsXNIP4bYp7OJ7pnEQ,6588
@@ -25,19 +25,19 @@ pycityagent/environment/simulator.py,sha256=XjcxbyBIbB3Ht9z087z_oWIPAN6pP5Eq1lyf
 pycityagent/environment/utils/__init__.py,sha256=1m4Q1EfGvNpUsa1bgQzzCyWhfkpElnskNImjjFD3Znc,237
 pycityagent/environment/utils/base64.py,sha256=hoREzQo3FXMN79pqQLO2jgsDEvudciomyKii7MWljAM,374
 pycityagent/environment/utils/const.py,sha256=3RMNy7_bE7-23K90j9DFW_tWEzu8s7hSTgKbV-3BFl4,5327
-pycityagent/environment/utils/geojson.py,sha256=Ieg8Bzw63kKhJlhDIOVDoh-wQO4Sbtoe47FtIOy5wWg,686
+pycityagent/environment/utils/geojson.py,sha256=LVHAdEhnZM8d0BoUnuPiIL_gaeXBIIglrLrfje5M0b4,661
 pycityagent/environment/utils/grpc.py,sha256=6EJwKXXktIWb1NcUiJzIRmfrY0S03QAXXGcCDHqAT00,1998
-pycityagent/environment/utils/map_utils.py,sha256=oqrRgQICC3SYw6gwjjPe_MAif7_t6dlrQpY8E32Fexs,5777
+pycityagent/environment/utils/map_utils.py,sha256=lYOEoCFFK6-e9N5txLMMq4HUlxMqc8Uw1YrGW5oJmgg,5749
 pycityagent/environment/utils/port.py,sha256=3OM6kSUt3PxvDUOlgyiendBtETaWU8Mzk_8H0TzTmYg,295
 pycityagent/environment/utils/protobuf.py,sha256=0jBvK_s96y_n7tuMbG22TOtQmg71SGV4ONDy2IGsU9o,1148
 pycityagent/llm/__init__.py,sha256=7klKEmCcDWJIu-F4DoAukSuKfDbLhdczrSIhpwow-sY,145
-pycityagent/llm/embedding.py,sha256=Y0xhm_Ny6cawqzlendXb-mAS2QAuuEez1UtTR5-Kb2Q,4293
-pycityagent/llm/llm.py,sha256=BtxBvPK4tb8QlZIfxO5XJ73lKXwF8L31LqVbejWB8eo,15121
+pycityagent/llm/embedding.py,sha256=2psX_EK67oPlYe77g43EYUYams4M9AiJqxpHTFHG0n8,4253
+pycityagent/llm/llm.py,sha256=vJaaGqVuyV-GlBxrnvGKZnMDlxeTT_sGUTdxz5tYwEE,15141
 pycityagent/llm/llmconfig.py,sha256=4Ylf4OFSBEFy8jrOneeX0HvPhWEaF5jGvy1HkXK08Ro,436
 pycityagent/llm/utils.py,sha256=hoNPhvomb1u6lhFX0GctFipw74hVKb7bvUBDqwBzBYw,160
 pycityagent/memory/__init__.py,sha256=Hs2NhYpIG-lvpwPWwj4DydB1sxtjz7cuA4iDAzCXnjI,243
 pycityagent/memory/const.py,sha256=6zpJPJXWoH9-yf4RARYYff586agCoud9BRn7sPERB1g,932
-pycityagent/memory/memory.py,sha256=FjKVL_MgNBnSc0sox2tuxLqXg9_MQQr9vYdRDHMdDL4,18183
+pycityagent/memory/memory.py,sha256=vJxHOI74aJDGZPFu2LbBr02ASfOYpig66fto6Gjr-6Q,18191
 pycityagent/memory/memory_base.py,sha256=euKZRCs4dbcKxjlZzpLCTnH066DAtRjj5g1JFKD40qQ,5633
 pycityagent/memory/profile.py,sha256=s4LnxSPGSjIGZXHXkkd8mMa6uYYZrytgyQdWjcaqGf4,5182
 pycityagent/memory/self_define.py,sha256=poPiexNhOLq_iTgK8s4mK_xoL_DAAcB8kMvInj7iE5E,5179
@@ -45,12 +45,12 @@ pycityagent/memory/state.py,sha256=5W0c1yJ-aaPpE74B2LEcw3Ygpm77tyooHv8NylyrozE,5
 pycityagent/memory/utils.py,sha256=wLNlNlZ-AY9VB8kbUIy0UQSYh26FOQABbhmKQkit5o8,850
 pycityagent/message/__init__.py,sha256=TCjazxqb5DVwbTu1fF0sNvaH_EPXVuj2XQ0p6W-QCLU,55
 pycityagent/message/messager.py,sha256=W_OVlNGcreHSBf6v-DrEnfNCXExB78ySr0w26MSncfU,2541
-pycityagent/metrics/__init__.py,sha256=IrlwXs_733b3PlMV99Ogh6R8sGCEalF2qIT2mQwdPjU,75
-pycityagent/metrics/mlflow_client.py,sha256=BH-YVMlYr-eTFdrCZjgMZWi8ptUEm7todOXCpYixAgU,3311
+pycityagent/metrics/__init__.py,sha256=X08PaBbGVAd7_PRGLREXWxaqm7nS82WBQpD1zvQzcqc,128
+pycityagent/metrics/mlflow_client.py,sha256=g_tHxWkWTDijtbGL74-HmiYzWVKb1y8-w12QrY9jL30,4449
 pycityagent/metrics/utils/const.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycityagent/simulation/__init__.py,sha256=jYaqaNpzM5M_e_ykISS_M-mIyYdzJXJWhgpfBpA6l5k,111
-pycityagent/simulation/agentgroup.py,sha256=H2E_YQ3ir3gQmPWsupHXA_LSBVkffzkXWl4UgZ9qLOc,13327
-pycityagent/simulation/simulation.py,sha256=ANf04GqaK7gLT50F_ZcP2UcpS_uGx12pLqkrK6eA2L8,21014
+pycityagent/simulation/agentgroup.py,sha256=6p8-OP2x_syaQ1pWLgll0LHs823OHnappuY-8XzL_LU,13383
+pycityagent/simulation/simulation.py,sha256=pxFdaNMLdJMcpW8gnOcxjfbz4gkpcvWsTR0rNP8rE9k,21675
 pycityagent/survey/__init__.py,sha256=rxwou8U9KeFSP7rMzXtmtp2fVFZxK4Trzi-psx9LPIs,153
 pycityagent/survey/manager.py,sha256=S5IkwTdelsdtZETChRcfCEczzwSrry_Fly9MY4s3rbk,1681
 pycityagent/survey/models.py,sha256=YE50UUt5qJ0O_lIUsSY6XFCGUTkJVNu_L1gAhaCJ2fs,3546
@@ -60,13 +60,13 @@ pycityagent/utils/decorators.py,sha256=Gk3r41hfk6awui40tbwpq3C7wC7jHaRmLRlcJFlLQ
 pycityagent/utils/parsers/__init__.py,sha256=AN2xgiPxszWK4rpX7zrqRsqNwfGF3WnCA5-PFTvbaKk,281
 pycityagent/utils/parsers/code_block_parser.py,sha256=Cs2Z_hm9VfNCpPPll1TwteaJF-HAQPs-3RApsOekFm4,1173
 pycityagent/utils/parsers/json_parser.py,sha256=FZ3XN1g8z4Dr2TFraUOoah1oQcze4fPd2m01hHoX0Mo,2917
-pycityagent/utils/parsers/parser_base.py,sha256=k6DVqwAMK3jJdOP4IeLE-aFPm3V2F-St5qRBuRdx4aU,1742
+pycityagent/utils/parsers/parser_base.py,sha256=KBKO4zLZPNdGjPAGqIus8LseZ8W3Tlt2y0QxqeCd25Q,1713
 pycityagent/utils/survey_util.py,sha256=Be9nptmu2JtesFNemPgORh_2GsN7rcDYGQS9Zfvc5OI,2169
 pycityagent/workflow/__init__.py,sha256=QNkUV-9mACMrR8c0cSKna2gC1mMZdxXbxWzjE-Uods0,621
 pycityagent/workflow/block.py,sha256=WkE2On97DCZS_9n8aIgT8wxv9Oaff4Fdf2tLqbKfMtE,6010
-pycityagent/workflow/prompt.py,sha256=tY69nDO8fgYfF_dOA-iceR8pAhkYmCqoox8uRPqEuGY,2956
-pycityagent/workflow/tool.py,sha256=uaB0dV35jA9v2UqWw9L8iPM-HJW5e9BlgFVgOMf9jvw,8201
-pycityagent/workflow/trigger.py,sha256=t5X_i0WtL32bipZSsq_E3UUyYYudYLxQUpvxbgClp2s,5683
-pycityagent-2.0.0a19.dist-info/METADATA,sha256=x0DbwNQ4Zj0rswWGIv6MA0zSBBEiExXLIdErYA0bhnE,7800
-pycityagent-2.0.0a19.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-pycityagent-2.0.0a19.dist-info/RECORD,,
+pycityagent/workflow/prompt.py,sha256=6jI0Rq54JLv3-IXqZLYug62vse10wTI83xvf4ZX42nk,2929
+pycityagent/workflow/tool.py,sha256=SGY18lT71hBLKagopirFbxRjPY_387Dobo9SUwjHIn0,8215
+pycityagent/workflow/trigger.py,sha256=Df-MOBEDWBbM-v0dFLQLXteLsipymT4n8vqexmK2GiQ,5643
+pycityagent-2.0.0a20.dist-info/METADATA,sha256=aygfImVFM4jG0nXVuunu5LZDNi4JyikBULFce0WGWeM,7848
+pycityagent-2.0.0a20.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+pycityagent-2.0.0a20.dist-info/RECORD,,