pycityagent 2.0.0a37__tar.gz → 2.0.0a38__tar.gz

{pycityagent-2.0.0a37 → pycityagent-2.0.0a38}/PKG-INFO

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

{pycityagent-2.0.0a37 → pycityagent-2.0.0a38}/pycityagent/simulation/simulation.py

@@ -20,10 +20,11 @@ from mosstool.map._map_util.const import AOI_START_ID
 from ..agent import Agent, InstitutionAgent
 from ..environment.simulator import Simulator
 from ..llm import SimpleEmbedding
-from ..memory import FaissQuery, Memory
+from ..memory import Memory
 from ..message.messager import Messager
 from ..metrics import init_mlflow_connection
 from ..survey import Survey
+from ..utils import TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
 from .agentgroup import AgentGroup
 from .storage.pg import PgWriter, create_pg_tables
 
@@ -96,7 +97,11 @@ class AgentSimulation:
             logger.warning("PostgreSQL is not enabled, NO POSTGRESQL DATABASE STORAGE")
             self._pgsql_dsn = ""
         else:
-            self._pgsql_dsn = _pgsql_config["data_source_name"]
+            self._pgsql_dsn = (
+                _pgsql_config["data_source_name"]
+                if "data_source_name" in _pgsql_config
+                else _pgsql_config["dsn"]
+            )
 
         # 添加实验信息相关的属性
         self._exp_created_time = datetime.now(timezone.utc)
@@ -535,7 +540,9 @@ class AgentSimulation:
         try:
             if self.enable_pgsql:
                 worker: ray.ObjectRef = self._pgsql_writers[0]  # type:ignore
-                pg_exp_info = {k: v for k, v in self._exp_info.items()}
+                pg_exp_info = {
+                    k: self._exp_info[k] for (k, _) in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                }
                 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

{pycityagent-2.0.0a37 → pycityagent-2.0.0a38}/pycityagent/simulation/storage/pg.py

@@ -1,4 +1,5 @@
 import asyncio
+import logging
 from collections import defaultdict
 from typing import Any
 
@@ -8,7 +9,9 @@ import ray
 from psycopg.rows import dict_row
 
 from ...utils.decorators import lock_decorator
-from ...utils.pg_query import PGSQL_DICT
+from ...utils.pg_query import PGSQL_DICT, TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+
+logger = logging.getLogger("pg")
 
 
 def create_pg_tables(exp_id: str, dsn: str):
@@ -27,10 +30,15 @@ def create_pg_tables(exp_id: str, dsn: str):
                 if not table_type == "experiment":
                     # delete table
                     cur.execute(f"DROP TABLE IF EXISTS {table_name}")  # type:ignore
+                    logger.debug(
+                        f"table:{table_name} sql: DROP TABLE IF EXISTS {table_name}"
+                    )
                     conn.commit()
                 # create table
                 for _exec_str in exec_strs:
-                    cur.execute(_exec_str.format(table_name=table_name))
+                    exec_str = _exec_str.format(table_name=table_name)
+                    cur.execute(exec_str)
+                    logger.debug(f"table:{table_name} sql: {exec_str}")
                 conn.commit()
 
 
@@ -50,6 +58,7 @@ class PgWriter:
             copy_sql = psycopg.sql.SQL(
                 "COPY {} (id, day, t, type, speaker, content, created_at) FROM STDIN"
             ).format(psycopg.sql.Identifier(table_name))
+            _rows: list[Any] = []
             async with aconn.cursor() as cur:
                 async with cur.copy(copy_sql) as copy:
                     for row in rows:
@@ -58,6 +67,8 @@ class PgWriter:
                             for (_type, r) in zip(_tuple_types, row)
                         ]
                         await copy.write_row(_row)
+                        _rows.append(_row)
+            logger.debug(f"table:{table_name} sql: {copy_sql} values: {_rows}")
 
     # @lock_decorator
     async def async_write_status(self, rows: list[tuple]):
@@ -67,6 +78,7 @@ class PgWriter:
             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))
+            _rows: list[Any] = []
             async with aconn.cursor() as cur:
                 async with cur.copy(copy_sql) as copy:
                     for row in rows:
@@ -75,6 +87,8 @@ class PgWriter:
                             for (_type, r) in zip(_tuple_types, row)
                         ]
                         await copy.write_row(_row)
+                        _rows.append(_row)
+            logger.debug(f"table:{table_name} sql: {copy_sql} values: {_rows}")
 
     # @lock_decorator
     async def async_write_profile(self, rows: list[tuple]):
@@ -84,6 +98,7 @@ class PgWriter:
             copy_sql = psycopg.sql.SQL("COPY {} (id, name, profile) FROM STDIN").format(
                 psycopg.sql.Identifier(table_name)
             )
+            _rows: list[Any] = []
             async with aconn.cursor() as cur:
                 async with cur.copy(copy_sql) as copy:
                     for row in rows:
@@ -92,6 +107,8 @@ class PgWriter:
                             for (_type, r) in zip(_tuple_types, row)
                         ]
                         await copy.write_row(_row)
+                        _rows.append(_row)
+            logger.debug(f"table:{table_name} sql: {copy_sql} values: {_rows}")
 
     # @lock_decorator
     async def async_write_survey(self, rows: list[tuple]):
@@ -101,6 +118,7 @@ class PgWriter:
             copy_sql = psycopg.sql.SQL(
                 "COPY {} (id, day, t, survey_id, result, created_at) FROM STDIN"
             ).format(psycopg.sql.Identifier(table_name))
