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

pycityagent/agent.py

@@ -1,17 +1,19 @@
 """智能体模板类及其定义"""
 
 import asyncio
+import json
 import logging
 import random
 import uuid
 from abc import ABC, abstractmethod
 from copy import deepcopy
-from datetime import datetime
+from datetime import datetime, timezone
 from enum import Enum
 from typing import Any, Optional
 from uuid import UUID
 
 import fastavro
+import ray
 from mosstool.util.format_converter import dict2pb
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
 
@@ -56,6 +58,7 @@ class Agent(ABC):
         mlflow_client: Optional[MlflowClient] = None,
         memory: Optional[Memory] = None,
         avro_file: Optional[dict[str, str]] = None,
+        copy_writer: Optional[ray.ObjectRef] = None,
     ) -> None:
         """
         Initialize the Agent.
@@ -70,6 +73,7 @@ class Agent(ABC):
             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.
+            copy_writer (ray.ObjectRef): The copy_writer of the agent. Defaults to None.
         """
         self._name = name
         self._type = type
@@ -88,6 +92,8 @@ class Agent(ABC):
         self._interview_history: list[dict] = []  # 存储采访历史
         self._person_template = PersonService.default_dict_person()
         self._avro_file = avro_file
+        self._pgsql_writer = copy_writer
+        self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
 
     def __getstate__(self):
         state = self.__dict__.copy()
@@ -143,6 +149,12 @@ class Agent(ABC):
         """
         self._avro_file = avro_file
 
+    def set_pgsql_writer(self, pgsql_writer: ray.ObjectRef):
+        """
+        Set the PostgreSQL copy writer of the agent.
+        """
+        self._pgsql_writer = pgsql_writer
+
     @property
     def uuid(self):
         """The Agent's UUID"""
@@ -198,6 +210,15 @@ class Agent(ABC):
             )
         return self._simulator
 
+    @property
+    def copy_writer(self):
+        """Pg Copy Writer"""
+        if self._pgsql_writer is None:
+            raise RuntimeError(
+                f"Copy Writer access before assignment, please `set_pgsql_writer` first!"
+            )
+        return self._pgsql_writer
+
     async def generate_user_survey_response(self, survey: dict) -> str:
         """生成回答 —— 可重写
         基于智能体的记忆和当前状态，生成对问卷调查的回答。
@@ -237,8 +258,8 @@ class Agent(ABC):
 
     async def _process_survey(self, survey: dict):
         survey_response = await self.generate_user_survey_response(survey)
-        if self._avro_file is None:
-            return
+        _date_time = datetime.now(timezone.utc)
+        # Avro
         response_to_avro = [
             {
                 "id": self._uuid,
@@ -246,11 +267,41 @@ class Agent(ABC):
                 "t": await self.simulator.get_simulator_second_from_start_of_day(),
                 "survey_id": survey["id"],
                 "result": survey_response,
-                "created_at": int(datetime.now().timestamp() * 1000),
+                "created_at": int(_date_time.timestamp() * 1000),
             }
         ]
-        with open(self._avro_file["survey"], "a+b") as f:
-            fastavro.writer(f, SURVEY_SCHEMA, response_to_avro, codec="snappy")
+        if self._avro_file is not None:
+            with open(self._avro_file["survey"], "a+b") as f:
+                fastavro.writer(f, SURVEY_SCHEMA, response_to_avro, codec="snappy")
+        # Pg
+        if self._pgsql_writer is not None:
+            if self._last_asyncio_pg_task is not None:
+                await self._last_asyncio_pg_task
+            _keys = [
+                "id",
+                "day",
+                "t",
+                "survey_id",
+                "result",
+            ]
+            _data_tuples: list[tuple] = []
+            # str to json
+            for _dict in response_to_avro:
+                res = _dict["result"]
+                _dict["result"] = json.dumps(
+                    {
+                        "result": res,
+                    }
+                )
+                _data_list = [_dict[k] for k in _keys]
+                # created_at
+                _data_list.append(_date_time)
+                _data_tuples.append(tuple(_data_list))
+            self._last_asyncio_pg_task = (
+                self._pgsql_writer.async_write_survey.remote(  # type:ignore
+                    _data_tuples
+                )
+            )
 
     async def generate_user_chat_response(self, question: str) -> str:
         """生成回答 —— 可重写
