pycityagent 2.0.0a19__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, Dict, List, Optional
+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
 
@@ -55,7 +57,8 @@ class Agent(ABC):
         simulator: Optional[Simulator] = None,
         mlflow_client: Optional[MlflowClient] = None,
         memory: Optional[Memory] = None,
-        avro_file: Optional[Dict[str, str]] = None,
+        avro_file: Optional[dict[str, str]] = None,
+        copy_writer: Optional[ray.ObjectRef] = None,
     ) -> None:
         """
         Initialize the Agent.
@@ -69,7 +72,8 @@ class Agent(ABC):
             simulator (Simulator, optional): The simulator object. Defaults to None.
             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.
+            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
@@ -85,9 +89,11 @@ class Agent(ABC):
         self._has_bound_to_simulator = False
         self._has_bound_to_economy = False
         self._blocked = False
-        self._interview_history: List[Dict] = []  # 存储采访历史
+        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()
@@ -137,12 +143,18 @@ class Agent(ABC):
         """
         self._exp_id = exp_id
 
-    def set_avro_file(self, avro_file: Dict[str, str]):
+    def set_avro_file(self, avro_file: dict[str, str]):
         """
         Set the avro file of the agent.
         """
         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
@@ -557,7 +660,7 @@ class InstitutionAgent(Agent):
             avro_file=avro_file,
         )
         # 添加响应收集器
-        self._gather_responses: Dict[str, asyncio.Future] = {}
+        self._gather_responses: dict[str, asyncio.Future] = {}
 
     async def bind_to_simulator(self):
         await self._bind_to_economy()
@@ -659,7 +762,7 @@ class InstitutionAgent(Agent):
                 }
             )
 
-    async def gather_messages(self, agent_uuids: list[str], target: str) -> List[dict]:
+    async def gather_messages(self, agent_uuids: list[str], target: str) -> list[dict]:
         """从多个智能体收集消息
 
         Args:
@@ -667,7 +770,7 @@ class InstitutionAgent(Agent):
             target: 要收集的信息类型
 
         Returns:
-            List[dict]: 收集到的所有响应
+            list[dict]: 收集到的所有响应
         """
         # 为每个agent创建Future
         futures = {}

pycityagent/economy/econ_client.py

@@ -316,3 +316,40 @@ class EconomyClient:
             await self._aio_stub.GetOrgEntityIds(request)
         )
         return list(response.org_ids)
+
+    async def add_delta_value(
+        self,
+        id: int,
+        key: str,
+        value: Any,
+    ) -> Any:
+        """
+        Add key-value pair
+
+        Args:
+        - id (int): the id of `Org` or `Agent`.
+        - key (str): the attribute to update. Can only be `inventory`, `price`, `interest_rate` and `currency`
+
+
+        Returns:
+        - Any
+        """
+        pascal_key = _snake_to_pascal(key)
+        _request_type = getattr(org_service, f"Add{pascal_key}Request")
+        _request_func = getattr(self._aio_stub, f"Add{pascal_key}")
+        _available_keys = {
+            "inventory",
+            "price",
+            "interest_rate",
+            "currency",
+        }
+        if key not in _available_keys:
+            raise ValueError(f"Invalid key `{key}`, can only be {_available_keys}!")
+        return await _request_func(
+            _request_type(
+                **{
+                    "org_id": id,
+                    f"delta_{key}": value,
+                }
+            )
+        )

pycityagent/environment/utils/geojson.py

@@ -1,9 +1,7 @@
-from typing import List
-
 __all__ = ["wrap_feature_collection"]
 
 
-def wrap_feature_collection(features: List[dict], name: str):
+def wrap_feature_collection(features: list[dict], name: str):
     """
     将 GeoJSON Feature 集合包装为 FeatureCollection
     Wrap GeoJSON Feature collection as FeatureCollection

pycityagent/environment/utils/map_utils.py

@@ -1,5 +1,5 @@
 import math
-from typing import Dict, List, Literal, Optional, Tuple, Union
+from typing import Literal,  Union
 
 import numpy as np
 
@@ -14,12 +14,12 @@ def point_on_line_given_distance(start_node, end_node, distance):
     return the coordinates of the point reached after traveling s units along the line, starting from start_point.
 
     Args:
-        start_point (tuple): Tuple of (x, y) representing the starting point on the line.
-        end_point (tuple): Tuple of (x, y) representing the ending point on the line.
+        start_point (tuple): tuple of (x, y) representing the starting point on the line.
+        end_point (tuple): tuple of (x, y) representing the ending point on the line.
         distance (float): Distance to travel along the line, starting from start_point.
 
     Returns:
-        tuple: Tuple of (x, y) representing the new point reached after traveling s units along the line.
+        tuple: tuple of (x, y) representing the new point reached after traveling s units along the line.
     """
 
     x1, y1 = start_node["x"], start_node["y"]
@@ -49,7 +49,7 @@ def point_on_line_given_distance(start_node, end_node, distance):
 
 
 def get_key_index_in_lane(
-    nodes: List[Dict[str, float]],
+    nodes: list[dict[str, float]],
     distance: float,
     direction: Union[Literal["front"], Literal["back"]],
 ) -> int:
@@ -61,10 +61,10 @@ def get_key_index_in_lane(
         _index_offset, _index_factor = len(_nodes) - 1, -1
     else:
         raise ValueError(f"Invalid direction type {direction}!")
-    _lane_points: List[Tuple[float, float, float]] = [
+    _lane_points: list[tuple[float, float, float]] = [
         (n["x"], n["y"], n.get("z", 0)) for n in _nodes
     ]
-    _line_lengths: List[float] = [0.0 for _ in range(len(_nodes))]
+    _line_lengths: list[float] = [0.0 for _ in range(len(_nodes))]
     _s = 0.0
     for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
         _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])
@@ -82,20 +82,20 @@ def get_key_index_in_lane(
 
 
 def get_xy_in_lane(
-    nodes: List[Dict[str, float]],
+    nodes: list[dict[str, float]],
     distance: float,
     direction: Union[Literal["front"], Literal["back"]],
-) -> Tuple[float, float]:
+) -> tuple[float, float]:
     if direction == "front":
         _nodes = [n for n in nodes]
     elif direction == "back":
         _nodes = [n for n in nodes[::-1]]
     else:
         raise ValueError(f"Invalid direction type {direction}!")
-    _lane_points: List[Tuple[float, float, float]] = [
+    _lane_points: list[tuple[float, float, float]] = [
         (n["x"], n["y"], n.get("z", 0)) for n in _nodes
     ]
-    _line_lengths: List[float] = [0.0 for _ in range(len(_nodes))]
+    _line_lengths: list[float] = [0.0 for _ in range(len(_nodes))]
     _s = 0.0
     for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
         _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])
@@ -122,7 +122,7 @@ def get_xy_in_lane(
 
 
 def get_direction_by_s(
-    nodes: List[Dict[str, float]],
+    nodes: list[dict[str, float]],
     distance: float,
     direction: Union[Literal["front"], Literal["back"]],
 ) -> float:
@@ -132,11 +132,11 @@ def get_direction_by_s(
         _nodes = [n for n in nodes[::-1]]
     else:
         raise ValueError(f"Invalid direction type {direction}!")
-    _lane_points: List[Tuple[float, float, float]] = [
+    _lane_points: list[tuple[float, float, float]] = [
         (n["x"], n["y"], n.get("z", 0)) for n in _nodes
     ]
-    _line_lengths: List[float] = [0.0 for _ in range(len(_nodes))]
-    _line_directions: List[Tuple[float, float]] = []
+    _line_lengths: list[float] = [0.0 for _ in range(len(_nodes))]
+    _line_directions: list[tuple[float, float]] = []
     _s = 0.0
     for i, (cur_p, next_p) in enumerate(zip(_lane_points[:-1], _lane_points[1:])):
         _s += math.hypot(next_p[0] - cur_p[0], next_p[1] - cur_p[1])

pycityagent/llm/embedding.py

@@ -1,7 +1,6 @@
 """简单的基于内存的embedding实现"""
 
 import numpy as np
-from typing import List, Dict, Optional
 import hashlib
 import json
 
@@ -22,34 +21,34 @@ class SimpleEmbedding:
         """
         self.vector_dim = vector_dim
         self.cache_size = cache_size
-        self._cache: Dict[str, np.ndarray] = {}
-        self._vocab: Dict[str, int] = {}  # 词汇表
-        self._idf: Dict[str, float] = {}  # 逆文档频率
+        self._cache: dict[str, np.ndarray] = {}
+        self._vocab: dict[str, int] = {}  # 词汇表
+        self._idf: dict[str, float] = {}  # 逆文档频率
         self._doc_count = 0  # 文档总数
 
     def _text_to_hash(self, text: str) -> str:
         """将文本转换为hash值"""
         return hashlib.md5(text.encode()).hexdigest()
 
-    def _tokenize(self, text: str) -> List[str]:
+    def _tokenize(self, text: str) -> list[str]:
         """简单的分词"""
         # 这里使用简单的空格分词，实际应用中可以使用更复杂的分词方法
         return text.lower().split()
 
-    def _update_vocab(self, tokens: List[str]):
+    def _update_vocab(self, tokens: list[str]):
         """更新词汇表"""
         for token in set(tokens):  # 使用set去重
             if token not in self._vocab:
                 self._vocab[token] = len(self._vocab)
 
-    def _update_idf(self, tokens: List[str]):
+    def _update_idf(self, tokens: list[str]):
         """更新IDF值"""
         self._doc_count += 1
         unique_tokens = set(tokens)
         for token in unique_tokens:
             self._idf[token] = self._idf.get(token, 0) + 1
 
-    def _calculate_tf(self, tokens: List[str]) -> Dict[str, float]:
+    def _calculate_tf(self, tokens: list[str]) -> dict[str, float]:
         """计算词频(TF)"""
         tf = {}
         total_tokens = len(tokens)
@@ -60,7 +59,7 @@ class SimpleEmbedding:
             tf[token] /= total_tokens
         return tf
 
-    def _calculate_tfidf(self, tokens: List[str]) -> np.ndarray:
+    def _calculate_tfidf(self, tokens: list[str]) -> np.ndarray:
         """计算TF-IDF向量"""
         vector = np.zeros(self.vector_dim)
         tf = self._calculate_tf(tokens)

pycityagent/llm/llm.py

@@ -14,7 +14,7 @@ import requests
 from dashscope import ImageSynthesis
 from PIL import Image
 from io import BytesIO
-from typing import Any, Optional, Union, List, Dict
+from typing import Any, Optional, Union
 from .llmconfig import *
 from .utils import *
 
@@ -117,8 +117,8 @@ Token Usage:
         presence_penalty: Optional[float] = None,
         timeout: int = 300,
         retries=3,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        tool_choice: Optional[Dict[str, Any]] = None,
+        tools: Optional[list[dict[str, Any]]] = None,
+        tool_choice: Optional[dict[str, Any]] = None,
     ):
         """
         异步版文本请求
@@ -227,9 +227,9 @@ Token Usage:
                     self.prompt_tokens_used += result_response.usage.prompt_tokens  # type: ignore
                     self.completion_tokens_used += result_response.usage.completion_tokens  # type: ignore
                     self.request_number += 1
-                    if tools and result_response.choices[0].message.tool_calls:
+                    if tools and result_response.choices[0].message.tool_calls:  # type: ignore
                         return json.loads(
-                            result_response.choices[0]
+                            result_response.choices[0]  # type: ignore
                             .message.tool_calls[0]
                             .function.arguments
                         )