pycityagent 2.0.0a18__py3-none-any.whl → 2.0.0a19__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, Dict, List, 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,6 +53,7 @@ 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,
     ) -> None:
@@ -68,6 +67,7 @@ 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.
         """
@@ -78,6 +78,7 @@ 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
@@ -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.
@@ -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] = {}
-        
+
     async def bind_to_simulator(self):
         await self._bind_to_economy()
 
@@ -624,22 +648,24 @@ 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]:
         """从多个智能体收集消息
-        
+
         Args:
             agent_uuids: 目标智能体UUID列表
             target: 要收集的信息类型
-            
+
         Returns:
             List[dict]: 收集到的所有响应
         """
@@ -648,7 +674,7 @@ class InstitutionAgent(Agent):
         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,11 @@ 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)

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/metrics/__init__.py

@@ -0,0 +1,5 @@
+from .mlflow_client import MlflowClient
+
+__all__ = [
+    "MlflowClient",
+]

pycityagent/metrics/mlflow_client.py

@@ -0,0 +1,109 @@
+import asyncio
+import logging
+import os
+import uuid
+from collections.abc import Sequence
+from typing import Any, Optional, Union
+
+import mlflow
+from mlflow.entities import (Dataset, DatasetInput, Document, Experiment,
+                             ExperimentTag, FileInfo, InputTag, LifecycleStage,
+                             LiveSpan, Metric, NoOpSpan, Param, Run, RunData,
+                             RunInfo, RunInputs, RunStatus, RunTag, SourceType,
+                             Span, SpanEvent, SpanStatus, SpanStatusCode,
+                             SpanType, Trace, TraceData, TraceInfo, ViewType)
+
+from ..utils.decorators import lock_decorator
+
+logger = logging.getLogger("mlflow")
+
+
+class MlflowClient:
+    """
+    - Mlflow client
+    """
+
+    def __init__(
+        self,
+        config: dict,
+        mlflow_run_name: Optional[str] = None,
+        experiment_name: Optional[str] = None,
+        experiment_description: Optional[str] = None,
+        experiment_tags: Optional[dict[str, Any]] = None,
+    ) -> None:
+        os.environ["MLFLOW_TRACKING_USERNAME"] = config.get("username", None)
+        os.environ["MLFLOW_TRACKING_PASSWORD"] = config.get("password", None)
+        self._mlflow_uri = uri = config["mlflow_uri"]
+        self._client = client = mlflow.MlflowClient(tracking_uri=uri)
+        self._run_uuid = run_uuid = str(uuid.uuid4())
+        self._lock = asyncio.Lock()
+        # run name
+        if mlflow_run_name is None:
+            mlflow_run_name = f"exp_{run_uuid}"
+
+        # exp name
+        if experiment_name is None:
+            experiment_name = f"run_{run_uuid}"
+
+        # tags
+        if experiment_tags is None:
+            experiment_tags = {}
+        if experiment_description is not None:
+            experiment_tags["mlflow.note.content"] = experiment_description
+
+        try:
+            self._experiment_id = experiment_id = client.create_experiment(
+                name=experiment_name,
+                tags=experiment_tags,
+            )
+        except Exception as e:
+            experiment = client.get_experiment_by_name(experiment_name)
+            if experiment is None:
+                raise e
+            self._experiment_id = experiment_id = experiment.experiment_id
+
+        self._run = run = client.create_run(
+            experiment_id=experiment_id, run_name=mlflow_run_name
+        )
+        self._run_id = run.info.run_id
+
+    @property
+    def client(
+        self,
+    ) -> mlflow.MlflowClient:
+        return self._client
+
+    @property
+    def run_id(
+        self,
+    ) -> str:
+        return self._run_id
+
+    @lock_decorator
+    async def log_batch(
+        self,
+        metrics: Sequence[Metric] = (),
+        params: Sequence[Param] = (),
+        tags: Sequence[RunTag] = (),
+    ):
+        self.client.log_batch(
+            run_id=self.run_id, metrics=metrics, params=params, tags=tags
+        )
+
+    @lock_decorator
+    async def log_metric(
+        self,
+        key: str,
+        value: float,
+        step: Optional[int] = None,
+        timestamp: Optional[int] = None,
+    ):
+        if timestamp is not None:
+            timestamp = int(timestamp)
+        self.client.log_metric(
+            run_id=self.run_id,
+            key=key,
+            value=value,
+            timestamp=timestamp,
+            step=step,
+        )

pycityagent/simulation/agentgroup.py

@@ -1,34 +1,51 @@
 import asyncio
-from datetime import datetime
 import json
 import logging
-from pathlib import Path
+import time
 import uuid
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+from uuid import UUID
+
 import fastavro
 import ray
-from uuid import UUID
-from pycityagent.agent import Agent, CitizenAgent, InstitutionAgent
-from pycityagent.economy.econ_client import EconomyClient
-from pycityagent.environment.simulator import Simulator
-from pycityagent.llm.llm import LLM
-from pycityagent.llm.llmconfig import LLMConfig
-from pycityagent.message import Messager
-from pycityagent.utils import STATUS_SCHEMA, PROFILE_SCHEMA, DIALOG_SCHEMA, SURVEY_SCHEMA, INSTITUTION_STATUS_SCHEMA
-from typing import Any
+
+from ..agent import Agent, CitizenAgent, InstitutionAgent
+from ..economy.econ_client import EconomyClient
+from ..environment.simulator import Simulator
+from ..llm.llm import LLM
+from ..llm.llmconfig import LLMConfig
+from ..message import Messager
+from ..metrics import MlflowClient
+from ..utils import (DIALOG_SCHEMA, INSTITUTION_STATUS_SCHEMA, PROFILE_SCHEMA,
+                     STATUS_SCHEMA, SURVEY_SCHEMA)
 
 logger = logging.getLogger("pycityagent")
 
+
 @ray.remote
 class AgentGroup:
-    def __init__(self, agents: list[Agent], config: dict, exp_id: str|UUID, enable_avro: bool, avro_path: Path, logging_level: int):
+    def __init__(
+        self,
+        agents: list[Agent],
+        config: dict,
+        exp_id: str | UUID,
+        exp_name: str,
+        enable_avro: bool,
+        avro_path: Path,
+        enable_pgsql: bool,
+        pgsql_args: tuple[str, str, str, str, str],
+        logging_level: int,
+    ):
         logger.setLevel(logging_level)
         self._uuid = str(uuid.uuid4())
         self.agents = agents
         self.config = config
         self.exp_id = exp_id
         self.enable_avro = enable_avro
-        self.avro_path = avro_path / f"{self._uuid}"
         if enable_avro:
+            self.avro_path = avro_path / f"{self._uuid}"
             self.avro_path.mkdir(parents=True, exist_ok=True)
             self.avro_file = {
                 "profile": self.avro_path / f"profile.avro",
@@ -36,7 +53,7 @@ class AgentGroup:
                 "status": self.avro_path / f"status.avro",
                 "survey": self.avro_path / f"survey.avro",
             }
-            
+
         self.messager = Messager(
             hostname=config["simulator_request"]["mqtt"]["server"],
             port=config["simulator_request"]["mqtt"]["port"],
@@ -63,21 +80,35 @@ class AgentGroup:
         else:
             self.economy_client = None
 
+        # Mlflow
+        _mlflow_config = config.get("metric_request", {}).get("mlflow")
+        if _mlflow_config:
+            logger.info(f"-----Creating Mlflow client in AgentGroup {self._uuid} ...")
+            self.mlflow_client = MlflowClient(
+                config=_mlflow_config,
+                mlflow_run_name=f"EXP_{exp_name}_{1000*int(time.time())}",
+                experiment_name=exp_name,
+            )
+        else:
+            self.mlflow_client = None
+
         for agent in self.agents:
-            agent.set_exp_id(self.exp_id) # type: ignore
+            agent.set_exp_id(self.exp_id)  # type: ignore
             agent.set_llm_client(self.llm)
             agent.set_simulator(self.simulator)
             if self.economy_client is not None:
                 agent.set_economy_client(self.economy_client)
+            if self.mlflow_client is not None:
+                agent.set_mlflow_client(self.mlflow_client)
             agent.set_messager(self.messager)
             if self.enable_avro:
-                agent.set_avro_file(self.avro_file) # type: ignore
+                agent.set_avro_file(self.avro_file)  # type: ignore
 
     async def init_agents(self):
         logger.debug(f"-----Initializing Agents in AgentGroup {self._uuid} ...")
         logger.debug(f"-----Binding Agents to Simulator in AgentGroup {self._uuid} ...")
         for agent in self.agents:
-            await agent.bind_to_simulator() # type: ignore
+            await agent.bind_to_simulator()  # type: ignore
         self.id2agent = {agent._uuid: agent for agent in self.agents}
         logger.debug(f"-----Binding Agents to Messager in AgentGroup {self._uuid} ...")
         await self.messager.connect()
@@ -104,7 +135,7 @@ class AgentGroup:
                     for agent in self.agents:
                         profile = await agent.memory._profile.export()
                         profile = profile[0]
-                        profile['id'] = agent._uuid
+                        profile["id"] = agent._uuid
                         profiles.append(profile)
                     fastavro.writer(f, PROFILE_SCHEMA, profiles)
 
@@ -139,7 +170,9 @@ class AgentGroup:
         return results
 
     async def update(self, target_agent_uuid: str, target_key: str, content: Any):
-        logger.debug(f"-----Updating {target_key} for agent {target_agent_uuid} in group {self._uuid}")
+        logger.debug(
+            f"-----Updating {target_key} for agent {target_agent_uuid} in group {self._uuid}"
+        )
         agent = self.id2agent[target_agent_uuid]
         await agent.memory.update(target_key, content)
 
@@ -147,7 +180,9 @@ class AgentGroup:
         logger.debug(f"-----Starting message dispatch for group {self._uuid}")
         while True:
             if not self.messager.is_connected():
-                logger.warning("Messager is not connected. Skipping message processing.")
+                logger.warning(
+                    "Messager is not connected. Skipping message processing."
+                )
 
             # Step 1: 获取消息
             messages = await self.messager.fetch_messages()
@@ -165,7 +200,7 @@ class AgentGroup:
 
                 # 提取 agent_id（主题格式为 "exps/{exp_id}/agents/{agent_uuid}/{topic_type}"）
                 _, _, _, agent_uuid, topic_type = topic.strip("/").split("/")
-                    
+
                 if agent_uuid in self.id2agent:
                     agent = self.id2agent[agent_uuid]
                     # topic_type: agent-chat, user-chat, user-survey, gather

pycityagent/simulation/simulation.py

@@ -4,22 +4,21 @@ import logging
 import os
 import random
 import uuid
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Optional, Union
 
 import pycityproto.city.economy.v2.economy_pb2 as economyv2
 import yaml
 from mosstool.map._map_util.const import AOI_START_ID
 
-from pycityagent.environment.simulator import Simulator
-from pycityagent.memory.memory import Memory
-from pycityagent.message.messager import Messager
-from pycityagent.survey import Survey
-
 from ..agent import Agent, InstitutionAgent
+from ..environment.simulator import Simulator
+from ..memory.memory import Memory
+from ..message.messager import Messager
+from ..survey import Survey
 from .agentgroup import AgentGroup
 
 logger = logging.getLogger("pycityagent")
@@ -50,15 +49,16 @@ class AgentSimulation:
             self.agent_class = [agent_class]
         self.logging_level = logging_level
         self.config = config
+        self.exp_name = exp_name
         self._simulator = Simulator(config["simulator_request"])
         self.agent_prefix = agent_prefix
-        self._agents: Dict[uuid.UUID, Agent] = {}
-        self._groups: Dict[str, AgentGroup] = {}  # type:ignore
-        self._agent_uuid2group: Dict[uuid.UUID, AgentGroup] = {}  # type:ignore
-        self._agent_uuids: List[uuid.UUID] = []
-        self._user_chat_topics: Dict[uuid.UUID, str] = {}
-        self._user_survey_topics: Dict[uuid.UUID, str] = {}
-        self._user_interview_topics: Dict[uuid.UUID, str] = {}
+        self._agents: dict[uuid.UUID, Agent] = {}
+        self._groups: dict[str, AgentGroup] = {}  # type:ignore
+        self._agent_uuid2group: dict[uuid.UUID, AgentGroup] = {}  # type:ignore
+        self._agent_uuids: list[uuid.UUID] = []
+        self._user_chat_topics: dict[uuid.UUID, str] = {}
+        self._user_survey_topics: dict[uuid.UUID, str] = {}
+        self._user_interview_topics: dict[uuid.UUID, str] = {}
         self._loop = asyncio.get_event_loop()
 
         self._messager = Messager(
@@ -69,18 +69,37 @@ class AgentSimulation:
         )
         asyncio.create_task(self._messager.connect())
 
-        self._enable_avro = config["storage"]["avro"]["enabled"]
+        # storage
+        _storage_config: dict[str, Any] = config.get("storage", {})
+        # avro
+        _avro_config: dict[str, Any] = _storage_config.get("avro", {})
+        self._enable_avro = _avro_config.get("enabled", False)
         if not self._enable_avro:
+            self._avro_path = None
             logger.warning("AVRO is not enabled, NO AVRO LOCAL STORAGE")
-        self._avro_path = Path(config["storage"]["avro"]["path"]) / f"{self.exp_id}"
-        self._avro_path.mkdir(parents=True, exist_ok=True)
-
-        self._enable_pgsql = config["storage"]["pgsql"]["enabled"]
-        self._pgsql_host = config["storage"]["pgsql"]["host"]
-        self._pgsql_port = config["storage"]["pgsql"]["port"]
-        self._pgsql_database = config["storage"]["pgsql"]["database"]
-        self._pgsql_user = config["storage"]["pgsql"]["user"]
-        self._pgsql_password = config["storage"]["pgsql"]["password"]
+        else:
+            self._avro_path = Path(_avro_config["path"]) / f"{self.exp_id}"
+            self._avro_path.mkdir(parents=True, exist_ok=True)
+
+        # pg
+        _pgsql_config: dict[str, Any] = _storage_config.get("pgsql", {})
+        self._enable_pgsql = _pgsql_config.get("enabled", False)
+        if not self._enable_pgsql:
+            logger.warning("PostgreSQL is not enabled, NO POSTGRESQL DATABASE STORAGE")
+            self._pgsql_args = ("", "", "", "", "")
+        else:
+            self._pgsql_host = _pgsql_config["host"]
+            self._pgsql_port = _pgsql_config["port"]
+            self._pgsql_database = _pgsql_config["database"]
+            self._pgsql_user = _pgsql_config.get("user", None)
+            self._pgsql_password = _pgsql_config.get("password", None)
+            self._pgsql_args: tuple[str, str, str, str, str] = (
+                self._pgsql_host,
+                self._pgsql_port,
+                self._pgsql_database,
+                self._pgsql_user,
+                self._pgsql_password,
+            )
 
         # 添加实验信息相关的属性
         self._exp_info = {
@@ -96,14 +115,34 @@ class AgentSimulation:
         }
 
         # 创建异步任务保存实验信息
-        self._exp_info_file = self._avro_path / "experiment_info.yaml"
-        with open(self._exp_info_file, "w") as f:
-            yaml.dump(self._exp_info, f)
+        if self._enable_avro:
+            assert self._avro_path is not None
+            self._exp_info_file = self._avro_path / "experiment_info.yaml"
+            with open(self._exp_info_file, "w") as f:
+                yaml.dump(self._exp_info, f)
+
+    @property
+    def enable_avro(
+        self,
+    ) -> bool:
+        return self._enable_avro
+
+    @property
+    def enable_pgsql(
+        self,
+    ) -> bool:
+        return self._enable_pgsql
 
     @property
-    def agents(self):
+    def agents(self) -> dict[uuid.UUID, Agent]:
         return self._agents
 
+    @property
+    def avro_path(
+        self,
+    ) -> Path:
+        return self._avro_path  # type:ignore
+
     @property
     def groups(self):
         return self._groups
@@ -122,13 +161,24 @@ class AgentSimulation:
         agents: list[Agent],
         config: dict,
         exp_id: str,
+        exp_name: str,
         enable_avro: bool,
         avro_path: Path,
+        enable_pgsql: bool,
+        pgsql_args: tuple[str, str, str, str, str],
         logging_level: int = logging.WARNING,
     ):
         """创建远程组"""
         group = AgentGroup.remote(
-            agents, config, exp_id, enable_avro, avro_path, logging_level
+            agents,
+            config,
+            exp_id,
+            exp_name,
+            enable_avro,
+            avro_path,
+            enable_pgsql,
+            pgsql_args,
+            logging_level,
         )
         return group_name, group, agents
 
@@ -174,7 +224,6 @@ class AgentSimulation:
                     memory_config_func.append(self.default_memory_config_institution)
                 else:
                     memory_config_func.append(self.default_memory_config_citizen)
-
         # 使用线程池并行创建 AgentGroup
         group_creation_params = []
         class_init_index = 0
@@ -226,8 +275,11 @@ class AgentSimulation:
                 agents,
                 self.config,
                 self.exp_id,
-                self._enable_avro,
-                self._avro_path,
+                self.exp_name,
+                self.enable_avro,
+                self.avro_path,
+                self.enable_pgsql,
+                self._pgsql_args,
                 self.logging_level,
             )
             creation_tasks.append((group_name, group, agents))
@@ -393,7 +445,7 @@ class AgentSimulation:
         return EXTRA_ATTRIBUTES, PROFILE, BASE
 
     async def send_survey(
-        self, survey: Survey, agent_uuids: Optional[List[uuid.UUID]] = None
+        self, survey: Survey, agent_uuids: Optional[list[uuid.UUID]] = None
     ):
         """发送问卷"""
         survey_dict = survey.to_dict()
@@ -410,7 +462,7 @@ class AgentSimulation:
             await self._messager.send_message(topic, payload)
 
     async def send_interview_message(
-        self, content: str, agent_uuids: Union[uuid.UUID, List[uuid.UUID]]
+        self, content: str, agent_uuids: Union[uuid.UUID, list[uuid.UUID]]
     ):
         """发送面试消息"""
         payload = {
@@ -438,10 +490,17 @@ class AgentSimulation:
     async def _save_exp_info(self) -> None:
         """异步保存实验信息到YAML文件"""
         try:
-            with open(self._exp_info_file, "w") as f:
-                yaml.dump(self._exp_info, f)
+            if self.enable_avro:
+                with open(self._exp_info_file, "w") as f:
+                    yaml.dump(self._exp_info, f)
         except Exception as e:
-            logger.error(f"保存实验信息失败: {str(e)}")
+            logger.error(f"Avro保存实验信息失败: {str(e)}")
+        try:
+            if self.enable_pgsql:
+                #   TODO
+                pass
+        except Exception as e:
+            logger.error(f"PostgreSQL保存实验信息失败: {str(e)}")
 
     async def _update_exp_status(self, status: int, error: str = "") -> None:
         """更新实验状态并保存"""
@@ -492,7 +551,6 @@ class AgentSimulation:
                     tasks = []
                     for group in self._groups.values():
                         tasks.append(group.run.remote())
-
                     # 等待所有group运行完成
                     await asyncio.gather(*tasks)
 

pycityagent/workflow/__init__.py

@@ -4,14 +4,16 @@
 This module contains classes for creating blocks and running workflows.
 """
 