@@ -290,34 +341,52 @@ class Agent(ABC):
         return response  # type:ignore
 
     async def _process_interview(self, payload: dict):
-        auros = [
-            {
-                "id": self._uuid,
-                "day": await self.simulator.get_simulator_day(),
-                "t": await self.simulator.get_simulator_second_from_start_of_day(),
-                "type": 2,
-                "speaker": "user",
-                "content": payload["content"],
-                "created_at": int(datetime.now().timestamp() * 1000),
-            }
-        ]
+        pg_list: list[tuple[dict, datetime]] = []
+        auros: list[dict] = []
+        _date_time = datetime.now(timezone.utc)
+        _interview_dict = {
+            "id": self._uuid,
+            "day": await self.simulator.get_simulator_day(),
+            "t": await self.simulator.get_simulator_second_from_start_of_day(),
+            "type": 2,
+            "speaker": "user",
+            "content": payload["content"],
+            "created_at": int(_date_time.timestamp() * 1000),
+        }
+        auros.append(_interview_dict)
+        pg_list.append((_interview_dict, _date_time))
         question = payload["content"]
         response = await self.generate_user_chat_response(question)
-        auros.append(
-            {
-                "id": self._uuid,
-                "day": await self.simulator.get_simulator_day(),
-                "t": await self.simulator.get_simulator_second_from_start_of_day(),
-                "type": 2,
-                "speaker": "",
-                "content": response,
-                "created_at": int(datetime.now().timestamp() * 1000),
-            }
-        )
-        if self._avro_file is None:
-            return
-        with open(self._avro_file["dialog"], "a+b") as f:
-            fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+        _date_time = datetime.now(timezone.utc)
+        _interview_dict = {
+            "id": self._uuid,
+            "day": await self.simulator.get_simulator_day(),
+            "t": await self.simulator.get_simulator_second_from_start_of_day(),
+            "type": 2,
+            "speaker": "",
+            "content": response,
+            "created_at": int(_date_time.timestamp() * 1000),
+        }
+        auros.append(_interview_dict)
+        pg_list.append((_interview_dict, _date_time))
+        # Avro
+        if self._avro_file is not None:
+            with open(self._avro_file["dialog"], "a+b") as f:
+                fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+        # Pg
+        if self._pgsql_writer is not None:
+            if self._last_asyncio_pg_task is not None:
+                await self._last_asyncio_pg_task
+            _keys = ["id", "day", "t", "type", "speaker", "content", "created_at"]
+            _data = [
+                tuple([_dict[k] if k != "created_at" else _date_time for k in _keys])
+                for _dict, _date_time in pg_list
+            ]
+            self._last_asyncio_pg_task = (
+                self._pgsql_writer.async_write_dialog.remote(  # type:ignore
+                    _data
+                )
+            )
 
     async def process_agent_chat_response(self, payload: dict) -> str:
         resp = f"Agent {self._uuid} received agent chat response: {payload}"
@@ -325,22 +394,39 @@ class Agent(ABC):
         return resp
 
     async def _process_agent_chat(self, payload: dict):
-        auros = [
-            {
-                "id": self._uuid,
-                "day": payload["day"],
-                "t": payload["t"],
-                "type": 1,
-                "speaker": payload["from"],
-                "content": payload["content"],
-                "created_at": int(datetime.now().timestamp() * 1000),
-            }
-        ]
+        pg_list: list[tuple[dict, datetime]] = []
+        auros: list[dict] = []
+        _date_time = datetime.now(timezone.utc)
+        _chat_dict = {
+            "id": self._uuid,
+            "day": payload["day"],
+            "t": payload["t"],
+            "type": 1,
+            "speaker": payload["from"],
+            "content": payload["content"],
+            "created_at": int(_date_time.timestamp() * 1000),
+        }
+        auros.append(_chat_dict)
+        pg_list.append((_chat_dict, _date_time))
         asyncio.create_task(self.process_agent_chat_response(payload))
-        if self._avro_file is None:
-            return
-        with open(self._avro_file["dialog"], "a+b") as f:
-            fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+        # Avro
+        if self._avro_file is not None:
+            with open(self._avro_file["dialog"], "a+b") as f:
+                fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+        # Pg
+        if self._pgsql_writer is not None:
+            if self._last_asyncio_pg_task is not None:
+                await self._last_asyncio_pg_task
+            _keys = ["id", "day", "t", "type", "speaker", "content", "created_at"]
+            _data = [
+                tuple([_dict[k] if k != "created_at" else _date_time for k in _keys])
+                for _dict, _date_time in pg_list
+            ]
+            self._last_asyncio_pg_task = (
+                self._pgsql_writer.async_write_dialog.remote(  # type:ignore
+                    _data
+                )
+            )
 
     # Callback functions for MQTT message
     async def handle_agent_chat_message(self, payload: dict):
@@ -384,21 +470,38 @@ class Agent(ABC):
             "t": await self.simulator.get_simulator_second_from_start_of_day(),
         }
         await self._send_message(to_agent_uuid, payload, "agent-chat")
-        auros = [
-            {
-                "id": self._uuid,
-                "day": await self.simulator.get_simulator_day(),
-                "t": await self.simulator.get_simulator_second_from_start_of_day(),
-                "type": 1,
-                "speaker": self._uuid,
-                "content": content,
-                "created_at": int(datetime.now().timestamp() * 1000),
-            }
-        ]
-        if self._avro_file is None:
-            return
-        with open(self._avro_file["dialog"], "a+b") as f:
-            fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+        pg_list: list[tuple[dict, datetime]] = []
+        auros: list[dict] = []
+        _date_time = datetime.now(timezone.utc)
+        _message_dict = {
+            "id": self._uuid,
+            "day": await self.simulator.get_simulator_day(),
+            "t": await self.simulator.get_simulator_second_from_start_of_day(),
+            "type": 1,
+            "speaker": self._uuid,
+            "content": content,
+            "created_at": int(datetime.now().timestamp() * 1000),
+        }
+        auros.append(_message_dict)
+        pg_list.append((_message_dict, _date_time))
+        # Avro
+        if self._avro_file is not None:
+            with open(self._avro_file["dialog"], "a+b") as f:
+                fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+        # Pg
+        if self._pgsql_writer is not None:
+            if self._last_asyncio_pg_task is not None:
+                await self._last_asyncio_pg_task
+            _keys = ["id", "day", "t", "type", "speaker", "content", "created_at"]
+            _data = [
+                tuple([_dict[k] if k != "created_at" else _date_time for k in _keys])
+                for _dict, _date_time in pg_list
+            ]
+            self._last_asyncio_pg_task = (
+                self._pgsql_writer.async_write_dialog.remote(  # type:ignore
+                    _data
+                )
+            )
 
     # Agent logic
     @abstractmethod

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,7 @@ class AgentGroup:
         enable_avro: bool,
         avro_path: Path,
         enable_pgsql: bool,
-        pgsql_copy_writer: ray.ObjectRef,
+        pgsql_writer: ray.ObjectRef,
         mlflow_run_id: str,
         logging_level: int,
     ):
@@ -45,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)
@@ -54,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"],
@@ -61,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
@@ -105,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} ...")
@@ -161,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")
 
@@ -218,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"]
@@ -248,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(),
@@ -274,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:

pycityagent/simulation/simulation.py

@@ -23,6 +23,7 @@ from ..message.messager import Messager
 from ..metrics import init_mlflow_connection
 from ..survey import Survey
 from .agentgroup import AgentGroup
+from .storage.pg import PgWriter, create_pg_tables
 
 logger = logging.getLogger("pycityagent")
 
@@ -63,6 +64,7 @@ class AgentSimulation:
         self._user_survey_topics: dict[uuid.UUID, str] = {}
         self._user_interview_topics: dict[uuid.UUID, str] = {}
         self._loop = asyncio.get_event_loop()
+        # self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
 
         self._messager = Messager(
             hostname=config["simulator_request"]["mqtt"]["server"],
@@ -89,22 +91,13 @@ class AgentSimulation:
         self._enable_pgsql = _pgsql_config.get("enabled", False)
         if not self._enable_pgsql:
             logger.warning("PostgreSQL is not enabled, NO POSTGRESQL DATABASE STORAGE")
-            self._pgsql_args = ("", "", "", "", "")
+            self._pgsql_dsn = ""
         else:
-            self._pgsql_host = _pgsql_config["host"]
-            self._pgsql_port = _pgsql_config["port"]
-            self._pgsql_database = _pgsql_config["database"]
-            self._pgsql_user = _pgsql_config.get("user", None)
-            self._pgsql_password = _pgsql_config.get("password", None)
-            self._pgsql_args: tuple[str, str, str, str, str] = (
-                self._pgsql_host,
-                self._pgsql_port,
-                self._pgsql_database,
-                self._pgsql_user,
-                self._pgsql_password,
-            )
+            self._pgsql_dsn = _pgsql_config["data_source_name"]
 
         # 添加实验信息相关的属性
+        self._exp_created_time = datetime.now(timezone.utc)
+        self._exp_updated_time = datetime.now(timezone.utc)
         self._exp_info = {
             "id": self.exp_id,
             "name": exp_name,
@@ -114,7 +107,8 @@ class AgentSimulation:
             "cur_t": 0.0,
             "config": json.dumps(config),
             "error": "",
-            "created_at": datetime.now(timezone.utc).isoformat(),
+            "created_at": self._exp_created_time.isoformat(),
+            "updated_at": self._exp_updated_time.isoformat(),
         }
 
         # 创建异步任务保存实验信息
@@ -168,7 +162,7 @@ class AgentSimulation:
         enable_avro: bool,
         avro_path: Path,
         enable_pgsql: bool,
-        pgsql_copy_writer: ray.ObjectRef,
+        pgsql_writer: ray.ObjectRef,
         mlflow_run_id: str = None,  # type: ignore
         logging_level: int = logging.WARNING,
     ):
@@ -181,7 +175,7 @@ class AgentSimulation:
             enable_avro,
             avro_path,
             enable_pgsql,
-            pgsql_copy_writer,
+            pgsql_writer,
             mlflow_run_id,
             logging_level,
         )
@@ -191,6 +185,7 @@ class AgentSimulation:
         self,
         agent_count: Union[int, list[int]],
         group_size: int = 1000,
