pycityagent 2.0.0a18__py3-none-any.whl → 2.0.0a20__py3-none-any.whl

pycityagent/agent.py

@@ -1,30 +1,28 @@
 """智能体模板类及其定义"""
 
-from abc import ABC, abstractmethod
 import asyncio
-from uuid import UUID
-from copy import deepcopy
-from datetime import datetime
-from enum import Enum
 import logging
 import random
 import uuid
-from typing import Dict, List, Optional,Any
+from abc import ABC, abstractmethod
+from copy import deepcopy
+from datetime import datetime
+from enum import Enum
+from typing import Any, Optional
+from uuid import UUID
 
 import fastavro
-
-from pycityagent.environment.sim.person_service import PersonService
 from mosstool.util.format_converter import dict2pb
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
-from pycityagent.utils import process_survey_for_llm
-
-from pycityagent.message.messager import Messager
-from pycityagent.utils import SURVEY_SCHEMA, DIALOG_SCHEMA
 
 from .economy import EconomyClient
 from .environment import Simulator
+from .environment.sim.person_service import PersonService
 from .llm import LLM
 from .memory import Memory
+from .message.messager import Messager
+from .metrics import MlflowClient
+from .utils import DIALOG_SCHEMA, SURVEY_SCHEMA, process_survey_for_llm
 
 logger = logging.getLogger("pycityagent")
 
@@ -55,8 +53,9 @@ class Agent(ABC):
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,
         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,
     ) -> None:
         """
         Initialize the Agent.
@@ -68,8 +67,9 @@ class Agent(ABC):
             economy_client (EconomyClient): The `EconomySim` client. Defaults to None.
             messager (Messager, optional): The messager object. Defaults to None.
             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.
         """
         self._name = name
         self._type = type
@@ -78,13 +78,14 @@ class Agent(ABC):
         self._economy_client = economy_client
         self._messager = messager
         self._simulator = simulator
+        self._mlflow_client = mlflow_client
         self._memory = memory
         self._exp_id = -1
         self._agent_id = -1
         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
 
@@ -112,6 +113,12 @@ class Agent(ABC):
         """
         self._simulator = simulator
 
+    def set_mlflow_client(self, mlflow_client: MlflowClient):
+        """
+        Set the mlflow_client of the agent.
+        """
+        self._mlflow_client = mlflow_client
+
     def set_economy_client(self, economy_client: EconomyClient):
         """
         Set the economy_client of the agent.
@@ -130,7 +137,7 @@ 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.
         """
@@ -164,6 +171,15 @@ class Agent(ABC):
             )
         return self._economy_client
 
+    @property
+    def mlflow_client(self):
+        """The Agent's MlflowClient"""
+        if self._mlflow_client is None:
+            raise RuntimeError(
+                f"MlflowClient access before assignment, please `set_mlflow_client` first!"
+            )
+        return self._mlflow_client
+
     @property
     def memory(self):
         """The Agent's Memory"""
@@ -218,19 +234,21 @@ class Agent(ABC):
         response = await self._llm_client.atext_request(dialog)  # type:ignore
 
         return response  # type:ignore
-    
+
     async def _process_survey(self, survey: dict):
         survey_response = await self.generate_user_survey_response(survey)
         if self._avro_file is None:
             return
-        response_to_avro = [{
-            "id": self._uuid,
-            "day": await self.simulator.get_simulator_day(),
-            "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),
-        }]
+        response_to_avro = [
+            {
+                "id": self._uuid,
+                "day": await self.simulator.get_simulator_day(),
+                "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),
+            }
+        ]
         with open(self._avro_file["survey"], "a+b") as f:
             fastavro.writer(f, SURVEY_SCHEMA, response_to_avro, codec="snappy")
 
@@ -270,28 +288,32 @@ class Agent(ABC):
         response = await self._llm_client.atext_request(dialog)  # type:ignore
 
         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),
-        }]
+        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),
+            }
+        ]
         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),
-        })
+        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:
@@ -303,15 +325,17 @@ 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),
-        }]
+        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),
+            }
+        ]
         asyncio.create_task(self.process_agent_chat_response(payload))
         if self._avro_file is None:
             return
@@ -341,18 +365,14 @@ class Agent(ABC):
         raise NotImplementedError
 
     # MQTT send message
-    async def _send_message(
-        self, to_agent_uuid: str, payload: dict, sub_topic: str
-    ):
+    async def _send_message(self, to_agent_uuid: str, payload: dict, sub_topic: str):
         """通过 Messager 发送消息"""
         if self._messager is None:
             raise RuntimeError("Messager is not set")
         topic = f"exps/{self._exp_id}/agents/{to_agent_uuid}/{sub_topic}"
         await self._messager.send_message(topic, payload)
 
-    async def send_message_to_agent(
-        self, to_agent_uuid: str, content: str
-    ):
+    async def send_message_to_agent(self, to_agent_uuid: str, content: str):
         """通过 Messager 发送消息"""
         if self._messager is None:
             raise RuntimeError("Messager is not set")
@@ -364,15 +384,17 @@ 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),
-        }]
+        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:
@@ -403,20 +425,22 @@ class CitizenAgent(Agent):
         name: str,
         llm_client: Optional[LLM] = None,
         simulator: Optional[Simulator] = None,
+        mlflow_client: Optional[MlflowClient] = None,
         memory: Optional[Memory] = None,
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,
         avro_file: Optional[dict] = None,
     ) -> None:
         super().__init__(
-            name,
-            AgentType.Citizen,
-            llm_client,
-            economy_client,
-            messager,
-            simulator,
-            memory,
-            avro_file,
+            name=name,
+            type=AgentType.Citizen,
+            llm_client=llm_client,
+            economy_client=economy_client,
+            messager=messager,
+            simulator=simulator,
+            mlflow_client=mlflow_client,
+            memory=memory,
+            avro_file=avro_file,
         )
 
     async def bind_to_simulator(self):
@@ -464,9 +488,7 @@ class CitizenAgent(Agent):
                 )
                 person_id = resp["person_id"]
                 await memory.update("id", person_id, protect_llm_read_only_fields=False)
-                logger.debug(
-                    f"Binding to Person `{person_id}` just added to Simulator"
-                )
+                logger.debug(f"Binding to Person `{person_id}` just added to Simulator")
                 # 防止模拟器还没有到prepare阶段导致get_person出错
             self._has_bound_to_simulator = True
             self._agent_id = person_id
@@ -517,24 +539,26 @@ class InstitutionAgent(Agent):
         name: str,
         llm_client: Optional[LLM] = None,
         simulator: Optional[Simulator] = None,
+        mlflow_client: Optional[MlflowClient] = None,
         memory: Optional[Memory] = None,
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,
         avro_file: Optional[dict] = None,
     ) -> None:
         super().__init__(
-            name,
-            AgentType.Institution,
-            llm_client,
-            economy_client,
-            messager,
-            simulator,
-            memory,
-            avro_file,
+            name=name,
+            type=AgentType.Institution,
+            llm_client=llm_client,
+            economy_client=economy_client,
+            mlflow_client=mlflow_client,
+            messager=messager,
+            simulator=simulator,
+            memory=memory,
+            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()
 
@@ -624,31 +648,33 @@ class InstitutionAgent(Agent):
         """处理收到的消息，识别发送者"""
         content = payload["content"]
         sender_id = payload["from"]
-        
+
         # 将响应存储到对应的Future中
         response_key = str(sender_id)
         if response_key in self._gather_responses:
-            self._gather_responses[response_key].set_result({
-                "from": sender_id,
-                "content": content,
-            })
+            self._gather_responses[response_key].set_result(
+                {
+                    "from": sender_id,
+                    "content": content,
+                }
+            )
 
-    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:
             agent_uuids: 目标智能体UUID列表
             target: 要收集的信息类型
-            
+
         Returns:
-            List[dict]: 收集到的所有响应
+            list[dict]: 收集到的所有响应
         """
         # 为每个agent创建Future
         futures = {}
         for agent_uuid in agent_uuids:
             futures[agent_uuid] = asyncio.Future()
             self._gather_responses[agent_uuid] = futures[agent_uuid]
-            
+
         # 发送gather请求
         payload = {
             "from": self._uuid,
@@ -656,7 +682,7 @@ class InstitutionAgent(Agent):
         }
         for agent_uuid in agent_uuids:
             await self._send_message(agent_uuid, payload, "gather")
-            
+
         try:
             # 等待所有响应
             responses = await asyncio.gather(*futures.values())

pycityagent/economy/econ_client.py

@@ -308,11 +308,48 @@ class EconomyClient:
         # current agent ids and org ids
         return (list(response.agent_ids), list(response.org_ids))
 
-    async def get_org_entity_ids(self, org_type: economyv2.OrgType)->list[int]:
+    async def get_org_entity_ids(self, org_type: economyv2.OrgType) -> list[int]:
         request = org_service.GetOrgEntityIdsRequest(
             type=org_type,
         )
         response: org_service.GetOrgEntityIdsResponse = (
             await self._aio_stub.GetOrgEntityIds(request)
         )
-        return list(response.org_ids)
+        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/simulator.py

@@ -5,7 +5,7 @@ import logging
 import os
 from collections.abc import Sequence
 from datetime import datetime, timedelta
-from typing import Any, Optional, Tuple, Union, cast
+from typing import Any, Optional, Union, cast
 
 from mosstool.type import TripMode
 from mosstool.util.format_converter import coll2pb
@@ -28,7 +28,7 @@ class Simulator:
     - Simulator Class
     """
 
-    def __init__(self, config, secure: bool = False) -> None:
+    def __init__(self, config:dict, secure: bool = False) -> None:
         self.config = config
         """
         - 模拟器配置

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
                         )