+            _rows: list[Any] = []
             async with aconn.cursor() as cur:
                 async with cur.copy(copy_sql) as copy:
                     for row in rows:
@@ -109,51 +127,54 @@ class PgWriter:
                             for (_type, r) in zip(_tuple_types, row)
                         ]
                         await copy.write_row(_row)
+                        _rows.append(_row)
+            logger.debug(f"table:{table_name} sql: {copy_sql} values: {_rows}")
 
     # @lock_decorator
     async def async_update_exp_info(self, exp_info: dict[str, Any]):
         # timestamp不做类型转换
-        TO_UPDATE_EXP_INFO_KEYS_AND_TYPES = [
-            ("id", None),
-            ("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_experiment"
         async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
-            async with aconn.cursor(row_factory=dict_row) as cur:                
-                await cur.execute("SELECT * FROM {table_name} WHERE id=%s".format(table_name = table_name),(self.exp_id,))# type:ignore
-                record_exists = await cur.fetchall()    
-                print("record_exists",record_exists)            
+            async with aconn.cursor(row_factory=dict_row) as cur:
+                exec_str = "SELECT * FROM {table_name} WHERE id=%s".format(
+                    table_name=table_name
+                ), (self.exp_id,)
+                await cur.execute(exec_str)  # type:ignore
+                logger.debug(f"table:{table_name} sql: {exec_str}")
+                record_exists = await cur.fetchall()
                 if record_exists:
                     # 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='{self.exp_id}'" # type:ignore
+                        f"UPDATE {{}} SET {columns} WHERE id='{self.exp_id}'"  # 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]
+                    ]
+                    logger.debug(
+                        f"table:{table_name} sql: {update_sql} values: {params}"
+                    )
                     await cur.execute(update_sql, params)
                 else:
                     # INSERT
-                    keys = ", ".join(key for key, _ in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES)
-                    placeholders = ", ".join(["%s"] * len(TO_UPDATE_EXP_INFO_KEYS_AND_TYPES))
+                    keys = ", ".join(
+                        key for key, _ in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                    )
+                    placeholders = ", ".join(
+                        ["%s"] * len(TO_UPDATE_EXP_INFO_KEYS_AND_TYPES)
+                    )
                     insert_sql = psycopg.sql.SQL(
-                        f"INSERT INTO {{}} ({keys}) VALUES ({placeholders})" # type:ignore
+                        f"INSERT INTO {{}} ({keys}) VALUES ({placeholders})"  # 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
                     ]
+                    logger.debug(
+                        f"table:{table_name} sql: {insert_sql} values: {params}"
+                    )
                     await cur.execute(insert_sql, params)
                 await aconn.commit()

pycityagent-2.0.0a38/pycityagent/utils/__init__.py

@@ -0,0 +1,11 @@
+from .avro_schema import (DIALOG_SCHEMA, INSTITUTION_STATUS_SCHEMA,
+                          PROFILE_SCHEMA, STATUS_SCHEMA, SURVEY_SCHEMA)
+from .pg_query import PGSQL_DICT, TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+from .survey_util import process_survey_for_llm
+
+__all__ = [
+    "PROFILE_SCHEMA", "DIALOG_SCHEMA", "STATUS_SCHEMA", "SURVEY_SCHEMA", "INSTITUTION_STATUS_SCHEMA",
+    "process_survey_for_llm",
+    "TO_UPDATE_EXP_INFO_KEYS_AND_TYPES",
+    "PGSQL_DICT",
+]

{pycityagent-2.0.0a37 → pycityagent-2.0.0a38}/pycityagent/utils/pg_query.py

@@ -78,3 +78,15 @@ PGSQL_DICT: dict[str, list[Any]] = {
         "CREATE INDEX {table_name}_day_t_idx ON {table_name} (day,t)",
     ],
 }
+TO_UPDATE_EXP_INFO_KEYS_AND_TYPES: list[tuple[str, Any]] = [
+    ("id", None),
+    ("name", str),
+    ("num_day", int),
+    ("status", int),
+    ("cur_day", int),
+    ("cur_t", float),
+    ("config", str),
+    ("error", str),
+    ("created_at", None),
+    ("updated_at", None),
+]

{pycityagent-2.0.0a37 → pycityagent-2.0.0a38}/pyproject.toml

@@ -3,7 +3,7 @@ in-project = true
 
 [tool.poetry]
 name = "pycityagent"
-version = "2.0.0a37"
+version = "2.0.0a38"
 description = "LLM-based城市环境agent构建库"
 authors = ["Yuwei Yan <pinkgranite86@gmail.com>","Junbo Yan <yanjb20thu@gmali.com>"]
 license = "MIT"

pycityagent-2.0.0a37/pycityagent/utils/__init__.py

@@ -1,7 +0,0 @@
-from .avro_schema import PROFILE_SCHEMA, DIALOG_SCHEMA, STATUS_SCHEMA, SURVEY_SCHEMA, INSTITUTION_STATUS_SCHEMA
-from .survey_util import process_survey_for_llm
-
-__all__ = [
-    "PROFILE_SCHEMA", "DIALOG_SCHEMA", "STATUS_SCHEMA", "SURVEY_SCHEMA", "INSTITUTION_STATUS_SCHEMA",
-    "process_survey_for_llm"
-]