+        pg_sql_writers: int = 32,
         memory_config_func: Optional[Union[Callable, list[Callable]]] = None,
     ) -> None:
         """初始化智能体
@@ -251,8 +246,8 @@ class AgentSimulation:
                     memory=memory,
                 )
 
-                self._agents[agent._uuid] = agent
-                self._agent_uuids.append(agent._uuid)
+                self._agents[agent._uuid] = agent  # type:ignore
+                self._agent_uuids.append(agent._uuid)  # type:ignore
 
             # 计算需要的组数,向上取整以处理不足一组的情况
             num_group = (agent_count_i + group_size - 1) // group_size
@@ -282,9 +277,23 @@ class AgentSimulation:
             )
         else:
             mlflow_run_id = None
+        # 建表
+        if self.enable_pgsql:
+            _num_workers = min(1, pg_sql_writers)
+            create_pg_tables(
+                exp_id=self.exp_id,
+                dsn=self._pgsql_dsn,
+            )
+            self._pgsql_writers = _workers = [
+                PgWriter.remote(self.exp_id, self._pgsql_dsn)
+                for _ in range(_num_workers)
+            ]
+        else:
+            _num_workers = 1
+            self._pgsql_writers = _workers = [None for _ in range(_num_workers)]
         # 收集所有创建组的参数
         creation_tasks = []
-        for group_name, agents in group_creation_params:
+        for i, (group_name, agents) in enumerate(group_creation_params):
             # 直接创建异步任务
             group = AgentGroup.remote(
                 agents,
@@ -294,10 +303,8 @@ class AgentSimulation:
                 self.enable_avro,
                 self.avro_path,
                 self.enable_pgsql,
-                # TODO:
-                # self._pgsql_copy_writer, # type:ignore
-                None,
-                mlflow_run_id,
+                _workers[i % _num_workers],  # type:ignore
+                mlflow_run_id,  # type:ignore
                 self.logging_level,
             )
             creation_tasks.append((group_name, group, agents))
@@ -469,11 +476,13 @@ class AgentSimulation:
         survey_dict = survey.to_dict()
         if agent_uuids is None:
             agent_uuids = self._agent_uuids
+        _date_time = datetime.now(timezone.utc)
         payload = {
             "from": "none",
             "survey_id": survey_dict["id"],
-            "timestamp": int(datetime.now().timestamp() * 1000),
+            "timestamp": int(_date_time.timestamp() * 1000),
             "data": survey_dict,
+            "_date_time": _date_time,
         }
         for uuid in agent_uuids:
             topic = self._user_survey_topics[uuid]
@@ -483,10 +492,12 @@ class AgentSimulation:
         self, content: str, agent_uuids: Union[uuid.UUID, list[uuid.UUID]]
     ):
         """发送面试消息"""
+        _date_time = datetime.now(timezone.utc)
         payload = {
             "from": "none",
             "content": content,
-            "timestamp": int(datetime.now().timestamp() * 1000),
+            "timestamp": int(_date_time.timestamp() * 1000),
+            "_date_time": _date_time,
         }
         if not isinstance(agent_uuids, Sequence):
             agent_uuids = [agent_uuids]
@@ -515,15 +526,29 @@ class AgentSimulation:
             logger.error(f"Avro保存实验信息失败: {str(e)}")
         try:
             if self.enable_pgsql:
-                #   TODO
-                pass
+                worker: ray.ObjectRef = self._pgsql_writers[0]  # type:ignore
+                # if self._last_asyncio_pg_task is not None:
+                #     await self._last_asyncio_pg_task
+                # self._last_asyncio_pg_task = (
+                #     worker.async_update_exp_info.remote(  # type:ignore
+                #         pg_exp_info
+                #     )
+                # )
+                pg_exp_info = {k: v for k, v in self._exp_info.items()}
+                pg_exp_info["created_at"] = self._exp_created_time
+                pg_exp_info["updated_at"] = self._exp_updated_time
+                await worker.async_update_exp_info.remote(  # type:ignore
+                    pg_exp_info
+                )
         except Exception as e:
             logger.error(f"PostgreSQL保存实验信息失败: {str(e)}")
 
     async def _update_exp_status(self, status: int, error: str = "") -> None:
+        self._exp_updated_time = datetime.now(timezone.utc)
         """更新实验状态并保存"""
         self._exp_info["status"] = status
         self._exp_info["error"] = error
+        self._exp_info["updated_at"] = self._exp_updated_time.isoformat()
         await self._save_exp_info()
 
     async def _monitor_exp_status(self, stop_event: asyncio.Event):

pycityagent/simulation/storage/pg.py

@@ -0,0 +1,139 @@
+import asyncio
+from collections import defaultdict
+from typing import Any
+
+import psycopg
+import psycopg.sql
+import ray
+from psycopg.rows import dict_row
+
+from ...utils.decorators import lock_decorator
+from ...utils.pg_query import PGSQL_DICT
+
+
+def create_pg_tables(exp_id: str, dsn: str):
+    for table_type, exec_strs in PGSQL_DICT.items():
+        table_name = f"socialcity_{exp_id.replace('-', '_')}_{table_type}"
+        # # debug str
+        # for _str in [f"DROP TABLE IF EXISTS {table_name}"] + [
+        #     _exec_str.format(table_name=table_name) for _exec_str in exec_strs
+        # ]:
+        #     print(_str)
+        with psycopg.connect(dsn) as conn:
+            with conn.cursor() as cur:
+                # delete table
+                cur.execute(f"DROP TABLE IF EXISTS {table_name}")  # type:ignore
+                conn.commit()
+                # create table
+                for _exec_str in exec_strs:
+                    cur.execute(_exec_str.format(table_name=table_name))
+                conn.commit()
+
+
+@ray.remote
+class PgWriter:
+    def __init__(self, exp_id: str, dsn: str):
+        self.exp_id = exp_id
+        self._dsn = dsn
+        # self._lock = asyncio.Lock()
+
+    # @lock_decorator
+    async def async_write_dialog(self, rows: list[tuple]):
+        _tuple_types = [str, int, float, int, str, str, str, None]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_agent_dialog"
+        # 将数据插入数据库
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            copy_sql = psycopg.sql.SQL(
+                "COPY {} (id, day, t, type, speaker, content, created_at) FROM STDIN"
+            ).format(psycopg.sql.Identifier(table_name))
+            async with aconn.cursor() as cur:
+                async with cur.copy(copy_sql) as copy:
+                    for row in rows:
+                        _row = [
+                            _type(r) if _type is not None else r
+                            for (_type, r) in zip(_tuple_types, row)
+                        ]
+                        await copy.write_row(_row)
+
+    # @lock_decorator
+    async def async_write_status(self, rows: list[tuple]):
+        _tuple_types = [str, int, float, float, float, int, str, str, None]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_agent_status"
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            copy_sql = psycopg.sql.SQL(
+                "COPY {} (id, day, t, lng, lat, parent_id, action, status, created_at) FROM STDIN"
+            ).format(psycopg.sql.Identifier(table_name))
+            async with aconn.cursor() as cur:
+                async with cur.copy(copy_sql) as copy:
+                    for row in rows:
+                        _row = [
+                            _type(r) if _type is not None else r
+                            for (_type, r) in zip(_tuple_types, row)
+                        ]
+                        await copy.write_row(_row)
+
+    # @lock_decorator
+    async def async_write_profile(self, rows: list[tuple]):
+        _tuple_types = [str, str, str]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_agent_profile"
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            copy_sql = psycopg.sql.SQL("COPY {} (id, name, profile) FROM STDIN").format(
+                psycopg.sql.Identifier(table_name)
+            )
+            async with aconn.cursor() as cur:
+                async with cur.copy(copy_sql) as copy:
+                    for row in rows:
+                        _row = [
+                            _type(r) if _type is not None else r
+                            for (_type, r) in zip(_tuple_types, row)
+                        ]
+                        await copy.write_row(_row)
+
+    # @lock_decorator
+    async def async_write_survey(self, rows: list[tuple]):
+        _tuple_types = [str, int, float, str, str, None]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_agent_survey"
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            copy_sql = psycopg.sql.SQL(
+                "COPY {} (id, day, t, survey_id, result, created_at) FROM STDIN"
+            ).format(psycopg.sql.Identifier(table_name))
+            async with aconn.cursor() as cur:
+                async with cur.copy(copy_sql) as copy:
+                    for row in rows:
+                        _row = [
+                            _type(r) if _type is not None else r
+                            for (_type, r) in zip(_tuple_types, row)
+                        ]
+                        await copy.write_row(_row)
+
+    # @lock_decorator
+    async def async_update_exp_info(self, exp_info: dict[str, Any]):
+        # timestamp不做类型转换
+        TO_UPDATE_EXP_INFO_KEYS_AND_TYPES = [
+            ("id", str),
+            ("name", str),
+            ("num_day", int),
+            ("status", int),
+            ("cur_day", int),
+            ("cur_t", float),
+            ("config", str),
+            ("error", str),
+            ("created_at", None),
+            ("updated_at", None),
+        ]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_experiment"
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            async with aconn.cursor(row_factory=dict_row) as cur:
+                # UPDATE
+                columns = ", ".join(
+                    f"{key} = %s" for key, _ in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                )
+                update_sql = psycopg.sql.SQL(
+                    f"UPDATE {{}} SET {columns} WHERE id = %s"  # type:ignore
+                ).format(psycopg.sql.Identifier(table_name))
+                params = [
+                    _type(exp_info[key]) if _type is not None else exp_info[key]
+                    for key, _type in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                ] + [self.exp_id]
+                await cur.execute(update_sql, params)
+                await aconn.commit()

pycityagent/utils/pg_query.py

@@ -0,0 +1,80 @@
+from typing import Any
+
+PGSQL_DICT: dict[str, list[Any]] = {
+    # Experiment
+    "experiment": [
+        """
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        id UUID PRIMARY KEY,
+        name TEXT,
+        num_day INT4,
+        status INT4, 
+        cur_day INT4,
+        cur_t FLOAT,
+        config TEXT,
+        error TEXT,
+        created_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+        updated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+    )
+""",
+    ],
+    # Agent Profile
+    "agent_profile": [
+        """
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        id UUID PRIMARY KEY,
+        name TEXT,
+        profile JSONB
+    )
+""",
+    ],
+    # Agent Dialog
+    "agent_dialog": [
+        """
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        id UUID,
+        day INT4,
+        t FLOAT,
+        type INT4,
+        speaker TEXT,
+        content TEXT,
+        created_at TIMESTAMPTZ
+    )
+""",
+        "CREATE INDEX {table_name}_id_idx ON {table_name} (id)",
+        "CREATE INDEX {table_name}_day_t_idx ON {table_name} (day,t)",
+    ],
+    # Agent Status
+    "agent_status": [
+        """
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        id UUID,
+        day INT4,
+        t FLOAT,
+        lng DOUBLE PRECISION,
+        lat DOUBLE PRECISION,
+        parent_id INT4,
+        action TEXT,
+        status JSONB,
+        created_at TIMESTAMPTZ
+    )
+""",
+        "CREATE INDEX {table_name}_id_idx ON {table_name} (id)",
+        "CREATE INDEX {table_name}_day_t_idx ON {table_name} (day,t)",
+    ],
+    # Agent Survey
+    "agent_survey": [
+        """
+    CREATE TABLE IF NOT EXISTS {table_name} (
+        id UUID,
+        day INT4,
+        t FLOAT,
+        survey_id UUID,
+        result JSONB,
+        created_at TIMESTAMPTZ
+    )
+""",
+        "CREATE INDEX {table_name}_id_idx ON {table_name} (id)",
+        "CREATE INDEX {table_name}_day_t_idx ON {table_name} (day,t)",
+    ],
+}

pycityagent/workflow/tool.py

@@ -1,6 +1,8 @@
 import time
+from collections import defaultdict
+from collections.abc import Callable, Sequence
 from typing import Any, Optional, Union
-from collections.abc import Callable
+
 from mlflow.entities import Metric
 
 from ..agent import Agent
@@ -190,33 +192,38 @@ class ResetAgentPosition(Tool):
 class ExportMlflowMetrics(Tool):
     def __init__(self, log_batch_size: int = 100) -> None:
         self._log_batch_size = log_batch_size
-        # TODO:support other log types
-        self.metric_log_cache: list[Metric] = []
+        # TODO: support other log types
+        self.metric_log_cache: dict[str, list[Metric]] = defaultdict(list)
 
     async def __call__(
         self,
-        metric: Union[Metric, dict],
+        metric: Union[Sequence[Union[Metric, dict]], Union[Metric, dict]],
         clear_cache: bool = False,
     ):
         agent = self.agent
         batch_size = self._log_batch_size
-        if len(self.metric_log_cache) > batch_size:
-            client = agent.mlflow_client
-            await client.log_batch(
-                metrics=self.metric_log_cache[:batch_size],
-            )
-            self.metric_log_cache = self.metric_log_cache[batch_size:]
-        else:
-            if isinstance(metric, Metric):
-                self.metric_log_cache.append(metric)
+        if not isinstance(metric, Sequence):
+            metric = [metric]
+        for _metric in metric:
+            if isinstance(_metric, Metric):
+                item = _metric
+                metric_key = item.key
             else:
-                _metric = Metric(
-                    key=metric["key"],
-                    value=metric["value"],
-                    timestamp=metric.get("timestamp", int(1000 * time.time())),
-                    step=metric["step"],
+                item = Metric(
+                    key=_metric["key"],
+                    value=_metric["value"],
+                    timestamp=_metric.get("timestamp", int(1000 * time.time())),
+                    step=_metric["step"],
                 )
-                self.metric_log_cache.append(_metric)
+                metric_key = _metric["key"]
+            self.metric_log_cache[metric_key].append(item)
+        for metric_key, _cache in self.metric_log_cache.items():
+            if len(_cache) > batch_size:
+                client = agent.mlflow_client
+                await client.log_batch(
+                    metrics=_cache[:batch_size],
+                )
+                _cache = _cache[batch_size:]
         if clear_cache:
             await self._clear_cache()
 
@@ -225,8 +232,9 @@ class ExportMlflowMetrics(Tool):
     ):
         agent = self.agent
         client = agent.mlflow_client
-        if len(self.metric_log_cache) > 0:
-            await client.log_batch(
-                metrics=self.metric_log_cache,
-            )
-            self.metric_log_cache = []
+        for metric_key, _cache in self.metric_log_cache.items():
+            if len(_cache) > 0:
+                await client.log_batch(
+                    metrics=_cache,
+                )
+                _cache = []

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

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

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

@@ -1,5 +1,5 @@
 pycityagent/__init__.py,sha256=EDxt3Su3lH1IMh9suNw7GeGL7UrXeWiZTw5KWNznDzc,637
-pycityagent/agent.py,sha256=HRFAG_iM1q3nvXtV0T-Dz01foOtty4IWka7h4WD97CU,24268
+pycityagent/agent.py,sha256=FMplKAgcz2Exkl8EiE2RwQ0Hd5U08krRZ3CFFLoF_4g,28450
 pycityagent/economy/__init__.py,sha256=aonY4WHnx-6EGJ4WKrx4S-2jAkYNLtqUA04jp6q8B7w,75
 pycityagent/economy/econ_client.py,sha256=GuHK9ZBnhqW3Z7F8ViDJn_iN73yOBbbwFyJv1wLEBDk,12211
 pycityagent/environment/__init__.py,sha256=awHxlOud-btWbk0FCS4RmGJ13W84oVCkbGfcrhKqihA,240
@@ -48,9 +48,10 @@ pycityagent/message/messager.py,sha256=W_OVlNGcreHSBf6v-DrEnfNCXExB78ySr0w26MSnc
 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=6p8-OP2x_syaQ1pWLgll0LHs823OHnappuY-8XzL_LU,13383
-pycityagent/simulation/simulation.py,sha256=pxFdaNMLdJMcpW8gnOcxjfbz4gkpcvWsTR0rNP8rE9k,21675
+pycityagent/simulation/__init__.py,sha256=P5czbcg2d8S0nbbnsQXFIhwzO4CennAhZM8OmKvAeYw,194
+pycityagent/simulation/agentgroup.py,sha256=5p68wNoEaog4nDym3xsCTporBWmxNiQ1crN3mbOHFsE,19788
+pycityagent/simulation/simulation.py,sha256=7Go_RkpkC_DuBWW21JPqlV2yXY754RqSkqzM0vTdteU,23008
+pycityagent/simulation/storage/pg.py,sha256=Ws04mUgRcbbvWi_eQm3PXYa6w7AQUbDPWhSU7HFtsD8,6026
 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
@@ -61,12 +62,13 @@ pycityagent/utils/parsers/__init__.py,sha256=AN2xgiPxszWK4rpX7zrqRsqNwfGF3WnCA5-
 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=KBKO4zLZPNdGjPAGqIus8LseZ8W3Tlt2y0QxqeCd25Q,1713
+pycityagent/utils/pg_query.py,sha256=h5158xcrxjUTR0nKwAaG1neFfTHPbN5guLmaXpC8yvs,1918
 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=6jI0Rq54JLv3-IXqZLYug62vse10wTI83xvf4ZX42nk,2929
-pycityagent/workflow/tool.py,sha256=SGY18lT71hBLKagopirFbxRjPY_387Dobo9SUwjHIn0,8215
+pycityagent/workflow/tool.py,sha256=xADxhNgVsjNiMxlhdwn3xGUstFOkLEG8P67ez8VmwSI,8555
 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,,
+pycityagent-2.0.0a21.dist-info/METADATA,sha256=sowWsIPV6PFjNPeQI30Pn0J1Fqz5KfZ7sMydvfaOAX0,7848
+pycityagent-2.0.0a21.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+pycityagent-2.0.0a21.dist-info/RECORD,,