-from .block import Block, log_and_check, log_and_check_with_memory, trigger_class
+from .block import (Block, log_and_check, log_and_check_with_memory,
+                    trigger_class)
 from .prompt import FormatPrompt
-from .tool import GetMap, SencePOI, Tool
-from .trigger import MemoryChangeTrigger, TimeTrigger, EventTrigger
+from .tool import ExportMlflowMetrics, GetMap, SencePOI, Tool
+from .trigger import EventTrigger, MemoryChangeTrigger, TimeTrigger
 
 __all__ = [
     "SencePOI",
     "Tool",
+    "ExportMlflowMetrics",
     "GetMap",
     "MemoryChangeTrigger",
     "TimeTrigger",

pycityagent/workflow/block.py

@@ -3,12 +3,11 @@ import functools
 import inspect
 from typing import Any, Callable, Coroutine, Optional, Union
 
-from pycityagent.environment.simulator import Simulator
-from pycityagent.workflow.trigger import EventTrigger
-
+from ..environment.simulator import Simulator
 from ..llm import LLM
 from ..memory import Memory
 from ..utils.decorators import record_call_aio
+from ..workflow.trigger import EventTrigger
 
 TRIGGER_INTERVAL = 1
 

pycityagent/workflow/tool.py

@@ -1,7 +1,11 @@
+import time
 from typing import Any, Callable, Dict, List, Optional, Union
 
+from mlflow.entities import Metric
+
 from ..agent import Agent
-from ..environment import LEVEL_ONE_PRE, POI_TYPE_DICT, AoiService, PersonService
+from ..environment import (LEVEL_ONE_PRE, POI_TYPE_DICT, AoiService,
+                           PersonService)
 from ..workflow import Block
 
 
@@ -139,7 +143,7 @@ class UpdateWithSimulator(Tool):
         if agent._simulator is None:
             return
         if not agent._has_bound_to_simulator:
-            await agent._bind_to_simulator() # type: ignore
+            await agent._bind_to_simulator()  # type: ignore
         simulator = agent.simulator
         memory = agent.memory
         person_id = await memory.get("id")
@@ -181,3 +185,48 @@ class ResetAgentPosition(Tool):
             lane_id=lane_id,
             s=s,
         )
+
+
+class ExportMlflowMetrics(Tool):
+    def __init__(self, log_batch_size: int = 100) -> None:
+        self._log_batch_size = log_batch_size
+        # TODO:support other log types
+        self.metric_log_cache: list[Metric] = []
+
+    async def __call__(
+        self,
+        metric: Union[Metric, dict],
+        clear_cache: bool = False,
+    ):
+        agent = self.agent
+        batch_size = self._log_batch_size
+        if len(self.metric_log_cache) > batch_size:
+            client = agent.mlflow_client
+            await client.log_batch(
+                metrics=self.metric_log_cache[:batch_size],
+            )
+            self.metric_log_cache = self.metric_log_cache[batch_size:]
+        else:
+            if isinstance(metric, Metric):
+                self.metric_log_cache.append(metric)
+            else:
+                _metric = Metric(
+                    key=metric["key"],
+                    value=metric["value"],
+                    timestamp=metric.get("timestamp", int(1000 * time.time())),
+                    step=metric["step"],
+                )
+                self.metric_log_cache.append(_metric)
+        if clear_cache:
+            await self._clear_cache()
+
+    async def _clear_cache(
+        self,
+    ):
+        agent = self.agent
+        client = agent.mlflow_client
+        if len(self.metric_log_cache) > 0:
+            await client.log_batch(
+                metrics=self.metric_log_cache,
+            )
+            self.metric_log_cache = []

{pycityagent-2.0.0a18.dist-info → pycityagent-2.0.0a19.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a18
+Version: 2.0.0a19
 Summary: LLM-based城市环境agent构建库
 License: MIT
 Author: Yuwei Yan
@@ -26,6 +26,7 @@ Requires-Dist: gradio (>=5.7.1,<6.0.0)
 Requires-Dist: grpcio (==1.67.1)
 Requires-Dist: langchain-core (>=0.3.28,<0.4.0)
 Requires-Dist: matplotlib (==3.8.3)
+Requires-Dist: mlflow (>=2.19.0,<3.0.0)
 Requires-Dist: mosstool (==1.0.24)
 Requires-Dist: networkx (==3.2.1)
 Requires-Dist: numpy (>=1.20.0,<2.0.0)

{pycityagent-2.0.0a18.dist-info → pycityagent-2.0.0a19.dist-info}/RECORD

@@ -1,7 +1,7 @@
 pycityagent/__init__.py,sha256=EDxt3Su3lH1IMh9suNw7GeGL7UrXeWiZTw5KWNznDzc,637
-pycityagent/agent.py,sha256=t9W9sKxtQ0EkMxL78kAjAu-rXigEK6gyLY0IEA4DbnQ,23143
+pycityagent/agent.py,sha256=gy9OlFpnEla7wK45k4m7_ySLp67mBKoG50vxUEmxbJE,24280
 pycityagent/economy/__init__.py,sha256=aonY4WHnx-6EGJ4WKrx4S-2jAkYNLtqUA04jp6q8B7w,75
-pycityagent/economy/econ_client.py,sha256=DE11Ng_NO_foW65A-LxFW0VED-HLrnn4GwUf_Xn-Tlg,11189
+pycityagent/economy/econ_client.py,sha256=EZDGxM7K83ucYZQ5qdv6HA-jhRCWbR1u5q-kLMqelKc,11192
 pycityagent/environment/__init__.py,sha256=awHxlOud-btWbk0FCS4RmGJ13W84oVCkbGfcrhKqihA,240
 pycityagent/environment/interact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycityagent/environment/interact/interact.py,sha256=ifxPPzuHeqLHIZ_6zvfXMoBOnBsXNIP4bYp7OJ7pnEQ,6588
@@ -21,7 +21,7 @@ pycityagent/environment/sim/person_service.py,sha256=nIvOsoBoqOTDYtsiThg07-4ZBgk
 pycityagent/environment/sim/road_service.py,sha256=phKTwTyhc_6Ht2mddEXpdENfl-lRXIVY0CHAlw1yHjI,1264
 pycityagent/environment/sim/sim_env.py,sha256=HI1LcS_FotDKQ6vBnx0e49prXSABOfA20aU9KM-ZkCY,4625
 pycityagent/environment/sim/social_service.py,sha256=6Iqvq6dz8H2jhLLdtaITc6Js9QnQw-Ylsd5AZgUj3-E,1993
-pycityagent/environment/simulator.py,sha256=K7IyhiGC9BxanW28bpML4M0YREdMp1h7yMoWBlbf3RY,12504
+pycityagent/environment/simulator.py,sha256=XjcxbyBIbB3Ht9z087z_oWIPAN6pP5Eq1lyf4W5atb8,12502
 pycityagent/environment/utils/__init__.py,sha256=1m4Q1EfGvNpUsa1bgQzzCyWhfkpElnskNImjjFD3Znc,237
 pycityagent/environment/utils/base64.py,sha256=hoREzQo3FXMN79pqQLO2jgsDEvudciomyKii7MWljAM,374
 pycityagent/environment/utils/const.py,sha256=3RMNy7_bE7-23K90j9DFW_tWEzu8s7hSTgKbV-3BFl4,5327
@@ -45,9 +45,12 @@ pycityagent/memory/state.py,sha256=5W0c1yJ-aaPpE74B2LEcw3Ygpm77tyooHv8NylyrozE,5
 pycityagent/memory/utils.py,sha256=wLNlNlZ-AY9VB8kbUIy0UQSYh26FOQABbhmKQkit5o8,850
 pycityagent/message/__init__.py,sha256=TCjazxqb5DVwbTu1fF0sNvaH_EPXVuj2XQ0p6W-QCLU,55
 pycityagent/message/messager.py,sha256=W_OVlNGcreHSBf6v-DrEnfNCXExB78ySr0w26MSncfU,2541
+pycityagent/metrics/__init__.py,sha256=IrlwXs_733b3PlMV99Ogh6R8sGCEalF2qIT2mQwdPjU,75
+pycityagent/metrics/mlflow_client.py,sha256=BH-YVMlYr-eTFdrCZjgMZWi8ptUEm7todOXCpYixAgU,3311
+pycityagent/metrics/utils/const.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycityagent/simulation/__init__.py,sha256=jYaqaNpzM5M_e_ykISS_M-mIyYdzJXJWhgpfBpA6l5k,111
-pycityagent/simulation/agentgroup.py,sha256=M19XWJRWyjMAYS0_RIOBQ2C7I1MuVYIaX3DgehGZL2Y,12541
-pycityagent/simulation/simulation.py,sha256=9VH-VaRufUcQcVhqXcQJTMtrquns8pbDpRTfK6736Io,19318
+pycityagent/simulation/agentgroup.py,sha256=H2E_YQ3ir3gQmPWsupHXA_LSBVkffzkXWl4UgZ9qLOc,13327
+pycityagent/simulation/simulation.py,sha256=ANf04GqaK7gLT50F_ZcP2UcpS_uGx12pLqkrK6eA2L8,21014
 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
@@ -59,11 +62,11 @@ pycityagent/utils/parsers/code_block_parser.py,sha256=Cs2Z_hm9VfNCpPPll1TwteaJF-
 pycityagent/utils/parsers/json_parser.py,sha256=FZ3XN1g8z4Dr2TFraUOoah1oQcze4fPd2m01hHoX0Mo,2917
 pycityagent/utils/parsers/parser_base.py,sha256=k6DVqwAMK3jJdOP4IeLE-aFPm3V2F-St5qRBuRdx4aU,1742
 pycityagent/utils/survey_util.py,sha256=Be9nptmu2JtesFNemPgORh_2GsN7rcDYGQS9Zfvc5OI,2169
-pycityagent/workflow/__init__.py,sha256=EyCcjB6LyBim-5iAOPe4m2qfvghEPqu1ZdGfy4KPeZ8,551
-pycityagent/workflow/block.py,sha256=6EmiRMLdOZC1wMlmLMIjfrp9TuiI7Gw4s3nnXVMbrnw,6031
+pycityagent/workflow/__init__.py,sha256=QNkUV-9mACMrR8c0cSKna2gC1mMZdxXbxWzjE-Uods0,621
+pycityagent/workflow/block.py,sha256=WkE2On97DCZS_9n8aIgT8wxv9Oaff4Fdf2tLqbKfMtE,6010
 pycityagent/workflow/prompt.py,sha256=tY69nDO8fgYfF_dOA-iceR8pAhkYmCqoox8uRPqEuGY,2956
-pycityagent/workflow/tool.py,sha256=_bCluIX8HTC8ZW6a-wrMB3Uhx2yzD8sM8XFDI3vd0MM,6642
+pycityagent/workflow/tool.py,sha256=uaB0dV35jA9v2UqWw9L8iPM-HJW5e9BlgFVgOMf9jvw,8201
 pycityagent/workflow/trigger.py,sha256=t5X_i0WtL32bipZSsq_E3UUyYYudYLxQUpvxbgClp2s,5683
-pycityagent-2.0.0a18.dist-info/METADATA,sha256=iajlG1hMfjQA2qBuTYITWemGBg4TEPU8hHpSU1l0MDs,7760
-pycityagent-2.0.0a18.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-pycityagent-2.0.0a18.dist-info/RECORD,,
+pycityagent-2.0.0a19.dist-info/METADATA,sha256=x0DbwNQ4Zj0rswWGIv6MA0zSBBEiExXLIdErYA0bhnE,7800
+pycityagent-2.0.0a19.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+pycityagent-2.0.0a19.dist-info/RECORD,,