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

pycityagent/memory/memory.py

@@ -2,7 +2,8 @@ import asyncio
 import logging
 from copy import deepcopy
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
+from typing import Any, Literal, Optional,  Union
+from collections.abc import Sequence,Callable
 
 import numpy as np
 from pyparsing import deque
@@ -27,10 +28,10 @@ class Memory:
 
     def __init__(
         self,
-        config: Optional[Dict[Any, Any]] = None,
-        profile: Optional[Dict[Any, Any]] = None,
-        base: Optional[Dict[Any, Any]] = None,
-        motion: Optional[Dict[Any, Any]] = None,
+        config: Optional[dict[Any, Any]] = None,
+        profile: Optional[dict[Any, Any]] = None,
+        base: Optional[dict[Any, Any]] = None,
+        motion: Optional[dict[Any, Any]] = None,
         activate_timestamp: bool = False,
         embedding_model: Any = None,
     ) -> None:
@@ -38,7 +39,7 @@ class Memory:
         Initializes the Memory with optional configuration.
 
         Args:
-            config (Optional[Dict[Any, Any]], optional):
+            config (Optional[dict[Any, Any]], optional):
                 A configuration dictionary for dynamic memory. The dictionary format is:
                 - Key: The name of the dynamic memory field.
                 - Value: Can be one of two formats:
@@ -46,24 +47,24 @@ class Memory:
                     2. A callable that returns the default value when invoked (useful for complex default values).
                 Note: If a key in `config` overlaps with predefined attributes in `PROFILE_ATTRIBUTES` or `STATE_ATTRIBUTES`, a warning will be logged, and the key will be ignored.
                 Defaults to None.
-            profile (Optional[Dict[Any, Any]], optional): profile attribute dict.
-            base (Optional[Dict[Any, Any]], optional): base attribute dict from City Simulator.
-            motion (Optional[Dict[Any, Any]], optional): motion attribute dict from City Simulator.
+            profile (Optional[dict[Any, Any]], optional): profile attribute dict.
+            base (Optional[dict[Any, Any]], optional): base attribute dict from City Simulator.
+            motion (Optional[dict[Any, Any]], optional): motion attribute dict from City Simulator.
             activate_timestamp (bool): Whether activate timestamp storage in MemoryUnit
             embedding_model (Any): The embedding model for memory search.
         """
-        self.watchers: Dict[str, List[Callable]] = {}
+        self.watchers: dict[str, list[Callable]] = {}
         self._lock = asyncio.Lock()
         self.embedding_model = embedding_model
 
         # 初始化embedding存储
         self._embeddings = {"state": {}, "profile": {}, "dynamic": {}}
 
-        _dynamic_config: Dict[Any, Any] = {}
-        _state_config: Dict[Any, Any] = {}
-        _profile_config: Dict[Any, Any] = {}
+        _dynamic_config: dict[Any, Any] = {}
+        _state_config: dict[Any, Any] = {}
+        _profile_config: dict[Any, Any] = {}
         # 记录哪些字段需要embedding
-        self._embedding_fields: Dict[str, bool] = {}
+        self._embedding_fields: dict[str, bool] = {}
 
         if config is not None:
             for k, v in config.items():
@@ -303,7 +304,7 @@ class Memory:
 
     async def update_batch(
         self,
-        content: Union[Dict, Sequence[Tuple[Any, Any]]],
+        content: Union[dict, Sequence[tuple[Any, Any]]],
         mode: Union[Literal["replace"], Literal["merge"]] = "replace",
         store_snapshot: bool = False,
         protect_llm_read_only_fields: bool = True,
@@ -312,7 +313,7 @@ class Memory:
         Updates multiple values in the memory at once.
 
         Args:
-            content (Union[Dict, Sequence[Tuple[Any, Any]]]): A dictionary or sequence of tuples containing the keys and values to update.
+            content (Union[dict, Sequence[tuple[Any, Any]]]): A dictionary or sequence of tuples containing the keys and values to update.
             mode (Union[Literal["replace"], Literal["merge"]], optional): Update mode. Defaults to "replace".
             store_snapshot (bool): Whether to store a snapshot of the memory after the update.
             protect_llm_read_only_fields (bool): Whether to protect non-self define fields from being updated.
@@ -321,9 +322,9 @@ class Memory:
             TypeError: If the content type is neither a dictionary nor a sequence of tuples.
         """
         if isinstance(content, dict):
