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

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/__init__.py

@@ -3,5 +3,6 @@
 """
 
 from .simulation import AgentSimulation
+from .storage.pg import PgWriter, create_pg_tables
 
-__all__ = ["AgentSimulation"]
+__all__ = ["AgentSimulation", "PgWriter", "create_pg_tables"]

pycityagent/simulation/agentgroup.py

@@ -3,7 +3,7 @@ import json
 import logging
 import time
 import uuid
-from datetime import datetime
+from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any
 from uuid import UUID
@@ -35,7 +35,8 @@ class AgentGroup:
         enable_avro: bool,
         avro_path: Path,
         enable_pgsql: bool,
-        pgsql_args: tuple[str, str, str, str, str],
+        pgsql_writer: ray.ObjectRef,
+        mlflow_run_id: str,
         logging_level: int,
     ):
         logger.setLevel(logging_level)
@@ -44,6 +45,7 @@ class AgentGroup:
         self.config = config
         self.exp_id = exp_id
         self.enable_avro = enable_avro
+        self.enable_pgsql = enable_pgsql
         if enable_avro:
             self.avro_path = avro_path / f"{self._uuid}"
             self.avro_path.mkdir(parents=True, exist_ok=True)
@@ -53,6 +55,8 @@ class AgentGroup:
                 "status": self.avro_path / f"status.avro",
                 "survey": self.avro_path / f"survey.avro",
             }
+        if self.enable_pgsql:
+            pass
 
         self.messager = Messager(
             hostname=config["simulator_request"]["mqtt"]["server"],
@@ -60,6 +64,8 @@ class AgentGroup:
             username=config["simulator_request"]["mqtt"].get("username", None),
             password=config["simulator_request"]["mqtt"].get("password", None),
         )
+        self._pgsql_writer = pgsql_writer
+        self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
         self.initialized = False
         self.id2agent = {}
         # Step:1 prepare LLM client
@@ -88,6 +94,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
@@ -103,6 +110,8 @@ class AgentGroup:
             agent.set_messager(self.messager)
             if self.enable_avro:
                 agent.set_avro_file(self.avro_file)  # type: ignore
+            if self.enable_pgsql:
+                agent.set_pgsql_writer(self._pgsql_writer)
 
     async def init_agents(self):
         logger.debug(f"-----Initializing Agents in AgentGroup {self._uuid} ...")
@@ -159,6 +168,20 @@ class AgentGroup:
             with open(filename, "wb") as f:
                 surveys = []
                 fastavro.writer(f, SURVEY_SCHEMA, surveys)
+
+        if self.enable_pgsql:
+            if not issubclass(type(self.agents[0]), InstitutionAgent):
+                profiles: list[Any] = []
+                for agent in self.agents:
+                    profile = await agent.memory._profile.export()
+                    profile = profile[0]
+                    profile["id"] = agent._uuid
+                    profiles.append(
+                        (agent._uuid, profile.get("name", ""), json.dumps(profile))
+                    )
+                await self._pgsql_writer.async_write_profile.remote(  # type:ignore
+                    profiles
+                )
         self.initialized = True
         logger.debug(f"-----AgentGroup {self._uuid} initialized")
 
@@ -216,11 +239,13 @@ class AgentGroup:
             await asyncio.sleep(0.5)
 
     async def save_status(self):
+        _statuses_time_list: list[tuple[dict, datetime]] = []
         if self.enable_avro:
             logger.debug(f"-----Saving status for group {self._uuid}")
             avros = []
             if not issubclass(type(self.agents[0]), InstitutionAgent):
                 for agent in self.agents:
+                    _date_time = datetime.now(timezone.utc)
                     position = await agent.memory.get("position")
                     lng = position["longlat_position"]["longitude"]
                     lat = position["longlat_position"]["latitude"]
@@ -246,13 +271,15 @@ class AgentGroup:
                         "tired": needs["tired"],
                         "safe": needs["safe"],
                         "social": needs["social"],
-                        "created_at": int(datetime.now().timestamp() * 1000),
+                        "created_at": int(_date_time.timestamp() * 1000),
                     }
                     avros.append(avro)
+                    _statuses_time_list.append((avro, _date_time))
                 with open(self.avro_file["status"], "a+b") as f:
                     fastavro.writer(f, STATUS_SCHEMA, avros, codec="snappy")
             else:
                 for agent in self.agents:
+                    _date_time = datetime.now(timezone.utc)
                     avro = {
                         "id": agent._uuid,
                         "day": await self.simulator.get_simulator_day(),
@@ -272,8 +299,109 @@ class AgentGroup:
                         "customers": await agent.memory.get("customers"),
                     }
                     avros.append(avro)
+                    _statuses_time_list.append((avro, _date_time))
                 with open(self.avro_file["status"], "a+b") as f:
                     fastavro.writer(f, INSTITUTION_STATUS_SCHEMA, avros, codec="snappy")
+        if self.enable_pgsql:
+            # data already acquired from Avro part
+            if len(_statuses_time_list) > 0:
+                for _status_dict, _date_time in _statuses_time_list:
+                    for key in ["lng", "lat", "parent_id"]:
+                        if key not in _status_dict:
+                            _status_dict[key] = -1
+                    for key in [
+                        "action",
+                    ]:
+                        if key not in _status_dict:
+                            _status_dict[key] = ""
+                    _status_dict["created_at"] = _date_time
+            else:
+                if not issubclass(type(self.agents[0]), InstitutionAgent):
+                    for agent in self.agents:
+                        _date_time = datetime.now(timezone.utc)
+                        position = await agent.memory.get("position")
+                        lng = position["longlat_position"]["longitude"]
+                        lat = position["longlat_position"]["latitude"]
+                        if "aoi_position" in position:
+                            parent_id = position["aoi_position"]["aoi_id"]
+                        elif "lane_position" in position:
+                            parent_id = position["lane_position"]["lane_id"]
+                        else:
+                            # BUG: 需要处理
+                            parent_id = -1
+                        needs = await agent.memory.get("needs")
+                        action = await agent.memory.get("current_step")
+                        action = action["intention"]
+                        _status_dict = {
+                            "id": agent._uuid,
+                            "day": await self.simulator.get_simulator_day(),
+                            "t": await self.simulator.get_simulator_second_from_start_of_day(),
+                            "lng": lng,
+                            "lat": lat,
+                            "parent_id": parent_id,
+                            "action": action,
+                            "hungry": needs["hungry"],
+                            "tired": needs["tired"],
+                            "safe": needs["safe"],
+                            "social": needs["social"],
+                            "created_at": _date_time,
+                        }
+                        _statuses_time_list.append((_status_dict, _date_time))
+                else:
+                    for agent in self.agents:
+                        _date_time = datetime.now(timezone.utc)
+                        _status_dict = {
+                            "id": agent._uuid,
+                            "day": await self.simulator.get_simulator_day(),
+                            "t": await self.simulator.get_simulator_second_from_start_of_day(),
+                            "lng": -1,
+                            "lat": -1,
+                            "parent_id": -1,
+                            "action": "",
+                            "type": await agent.memory.get("type"),
+                            "nominal_gdp": await agent.memory.get("nominal_gdp"),
+                            "real_gdp": await agent.memory.get("real_gdp"),
+                            "unemployment": await agent.memory.get("unemployment"),
+                            "wages": await agent.memory.get("wages"),
+                            "prices": await agent.memory.get("prices"),
+                            "inventory": await agent.memory.get("inventory"),
+                            "price": await agent.memory.get("price"),
+                            "interest_rate": await agent.memory.get("interest_rate"),
+                            "bracket_cutoffs": await agent.memory.get(
+                                "bracket_cutoffs"
+                            ),
+                            "bracket_rates": await agent.memory.get("bracket_rates"),
+                            "employees": await agent.memory.get("employees"),
+                            "customers": await agent.memory.get("customers"),
+                            "created_at": _date_time,
+                        }
+                        _statuses_time_list.append((_status_dict, _date_time))
+        to_update_statues: list[tuple] = []
+        for _status_dict, _ in _statuses_time_list:
+            BASIC_KEYS = [
+                "id",
+                "day",
+                "t",
+                "lng",
+                "lat",
+                "parent_id",
+                "action",
+                "created_at",
+            ]
+            _data = [_status_dict[k] for k in BASIC_KEYS if k != "created_at"]
+            _other_dict = json.dumps(
+                {k: v for k, v in _status_dict.items() if k not in BASIC_KEYS}
+            )
+            _data.append(_other_dict)
+            _data.append(_status_dict["created_at"])
+            to_update_statues.append(tuple(_data))
+        if self._last_asyncio_pg_task is not None:
+            await self._last_asyncio_pg_task
+        self._last_asyncio_pg_task = (
+            self._pgsql_writer.async_write_status.remote(  # type:ignore
+                to_update_statues
+            )
+        )
 
     async def step(self):
         if not self.initialized: