pycityagent 2.0.0a20__tar.gz → 2.0.0a21__tar.gz

{pycityagent-2.0.0a20 → pycityagent-2.0.0a21}/PKG-INFO

@@ -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 → pycityagent-2.0.0a21}/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-2.0.0a21/pycityagent/simulation/__init__.py

@@ -0,0 +1,8 @@
+"""
+城市智能体模拟器模块
+"""
+
+from .simulation import AgentSimulation
+from .storage.pg import PgWriter, create_pg_tables
+
+__all__ = ["AgentSimulation", "PgWriter", "create_pg_tables"]

{pycityagent-2.0.0a20 → pycityagent-2.0.0a21}/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: