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

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,8 +20,10 @@ 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
+from .storage.pg import PgWriter, create_pg_tables
 
 logger = logging.getLogger("pycityagent")
 
@@ -60,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"],
@@ -86,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,
@@ -111,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(),
         }
 
         # 创建异步任务保存实验信息
@@ -165,7 +162,8 @@ class AgentSimulation:
         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 = None,  # type: ignore
         logging_level: int = logging.WARNING,
     ):
         """创建远程组"""
@@ -177,7 +175,8 @@ class AgentSimulation:
             enable_avro,
             avro_path,
             enable_pgsql,
-            pgsql_args,
+            pgsql_writer,
+            mlflow_run_id,
             logging_level,
         )
         return group_name, group, agents
@@ -186,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:
         """初始化智能体
@@ -246,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
@@ -267,9 +267,33 @@ 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
+        # 建表
+        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,
@@ -279,7 +303,8 @@ class AgentSimulation:
                 self.enable_avro,
                 self.avro_path,
                 self.enable_pgsql,
-                self._pgsql_args,
+                _workers[i % _num_workers],  # type:ignore
+                mlflow_run_id,  # type:ignore
                 self.logging_level,
             )
             creation_tasks.append((group_name, group, agents))
@@ -451,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]
@@ -465,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]
@@ -497,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/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/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/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,5 +1,7 @@
 import time
-from typing import Any, Callable, Dict, List, Optional, Union
+from collections import defaultdict
+from collections.abc import Callable, Sequence
+from typing import Any, Optional, Union
 
 from mlflow.entities import Metric
 
@@ -76,7 +78,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.
@@ -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"],
+                )
+                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],
                 )
-                self.metric_log_cache.append(_metric)
+                _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/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.0a21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a19
+Version: 2.0.0a21
 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)