-            _list_content: List[Tuple[Any, Any]] = [(k, v) for k, v in content.items()]
+            _list_content: list[tuple[Any, Any]] = [(k, v) for k, v in content.items()]
         elif isinstance(content, Sequence):
-            _list_content: List[Tuple[Any, Any]] = [(k, v) for k, v in content]
+            _list_content: list[tuple[Any, Any]] = [(k, v) for k, v in content]
         else:
             raise TypeError(f"Invalid content type `{type(content)}`!")
         for k, v in _list_content[:1]:
@@ -353,12 +354,12 @@ class Memory:
     @lock_decorator
     async def export(
         self,
-    ) -> Tuple[Sequence[Dict], Sequence[Dict], Sequence[Dict]]:
+    ) -> tuple[Sequence[dict], Sequence[dict], Sequence[dict]]:
         """
         Exports the current state of all memory sections.
 
         Returns:
-            Tuple[Sequence[Dict], Sequence[Dict], Sequence[Dict]]: A tuple containing the exported data of profile, state, and dynamic memory sections.
+            tuple[Sequence[dict], Sequence[dict], Sequence[dict]]: A tuple containing the exported data of profile, state, and dynamic memory sections.
         """
         return (
             await self._profile.export(),
@@ -369,14 +370,14 @@ class Memory:
     @lock_decorator
     async def load(
         self,
-        snapshots: Tuple[Sequence[Dict], Sequence[Dict], Sequence[Dict]],
+        snapshots: tuple[Sequence[dict], Sequence[dict], Sequence[dict]],
         reset_memory: bool = True,
     ) -> None:
         """
         Import the snapshot memories of all sections.
 
         Args:
-            snapshots (Tuple[Sequence[Dict], Sequence[Dict], Sequence[Dict]]): The exported snapshots.
+            snapshots (tuple[Sequence[dict], Sequence[dict], Sequence[dict]]): The exported snapshots.
             reset_memory (bool): Whether to reset previous memory.
         """
         _profile_snapshot, _state_snapshot, _dynamic_snapshot = snapshots

pycityagent/metrics/__init__.py

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

pycityagent/metrics/mlflow_client.py

@@ -0,0 +1,147 @@
+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")
+
+
+def init_mlflow_connection(
+    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,
+) -> tuple[str, tuple[str, mlflow.MlflowClient, Run, str]]:
+
+    os.environ["MLFLOW_TRACKING_USERNAME"] = config.get("username", None)
+    os.environ["MLFLOW_TRACKING_PASSWORD"] = config.get("password", None)
+
+    run_uuid = str(uuid.uuid4())
+    # 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
+
+    uri = config["mlflow_uri"]
+    client = mlflow.MlflowClient(tracking_uri=uri)
+
+    # experiment
+    try:
+        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
+        experiment_id = experiment.experiment_id
+
+    # run
+    run = client.create_run(experiment_id=experiment_id, run_name=mlflow_run_name)
+
+    run_id = run.info.run_id
+
+    return run_id, (uri, client, run, run_uuid)
+
+
+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,
+        run_id: Optional[str] = None,
+    ) -> None:
+        if run_id is None:
+            self._run_id, (
+                self._mlflow_uri,
+                self._client,
+                self._run,
+                self._run_uuid,
+            ) = init_mlflow_connection(
+                config=config,
+                mlflow_run_name=mlflow_run_name,
+                experiment_name=experiment_name,
+                experiment_description=experiment_description,
+                experiment_tags=experiment_tags,
+            )
+        else:
+            self._mlflow_uri = uri = config["mlflow_uri"]
+            os.environ["MLFLOW_TRACKING_USERNAME"] = config.get("username", None)
+            os.environ["MLFLOW_TRACKING_PASSWORD"] = config.get("password", None)
+            self._client = client = mlflow.MlflowClient(tracking_uri=uri)
+            self._run = client.get_run(run_id=run_id)
+            self._run_id = run_id
+            self._run_uuid = run_uuid = str(uuid.uuid4())
+        self._lock = asyncio.Lock()
+
+    @property
+    def client(
+        self,
+    ) -> mlflow.MlflowClient:
+        return self._client
+
+    @property
+    def run_id(
+        self,
+    ) -> str:
+        assert self._run_id is not None
+        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,52 @@
 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_copy_writer: ray.ObjectRef,
+        mlflow_run_id: 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 +54,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 +81,36 @@ 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,
+                run_id=mlflow_run_id,
+            )
+        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 +137,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 +172,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 +182,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 +202,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