pycityagent 2.0.0a94__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a95__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/message/messager.py

@@ -7,8 +7,6 @@ from typing import Any, Optional, Union
 import ray
 from aiomqtt import Client
 
-from .message_interceptor import MessageInterceptor
-
 __all__ = [
     "Messager",
 ]
@@ -69,10 +67,10 @@ class Messager:
             - `Union[None, ray.ObjectRef]`: The message interceptor reference.
         """
         return self._message_interceptor
-    
+
     def get_log_list(self):
         return self._log_list
-    
+
     def clear_log_list(self):
         self._log_list = []
 
@@ -213,7 +211,7 @@ class Messager:
             "from_uuid": from_uuid,
             "to_uuid": to_uuid,
             "start_time": start_time,
-            "consumption": 0
+            "consumption": 0,
         }
         message = json.dumps(payload, default=str)
         interceptor = self.message_interceptor

pycityagent/metrics/__init__.py

@@ -1,4 +1,4 @@
-from .mlflow_client import MlflowClient,init_mlflow_connection
+from .mlflow_client import MlflowClient, init_mlflow_connection
 
 __all__ = [
     "MlflowClient",

pycityagent/metrics/mlflow_client.py

@@ -6,22 +6,21 @@ 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 mlflow.entities import Metric, Param, Run, RunTag
 
+from ..configs import MlflowConfig
 from ..utils.decorators import lock_decorator
 
-__all__ = ["init_mlflow_connection","MlflowClient",]
+__all__ = [
+    "init_mlflow_connection",
+    "MlflowClient",
+]
 
 logger = logging.getLogger("mlflow")
 
 
 def init_mlflow_connection(
-    config: dict,
+    config: MlflowConfig,
     experiment_uuid: str,
     mlflow_run_name: Optional[str] = None,
     experiment_name: Optional[str] = None,
@@ -32,7 +31,7 @@ def init_mlflow_connection(
     Initialize an MLflow connection with a new or existing experiment and run.
 
     - **Args**:
-        config (dict): Configuration dictionary containing MLflow credentials and URI.
+        config (MlflowConfig): Configuration containing MLflow credentials and URI.
         experiment_uuid (str): A unique identifier for the experiment.
         mlflow_run_name (str, optional): Name of the MLflow run. Defaults to a generated name.
         experiment_name (str, optional): Name of the experiment. Defaults to a generated name.
@@ -42,8 +41,10 @@ def init_mlflow_connection(
     - **Returns**:
         tuple: A tuple containing the run_id and another tuple with the MLflow URI, client, run object, and run UUID.
     """
-    os.environ["MLFLOW_TRACKING_USERNAME"] = config.get("username", None)
-    os.environ["MLFLOW_TRACKING_PASSWORD"] = config.get("password", None)
+    if config.username is not None:
+        os.environ["MLFLOW_TRACKING_USERNAME"] = config.username
+    if config.password is not None:
+        os.environ["MLFLOW_TRACKING_PASSWORD"] = config.password
 
     run_uuid = str(uuid.uuid4())
     # run name
@@ -61,7 +62,7 @@ def init_mlflow_connection(
         experiment_tags["mlflow.note.content"] = experiment_description
         experiment_tags["experiment_id"] = experiment_uuid
 
-    uri = config["mlflow_uri"]
+    uri = config.mlflow_uri
     client = mlflow.MlflowClient(tracking_uri=uri)
 
     # experiment
@@ -89,7 +90,7 @@ class MlflowClient:
 
     def __init__(
         self,
-        config: dict,
+        config: MlflowConfig,
         experiment_uuid: str,
         mlflow_run_name: Optional[str] = None,
         experiment_name: Optional[str] = None,
@@ -101,7 +102,7 @@ class MlflowClient:
         Initialize the MlflowClient.
 
         - **Args**:
-            config (dict): Configuration dictionary containing MLflow credentials and URI.
+            config (dict): Configuration containing MLflow credentials and URI.
             experiment_uuid (str): A unique identifier for the experiment.
             mlflow_run_name (str, optional): Name of the MLflow run. Defaults to a generated name.
             experiment_name (str, optional): Name of the experiment. Defaults to a generated name.
@@ -124,9 +125,11 @@ class MlflowClient:
                 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._mlflow_uri = uri = config.mlflow_uri
+            if config.username is not None:
+                os.environ["MLFLOW_TRACKING_USERNAME"] = config.username
+            if config.password is not None:
+                os.environ["MLFLOW_TRACKING_PASSWORD"] = config.password
             self._client = client = mlflow.MlflowClient(tracking_uri=uri)
             self._run = client.get_run(run_id=run_id)
             self._run_id = run_id

pycityagent/simulation/agentgroup.py

@@ -15,10 +15,9 @@ import ray
 from langchain_core.embeddings import Embeddings
 
 from ..agent import Agent, InstitutionAgent
-from ..economy.econ_client import EconomyClient
-from ..environment import Simulator
+from ..configs import SimConfig
+from ..environment import EconomyClient, Simulator
 from ..llm.llm import LLM
-from ..llm.llmconfig import LLMConfig
 from ..memory import FaissQuery, Memory
 from ..message import Messager
 from ..metrics import MlflowClient
@@ -36,7 +35,7 @@ class AgentGroup:
         agent_class: Union[type[Agent], list[type[Agent]]],
         number_of_agents: Union[int, list[int]],
         memory_config_function_group: dict[type[Agent], Callable],
-        config: dict,
+        config: SimConfig,
         map_ref: ray.ObjectRef,
         exp_name: str,
         exp_id: Union[str, UUID],
@@ -65,7 +64,7 @@ class AgentGroup:
             - `agent_class` (Union[Type[Agent], List[Type[Agent]]]): A single or list of agent classes to instantiate.
             - `number_of_agents` (Union[int, List[int]]): Number of instances to create for each agent class.
             - `memory_config_function_group` (Dict[Type[Agent], Callable]): Functions to configure memory for each agent type.
-            - `config` (dict): Configuration settings for the agent group.
+            - `config` (SimConfig): Configuration settings for the agent group.
             - `map_ref` (ray.ObjectRef): Reference to the map object.
             - `exp_name` (str): Name of the experiment.
             - `exp_id` (str | UUID): Identifier for the experiment.
@@ -109,11 +108,11 @@ class AgentGroup:
         if self.enable_pgsql:
             pass
         # Mlflow
-        _mlflow_config = config.get("metric_request", {}).get("mlflow")
-        if _mlflow_config:
+        metric_config = config.prop_metric_request
+        if metric_config is not None and metric_config.mlflow is not None:
             logger.info(f"-----Creating Mlflow client in AgentGroup {self._uuid} ...")
             self.mlflow_client = MlflowClient(
-                config=_mlflow_config,
+                config=metric_config.mlflow,
                 experiment_uuid=self.exp_id,  # type:ignore
                 mlflow_run_name=f"{exp_name}_{1000*int(time.time())}",
                 experiment_name=exp_name,
@@ -123,12 +122,13 @@ class AgentGroup:
             self.mlflow_client = None
 
         # prepare Messager
-        if "mqtt" in config["simulator_request"]:
+        mqtt_config = config.prop_mqtt
+        if mqtt_config is not None:
             self.messager = Messager.remote(
-                hostname=config["simulator_request"]["mqtt"]["server"],  # type:ignore
-                port=config["simulator_request"]["mqtt"]["port"],
-                username=config["simulator_request"]["mqtt"].get("username", None),
-                password=config["simulator_request"]["mqtt"].get("password", None),
+                hostname=mqtt_config.server,  # type:ignore
+                port=mqtt_config.port,
+                username=mqtt_config.username,
+                password=mqtt_config.password,
             )
         else:
             self.messager = None
@@ -140,23 +140,21 @@ class AgentGroup:
         self.initialized = False
         self.id2agent = {}
         # prepare LLM client
-        llmConfig = LLMConfig(config["llm_request"])
         logger.info(f"-----Creating LLM client in AgentGroup {self._uuid} ...")
-        self.llm = LLM(llmConfig)
+        self.llm = LLM(config.prop_llm_request)
         self.llm.set_semaphore(llm_semaphore)
 
         # prepare Simulator
         logger.info(f"-----Initializing Simulator in AgentGroup {self._uuid} ...")
-        self.simulator = Simulator(config["simulator_request"])
+        self.simulator = Simulator(config)
+        self.simulator.set_environment(environment)
         self.simulator.set_map(map_ref)
         self.projector = pyproj.Proj(
             ray.get(self.simulator.map.get_projector.remote())  # type:ignore
         )
         # prepare Economy client
         logger.info(f"-----Creating Economy client in AgentGroup {self._uuid} ...")
-        self.economy_client = EconomyClient(
-            config["simulator_request"]["simulator"]["server"]
-        )
+        self.economy_client = EconomyClient(config.prop_simulator_server_address)
 
         # set FaissQuery
         if self.embedding_model is not None: