refactor monitor

src/agentscope/_init.py

@@ -10,10 +10,12 @@
 from ._runtime import Runtime
 from .file_manager import file_manager
 from .utils.logging_utils import LOG_LEVEL, setup_logger
+from .utils.monitor import MonitorFactory
 from .models import read_model_configs
+from .constants import _DEFAULT_DIR
+from .constants import _DEFAULT_LOG_LEVEL
+
 
-_DEFAULT_DIR = "./runs"
-_DEFAULT_LOG_LEVEL = "INFO"
 _INIT_SETTINGS = {}
 
 
@@ -85,6 +87,9 @@ def init(
     dir_log = str(file_manager.dir_log) if save_log else None
     setup_logger(dir_log, logger_level)
 
+    # Set monitor
+    _ = MonitorFactory.get_monitor(db_path=file_manager.path_db)
+
     # Load config and init agent by configs
     if agent_configs is not None:
         if isinstance(agent_configs, str):

src/agentscope/configs/model_config.py

@@ -3,6 +3,7 @@
 from typing import Any
 from ..constants import _DEFAULT_MAX_RETRIES
 from ..constants import _DEFAULT_MESSAGES_KEY
+from ..constants import _DEFAULT_API_BUDGET
 
 
 class CfgBase(dict):
@@ -57,6 +58,9 @@ class OpenAICfg(CfgBase):
     """The arguments used in openai api generation, e.g. `temperature`,
     `seed`."""
 
+    budget: float = _DEFAULT_API_BUDGET
+    """The total budget using this model. Set to `None` means no limit."""
+
 
 class PostApiCfg(CfgBase):
     """The config for Post API. The final request post will be
@@ -113,3 +117,6 @@ class PostApiCfg(CfgBase):
     """The key of the prompt messages in `requests.post()`,
     e.g. `request.post(json={${messages_key}: messages, **json_args})`. For
     huggingface and modelscope inference API, the key is `inputs`"""
+
+    budget: float = _DEFAULT_API_BUDGET
+    """The total budget using this model. Set to `None` means no limit."""

src/agentscope/constants.py

@@ -16,13 +16,17 @@
 _DEFAULT_SUBDIR_FILE = "file"
 _DEFAULT_SUBDIR_INVOKE = "invoke"
 _DEFAULT_IMAGE_NAME = "image_{}_{}.png"
+_DEFAULT_SQLITE_DB_PATH = "agentscope.db"
 # for model wrapper
 _DEFAULT_MAX_RETRIES = 3
 _DEFAULT_MESSAGES_KEY = "inputs"
 _DEFAULT_RETRY_INTERVAL = 1
+_DEFAULT_API_BUDGET = None
 # for execute python
 _DEFAULT_PYPI_MIRROR = "http://mirrors.aliyun.com/pypi/simple/"
 _DEFAULT_TRUSTED_HOST = "mirrors.aliyun.com"
+# for monitor
+_DEFAULT_MONITOR_TABLE_NAME = "monitor_metrics"
 # for summarization
 _DEFAULT_SUMMARIZATION_PROMPT = """
 TEXT: {}

src/agentscope/file_manager.py

@@ -14,6 +14,7 @@
     _DEFAULT_SUBDIR_CODE,
     _DEFAULT_SUBDIR_FILE,
     _DEFAULT_SUBDIR_INVOKE,
+    _DEFAULT_SQLITE_DB_PATH,
     _DEFAULT_IMAGE_NAME,
 )
 
@@ -48,6 +49,10 @@ def _get_and_create_subdir(self, subdir: str) -> str:
             os.makedirs(path)
         return path
 
+    def _get_file_path(self, file_name: str) -> str:
+        """Get the path of the file."""
+        return os.path.join(self.dir, Runtime.runtime_id, file_name)
+
     @property
     def dir_log(self) -> str:
         """The directory for saving logs."""
@@ -69,11 +74,17 @@ def dir_invoke(self) -> str:
         """The directory for saving api invocations."""
         return self._get_and_create_subdir(_DEFAULT_SUBDIR_INVOKE)
 
+    @property
+    def path_db(self) -> str:
+        """The path to the sqlite db file."""
+        return self._get_file_path(_DEFAULT_SQLITE_DB_PATH)
+
     def init(self, save_dir: str, save_api_invoke: bool = False) -> None:
         """Set the directory for saving files."""
         self.dir = save_dir
-        if not os.path.exists(save_dir):
-            os.makedirs(save_dir)
+        runtime_dir = os.path.join(save_dir, Runtime.runtime_id)
+        if not os.path.exists(runtime_dir):
+            os.makedirs(runtime_dir)
 
         self.save_api_invoke = save_api_invoke
 

src/agentscope/models/openai_model.py

@@ -12,7 +12,8 @@
 except ImportError:
     openai = None
 
-from ..utils import MonitorFactory
+from ..utils.monitor import MonitorFactory
+from ..utils.monitor import get_full_name
 from ..utils import QuotaExceededError
 from ..utils.token_utils import get_openai_max_length
 
@@ -28,6 +29,7 @@ def __init__(
         organization: str = None,
         client_args: dict = None,
         generate_args: dict = None,
+        budget: float = None,
     ) -> None:
         """Initialize the openai client.
 
@@ -49,6 +51,9 @@ def __init__(
             generate_args (`dict`, default `None`):
                 The extra keyword arguments used in openai api generation,
                 e.g. `temperature`, `seed`.
+            budget (`float`, default `None`):
+                The total budget using this model. Set to `None` means no
+                limit.
         """
         super().__init__(name)
 
@@ -77,8 +82,18 @@ def __init__(
 
         # Set monitor accordingly
         self.monitor = None
+        self.budget = budget
+        self._register_budget()
         self._register_default_metrics()
 
+    def _register_budget(self) -> None:
+        self.monitor = MonitorFactory.get_monitor()
+        self.monitor.register_budget(
+            model_name=self.model_name,
+            value=self.budget,
+            prefix=self.model_name,
+        )
+
     def _register_default_metrics(self) -> None:
         """Register metrics to the monitor."""
         raise NotImplementedError(
@@ -95,7 +110,7 @@ def _metric(self, metric_name: str) -> str:
         Returns:
             `str`: Metric name of this wrapper.
         """
-        return f"{self.__class__.__name__}.{self.model_name}.{metric_name}"
+        return get_full_name(name=metric_name, prefix=self.model_name)
 
 
 class OpenAIChatWrapper(OpenAIWrapper):
@@ -193,7 +208,10 @@ def __call__(
 
         # step5: update monitor accordingly
         try:
-            self.monitor.update(**response.usage.model_dump())
+            self.monitor.update(
+                response.usage.model_dump(),
+                prefix=self.model_name,
+            )
         except QuotaExceededError as e:
             # TODO: optimize quota exceeded error handling process
             logger.error(e.message)