pycityagent 2.0.0a28__py3-none-any.whl → 2.0.0a29__py3-none-any.whl

pycityagent/simulation/storage/pg.py

@@ -106,7 +106,6 @@ class PgWriter:
                         ]
                         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 = [
@@ -124,16 +123,38 @@ class PgWriter:
         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 cur.execute(
+                    "SELECT * FROM {table_name}".format(table_name=table_name) # type:ignore
+                )  
+                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}"  # 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)
+                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)
+                    )
+                    insert_sql = psycopg.sql.SQL(
+                        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
+                    ]
+                    await cur.execute(insert_sql, params)
                 await aconn.commit()

{pycityagent-2.0.0a28.dist-info → pycityagent-2.0.0a29.dist-info}/METADATA

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

{pycityagent-2.0.0a28.dist-info → pycityagent-2.0.0a29.dist-info}/RECORD

@@ -52,7 +52,7 @@ pycityagent/metrics/utils/const.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3
 pycityagent/simulation/__init__.py,sha256=P5czbcg2d8S0nbbnsQXFIhwzO4CennAhZM8OmKvAeYw,194
 pycityagent/simulation/agentgroup.py,sha256=0uMoRBektdK9WEXmqYT2gASg42CtS6JjaQ7slP8IWkY,22011
 pycityagent/simulation/simulation.py,sha256=9kkdgXSEOAN8wiewVFyORksti4IdVNU0opObV6ZYa9k,23344
-pycityagent/simulation/storage/pg.py,sha256=Ws04mUgRcbbvWi_eQm3PXYa6w7AQUbDPWhSU7HFtsD8,6026
+pycityagent/simulation/storage/pg.py,sha256=xJBR0ROCZpGs6d8ZDtV4iW0awatDvTDPnn6iMozDJfA,7083
 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
@@ -70,6 +70,6 @@ pycityagent/workflow/block.py,sha256=l-z9iJo9_USZQRyj4TLMfihK0-tnNDG0a6jVk9WhG0o
 pycityagent/workflow/prompt.py,sha256=6jI0Rq54JLv3-IXqZLYug62vse10wTI83xvf4ZX42nk,2929
 pycityagent/workflow/tool.py,sha256=xADxhNgVsjNiMxlhdwn3xGUstFOkLEG8P67ez8VmwSI,8555
 pycityagent/workflow/trigger.py,sha256=Df-MOBEDWBbM-v0dFLQLXteLsipymT4n8vqexmK2GiQ,5643
-pycityagent-2.0.0a28.dist-info/METADATA,sha256=2M6sidbsEZfhgXLIKYTrhumMs4LM0Qq6MWDJVMWBA-s,8033
-pycityagent-2.0.0a28.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-pycityagent-2.0.0a28.dist-info/RECORD,,
+pycityagent-2.0.0a29.dist-info/METADATA,sha256=8zdE1yf3zCJPa0d_M5CYPXULA1fG8dkeveA5EmexLJ0,8033
+pycityagent-2.0.0a29.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+pycityagent-2.0.0a29.dist-info/RECORD,,